#!/usr/bin/env python3
"""
Antipublic CLI — standalone tool for antipublic.one API v2.

Setup:
    export ANTIPUBLIC_KEY="your_bearer_token"
    # or create .env file with ANTIPUBLIC_KEY=your_bearer_token

Usage:
    ./antipublic.py email user@example.com
    ./antipublic.py password secretpass123
    ./antipublic.py domain corp.com
    ./antipublic.py batch emails.txt
    ./antipublic.py batch emails.txt --only-passwords
    ./antipublic.py quota

Options:
    --strict          Exact match (default: prefix match)
    --all             Substring match (Premium)
    --dedup           Deduplicate by password (Premium, email/domain search)
    --passwords-only  Return only passwords, no logins (batch mode)
    --all-pages       Auto-paginate to collect all results
    --json            Output raw JSON
    --output FILE     Save results to file (appends .json for --json)
    --limit N         Limit results per email in batch mode (default 1000)
    --delay N         Delay between requests in seconds (default 3)
    --quiet           Suppress banner and status, only output results

Environment:
    ANTIPUBLIC_KEY    Bearer token (required)
"""

import argparse
import json
import os
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

try:
    import requests
except ImportError:
    print("Error: 'requests' package required. Install: pip install requests")
    sys.exit(1)


API_BASE = "https://antipublic.one/api/v2"
VERSION = "1.0.0"

NO_COLOR = os.environ.get("NO_COLOR") or not sys.stdout.isatty()


def c(code: str, text: str) -> str:
    if NO_COLOR:
        return text
    codes = {"r": "31", "g": "32", "y": "33", "b": "34", "m": "35", "c": "36", "w": "37", "d": "90"}
    return f"\033[{codes.get(code, '0')}m{text}\033[0m"


def bold(text: str) -> str:
    if NO_COLOR:
        return text
    return f"\033[1m{text}\033[0m"


def load_token() -> str:
    token = os.environ.get("ANTIPUBLIC_KEY", "")
    if token:
        return token
    for env_path in [Path(".env"), Path.home() / ".env"]:
        if env_path.exists():
            for line in env_path.read_text().splitlines():
                line = line.strip()
                if line.startswith("ANTIPUBLIC_KEY="):
                    return line.split("=", 1)[1].strip().strip("'\"")
    return ""


class AntipublicClient:
    def __init__(self, token: str, delay: float = 3.0):
        self.token = token
        self.delay = delay
        self.headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        }
        self._last_request = 0.0
        self.last_quota = None

    def _throttle(self):
        elapsed = time.time() - self._last_request
        if elapsed < self.delay and self._last_request > 0:
            time.sleep(self.delay - elapsed)
        self._last_request = time.time()

    def search(self, search_by: str, value: str, direction: str = "start",
               group_by: str = None, fmt: str = None, page_token: str = None) -> dict:
        self._throttle()
        payload = {
            "searchBy": search_by,
            "query": {search_by: value},
            "direction": {search_by: direction},
        }
        if group_by:
            payload["groupBy"] = group_by
        if fmt:
            payload["format"] = fmt
        if page_token:
            payload["pageToken"] = page_token
        r = requests.post(f"{API_BASE}/search", headers=self.headers, json=payload, timeout=30)
        data = self._parse(r)
        self.last_quota = data.get("availableQueries")
        return data

    def search_all_pages(self, search_by: str, value: str, direction: str = "start",
                         group_by: str = None, fmt: str = None) -> dict:
        all_results = []
        page_num = 0
        page_token = None
        first_response = None
        while True:
            page_num += 1
            data = self.search(search_by, value, direction, group_by, fmt, page_token)
            if first_response is None:
                first_response = data
            results = self._get_results(data)
            all_results.extend(results)
            page_token = (data.get("page") or {}).get("nextPageToken")
            if not page_token:
                break
            if not NO_COLOR:
                print(c("d", f"  ... page {page_num}, {len(all_results)} results so far"))
        first_response["results"] = all_results
        first_response["resultCount"] = len(all_results)
        return first_response

    def email_passwords(self, emails: list, limit: int = 1000, only_password: bool = False) -> dict:
        self._throttle()
        payload = {"emails": emails, "limit": limit}
        if only_password:
            payload["onlyPassword"] = True
        r = requests.post(f"{API_BASE}/emailPasswords", headers=self.headers, json=payload, timeout=60)
        data = self._parse(r)
        self.last_quota = data.get("availableQueries")
        return data

    def _parse(self, r: requests.Response) -> dict:
        try:
            return r.json()
        except Exception:
            return {"_error": f"HTTP {r.status_code}", "_raw": r.text[:2000]}

    def _get_results(self, data: dict) -> list:
        return data.get("results", data.get("data", data.get("lines", []))) or []


