import argparse
import os
import time
import threading
from typing import List, Optional
from concurrent.futures import ThreadPoolExecutor, as_completed

from dns_client import make_resolver, query_mx, query_txt
from domain_utils import normalize_domain, extract_etld1_and_subdomain
from syntax_analyzer import analyze_syntax
from email_posture import analyze_dmarc, analyze_spf, probe_dkim_selectors
from scoring import score, level
from reporting import print_table, write_json, write_csv, prepare_output_path
from models import (
    Result,
    SyntaxInfo,
    MxInfo,
    DmarcInfo,
    SpfInfo,
    DkimInfo,
)

# Time helpers
def utc_timestamp() -> str:
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


# Thread-local DNS resolver (IMPORTANT for thread safety)
_thread_local = threading.local()

def get_thread_resolver(dns_servers, timeout):
    if not hasattr(_thread_local, "resolver"):
        _thread_local.resolver = make_resolver(dns_servers, timeout)
    return _thread_local.resolver


# Single-domain analysis (called by each thread)
def analyze_one(domain: str, dns_servers, timeout, brands: List[str], dkim_probe: bool) -> Result:
    resolver = get_thread_resolver(dns_servers, timeout)

    syn = analyze_syntax(domain, brands)
    etld1 = syn.registrable_domain

    mx_hosts = query_mx(resolver, etld1)
    mx_info = MxInfo(present=bool(mx_hosts), exchanges=mx_hosts)

    dmarc_txt = query_txt(resolver, f"_dmarc.{etld1}")
    dmarc_info = analyze_dmarc(dmarc_txt)

    spf_txt = query_txt(resolver, etld1)
    spf_info = analyze_spf(spf_txt)

    dkim_info = probe_dkim_selectors(resolver, etld1, enabled=dkim_probe)

    risk_score, reasons = score(syn, mx_info, dmarc_info, spf_info, dkim_info)

    return Result(
        domain_input=domain,
        etld_plus_one=etld1,
        risk_score=risk_score,
        risk_level=level(risk_score),
        reasons=reasons,
        syntax=syn,
        mx=mx_info,
        dmarc=dmarc_info,
        spf=spf_info,
        dkim=dkim_info,
        timestamp_utc=utc_timestamp(),
    )


# ----------------------------------------------------------------------
# CLI
# ----------------------------------------------------------------------

def load_lines(path: str) -> List[str]:
    with open(path, "r", encoding="utf-8") as f:
        return [line.strip() for line in f if line.strip() and not line.startswith("#")]


def main(argv: Optional[List[str]] = None) -> int:
    ap = argparse.ArgumentParser(
        description="Impersonation / lookalike / brand spoofing risk scorer (eTLD+1 based)"
    )

    ap.add_argument("domain", nargs="?", help="Single domain/FQDN/URL to analyze")
    ap.add_argument("-i", "--input", help="File with domains/FQDNs (one per line)")
    ap.add_argument("--brands", help="File with brand domains (one per line)")
    ap.add_argument("--dns", nargs="*", help="Custom DNS resolvers (e.g. 1.1.1.1 8.8.8.8)")
    ap.add_argument("--timeout", type=float, default=2.5, help="DNS timeout (seconds)")
    ap.add_argument("--dkim-probe", action="store_true", help="Light DKIM selector probing")
    ap.add_argument("--workers", type=int, default=50, help="Number of parallel threads")

    ap.add_argument(
        "--format", choices=["table", "json", "csv"], default="table", help="Output format"
    )
    ap.add_argument("-o", "--output", help="Output file path (optional)")

    args = ap.parse_args(argv)

    if not args.domain and not args.input:
        ap.error("Provide a domain or an input file")

    # --------------------------------------------------
    # Load inputs
    # --------------------------------------------------

    domains: List[str] = []
    if args.domain:
        domains.append(args.domain)
    if args.input:
        domains.extend(load_lines(args.input))

    brands: List[str] = []
    if args.brands:
        brands = load_lines(args.brands)

    # --------------------------------------------------
    # Parallel execution
    # --------------------------------------------------

    #results: List[Result] = []

    def task(domain: str):
        return analyze_one(
            domain=domain,
            dns_servers=args.dns,
            timeout=args.timeout,
            brands=brands,
            dkim_probe=args.dkim_probe,
        )

    from concurrent.futures import ThreadPoolExecutor, as_completed

    results = []
    total = len(domains)
    done = 0

    with ThreadPoolExecutor(max_workers=args.workers) as executor:
        futures = {executor.submit(task, d): d for d in domains}

        for future in as_completed(futures):
            domain = futures[future]
            try:
                results.append(future.result())
            except Exception as e:
                etld1, _, _ = extract_etld1_and_subdomain(normalize_domain(domain))
                results.append(
                    Result(
                        domain_input=domain,
                        etld_plus_one=etld1,
                        risk_score=0,
                        risk_level="LOW",
                        reasons=[f"error:{type(e).__name__}"],
                        syntax=SyntaxInfo(
                            input_domain=domain,
                            fqdn_provided=False,
                            registrable_domain=etld1,
                            subdomain="",
                        ),
                        mx=MxInfo(),
                        dmarc=DmarcInfo(parse_error=str(e)),
                        spf=SpfInfo(parse_error=str(e)),
                        dkim=DkimInfo(),
                        timestamp_utc=utc_timestamp(),
                    )
                )

            # -------- progress counter --------
            done += 1
            if done % 10 == 0 or done == total:
                print(f"[progress] {done}/{total} domains analyzed")


    # --------------------------------------------------
    # Output
    # --------------------------------------------------

    if args.format == "table":
        print_table(results)

    elif args.format == "json":
        if args.output:
            write_json(results, args.output)
        else:
            out_dir = prepare_output_path("reports")
            path = os.path.join(out_dir, "report.json")
            write_json(results, path)
            print(f"[+] JSON report saved to {path}")

    elif args.format == "csv":
        if args.output:
            write_csv(results, args.output)
        else:
            out_dir = prepare_output_path("reports")
            path = os.path.join(out_dir, "report.csv")
            write_csv(results, path)
            print(f"[+] CSV report saved to {path}")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())