def print_banner(quiet: bool):
    if quiet:
        return
    print(bold("Antipublic CLI") + c("d", f" v{VERSION}"))
    print()


def print_quota(client: AntipublicClient, quiet: bool):
    if quiet or client.last_quota is None:
        return
    q = client.last_quota
    color = "g" if q > 10000 else ("y" if q > 1000 else "r")
    print(c("d", f"\n  quota: ") + c(color, f"{q:,}") + c("d", " queries remaining"))


def format_result_line(line: str) -> tuple:
    if ":" not in line:
        return line, ""
    email, password = line.split(":", 1)
    return email, password


def print_search_results(data: dict, search_by: str, value: str, quiet: bool, json_mode: bool):
    if json_mode:
        print(json.dumps(data, indent=2, ensure_ascii=False))
        return

    if "_error" in data:
        print(c("r", f"  Error: {data['_error']}"))
        if "_raw" in data:
            print(c("d", f"  {data['_raw'][:200]}"))
        return

    if not data.get("success", True):
        print(c("r", f"  Error: {data.get('error', 'unknown')}"))
        return

    results = data.get("results", data.get("data", data.get("lines", []))) or []
    count = data.get("resultCount", len(results))
    total = (data.get("page") or {}).get("totalResults", count)

    if not quiet:
        label = {
            "email": f"email {c('c', value)}",
            "password": f"password {c('m', value)}",
            "domain": f"domain {c('y', value)}",
        }.get(search_by, value)
        print(f"  {label}  →  ", end="")

    if count == 0:
        print(c("d", "no results"))
        return

    hit_label = f"{c('g', str(count))} results"
    if total > count:
        hit_label += c("d", f" (of {total} total)")
    print(hit_label)

    if search_by == "email":
        passwords = []
        for line in results:
            _, pw = format_result_line(line)
            passwords.append(pw)
        unique = sorted(set(passwords))
        print(c("d", f"  unique passwords: {len(unique)}"))
        for pw in unique:
            print(f"    {pw}")
    elif search_by == "password":
        emails = []
        for line in results:
            em, _ = format_result_line(line)
            emails.append(em)
        unique = sorted(set(emails))
        print(c("d", f"  accounts using this password: {len(unique)}"))
        for em in unique:
            print(f"    {em}")
    elif search_by == "domain":
        pairs = {}
        for line in results:
            em, pw = format_result_line(line)
            pairs.setdefault(em, []).append(pw)
        print(c("d", f"  accounts: {len(pairs)}"))
        for em in sorted(pairs):
            pws = pairs[em]
            if len(pws) == 1:
                print(f"    {em}: {pws[0]}")
            else:
                print(f"    {em}: {', '.join(pws)}")


def cmd_search(args, client: AntipublicClient):
    direction = "strict" if args.strict else ("all" if args.all else "start")
    group_by = "password" if args.dedup else None

    if args.all_pages:
        data = client.search_all_pages(args.type, args.query, direction, group_by)
    else:
        data = client.search(args.type, args.query, direction, group_by)

    print_search_results(data, args.type, args.query, args.quiet, args.json)
    print_quota(client, args.quiet)

    if args.output:
        save_results(args.output, data, args.json)

    return data


def cmd_batch(args, client: AntipublicClient):
    filepath = Path(args.file)
    if not filepath.exists():
        print(c("r", f"  File not found: {filepath}"))
        sys.exit(1)

    lines = [l.strip() for l in filepath.read_text().splitlines() if l.strip() and not l.startswith("#")]
    if not lines:
        print(c("r", "  File is empty"))
        sys.exit(1)

    if not args.quiet:
        print(f"  Loaded {c('c', str(len(lines)))} emails from {filepath.name}")
        print()

    chunk_size = 50
    all_results = []
    for i in range(0, len(lines), chunk_size):
        chunk = lines[i:i + chunk_size]
        if not args.quiet:
            start = i + 1
            end = min(i + chunk_size, len(lines))
            print(c("d", f"  [{start}-{end}/{len(lines)}] querying /emailPasswords ..."))

        data = client.email_passwords(chunk, args.limit, args.passwords_only)

        if args.json:
            print(json.dumps(data, indent=2, ensure_ascii=False))
        else:
            results = data.get("results", []) or []
            all_results.extend(results)

            if not results:
                print(c("d", "    no results"))
            else:
                grouped = {}
                for line in results:
                    if args.passwords_only:
                        grouped.setdefault("passwords", []).append(line)
                    else:
                        em, pw = format_result_line(line)
                        grouped.setdefault(em, []).append(pw)

                if args.passwords_only:
                    for pw in grouped.get("passwords", []):
                        print(f"    {pw}")
                else:
                    for em in sorted(grouped):
                        pws = grouped[em]
                        unique = sorted(set(pws))
                        print(f"    {c('c', em)}: {', '.join(unique)}")

    print_quota(client, args.quiet)

    if args.output:
        save_results(args.output, all_results if not args.json else data, args.json)

    if not args.quiet:
        unique_passwords = set()
        for line in all_results:
            if ":" in line:
                unique_passwords.add(line.split(":", 1)[1])
            else:
                unique_passwords.add(line)
        print(f"\n  Total: {c('g', str(len(all_results)))} results, "
              f"{c('g', str(len(unique_passwords)))} unique passwords")


def cmd_quota(args, client: AntipublicClient):
    data = client.search("email", "quota-check-probe@test.invalid", "strict")
    q = data.get("availableQueries", "?")
    quota_used = data.get("quotaUsed", {})
    cost = quota_used.get("detail", {})

    if args.json:
        print(json.dumps({"availableQueries": q, "quotaUsed": quota_used}, indent=2))
        return

    print(f"  /search queries remaining:  {c('g', str(q))}")
    if cost:
        print(c("d", f"  (probe cost: results={cost.get('results', '?')}, cpu={cost.get('cpu', '?')})"))


def save_results(filepath: str, data, json_mode: bool):
    p = Path(filepath)
    if json_mode or p.suffix == ".json":
        if not p.suffix:
            p = p.with_suffix(".json")
        p.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
    else:
        lines = []
        if isinstance(data, dict):
            for line in data.get("results", data.get("data", data.get("lines", []))) or []:
                lines.append(str(line))
        elif isinstance(data, list):
            lines = [str(x) for x in data]
        p.write_text("\n".join(lines) + "\n", encoding="utf-8")
    print(c("d", f"\n  Saved → {p}"))


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="antipublic",
        description="Antipublic API v2 CLI tool",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""examples:
  %(prog)s email user@example.com              Search by email
  %(prog)s email user@example.com --strict     Exact match only
  %(prog)s password secretpass123              Reverse search by password
  %(prog)s domain corp.com --all-pages         All pages for domain
  %(prog)s batch emails.txt                    Bulk /emailPasswords lookup
  %(prog)s batch emails.txt --passwords-only   Only passwords, no logins
  %(prog)s quota                               Check remaining quota""",
    )

    p.add_argument("--token", help="API token (overrides env)")
    p.add_argument("--delay", type=float, default=3.0, help="Delay between requests (default: 3s)")

    sub = p.add_subparsers(dest="command")

    def add_common(s):
        s.add_argument("--json", action="store_true", help="Output raw JSON")
        s.add_argument("--output", "-o", help="Save results to file")
        s.add_argument("--quiet", "-q", action="store_true", help="Minimal output")

    for cmd_name in ("email", "password", "domain"):
        s = sub.add_parser(cmd_name, help=f"Search by {cmd_name}")
        s.add_argument("query", help=f"The {cmd_name} to search for")
        s.add_argument("--strict", action="store_true", help="Exact match")
        s.add_argument("--all", action="store_true", help="Substring match (Premium)")
        s.add_argument("--dedup", action="store_true", help="Deduplicate by password (Premium)")
        s.add_argument("--all-pages", action="store_true", help="Fetch all pages")
        add_common(s)

    batch = sub.add_parser("batch", help="Bulk email→passwords via /emailPasswords")
    batch.add_argument("file", help="File with emails (one per line)")
    batch.add_argument("--limit", type=int, default=1000, help="Results per email (default: 1000)")
    batch.add_argument("--passwords-only", action="store_true", help="Return only passwords")
    add_common(batch)

    quota = sub.add_parser("quota", help="Check remaining API quota")
    add_common(quota)

    return p


def main():
    parser = build_parser()
    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        sys.exit(0)

    print_banner(args.quiet)

    token = args.token or load_token()
    if not token:
        print(c("r", "  No API token found."))
        print(c("d", "  Set ANTIPUBLIC_KEY env var, create .env file, or use --token"))
        sys.exit(1)

    client = AntipublicClient(token, args.delay)

    if args.command in ("email", "password", "domain"):
        args.type = args.command
        cmd_search(args, client)
    elif args.command == "batch":
        cmd_batch(args, client)
    elif args.command == "quota":
        cmd_quota(args, client)


if __name__ == "__main__":
    main()