from __future__ import annotations

import argparse
import concurrent.futures
import datetime
import json
import sys
from typing import IO

from .aggregate import Aggregate, PhaseStats, summarize
from .duration import parse_duration
from .probe import Options, Result, VerboseDetail, measure

# ── exit codes (mirrors Go) ───────────────────────────────────────────────────

EXIT_OK      = 0
EXIT_USAGE   = 1
EXIT_DNS     = 2
EXIT_CONNECT = 3
EXIT_TIMEOUT = 4
EXIT_TLS     = 5
EXIT_HTTP    = 6

_PHASE_EXIT: dict[str, int] = {
    "dns":     EXIT_DNS,
    "timeout": EXIT_TIMEOUT,
    "tls":     EXIT_TLS,
}


def _phase_code(fail_phase: str) -> int:
    return _PHASE_EXIT.get(fail_phase, EXIT_CONNECT)


# ── display constants ─────────────────────────────────────────────────────────

_SINGLE_SEP = "  " + "─" * 29
_AGG_SEP    = "  " + "─" * 49

_PHASE_LABELS = [
    ("dns",      "DNS lookup    "),
    ("connect",  "TCP connect   "),
    ("tls",      "TLS handshake "),
    ("ttfb",     "Server (TTFB) "),
    ("transfer", "Transfer      "),
]

# Response headers shown in text verbose block, in priority order.
# In JSON verbose mode all parsed headers are included.
_VERBOSE_HEADERS = [
    "Location",
    "Server",
    "Content-Type",
    "X-Cache",
    "CF-Cache-Status",
    "Cache-Control",
    "Via",
    "X-Powered-By",
]

# ── argparse with injectable streams ─────────────────────────────────────────


class _ArgExit(Exception):
    def __init__(self, code: int) -> None:
        self.code = code


class _Parser(argparse.ArgumentParser):
    """ArgumentParser that writes to injected streams and raises instead of exiting."""

    def __init__(self, *args, out: IO, err: IO, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self._out = out
        self._err = err

    def _print_message(self, message: str, file=None) -> None:
        if message:
            (file if file is not None else self._out).write(message)

    def print_help(self, file=None) -> None:
        self._print_message(self.format_help(), self._out)

    def print_usage(self, file=None) -> None:
        self._print_message(self.format_usage(), self._err)

    def error(self, message: str) -> None:
        self.print_usage()
        self._err.write(f"{self.prog}: error: {message}\n")
        raise _ArgExit(EXIT_USAGE)

    def exit(self, status: int = 0, message: str | None = None) -> None:
        if message:
            self._err.write(message)
        raise _ArgExit(int(status) if status else EXIT_OK)


# ── sampling ──────────────────────────────────────────────────────────────────


def _run_samples(
    url: str, count: int, opts: Options
) -> tuple[list[Result], list[Result]]:
    succeeded, failed = [], []
    for _ in range(count):
        r = measure(url, opts)
        (failed if r.err else succeeded).append(r)
    return succeeded, failed


# ── verbose rendering ─────────────────────────────────────────────────────────


def _cert_expired(expiry: str) -> bool:
    try:
        return datetime.date.fromisoformat(expiry) < datetime.date.today()
    except (ValueError, TypeError):
        return False


def _print_verbose_block(detail: VerboseDetail, out: IO) -> None:
    """Write the verbose detail block (IP, TLS, cert, headers) after timing rows."""
    def _row(label: str, value: str) -> None:
        out.write(f"  {label:<14} :  {value}\n")

    if detail.resolved_ip:
        _row("IP", detail.resolved_ip)

    if detail.tls_version:
        parts = [detail.tls_version]
        if detail.tls_cipher:
            parts.append(detail.tls_cipher)
        if detail.tls_bits:
            parts.append(f"{detail.tls_bits} bit")
        _row("TLS", "  ".join(parts))

    if detail.cert:
        c = detail.cert
        label = "Cert (unvrf.)" if not c.verified else "Cert"
        cert_parts = []
        if c.cn:
            cert_parts.append(f"CN={c.cn}")
        if c.expiry:
            tag = "EXPIRED" if _cert_expired(c.expiry) else "valid until"
            cert_parts.append(f"{tag} {c.expiry}")
        if c.issuer_cn:
            cert_parts.append(c.issuer_cn)
        _row(label, "  ".join(cert_parts))

    for name in _VERBOSE_HEADERS:
        value = detail.headers.get(name)
        if value:
            _row(name, value)


# ── text rendering ────────────────────────────────────────────────────────────


def _status_str(status_code: int, fail_flag: bool) -> str:
    s = str(status_code)
    if fail_flag and status_code >= 400:
        s += " ✗"
    return s


def _print_single(r: Result, fail_flag: bool, out: IO) -> None:
    status = _status_str(r.status_code, fail_flag)
    out.write(f"{r.url}  ({status})\n")
    for attr, label in _PHASE_LABELS:
        ph = getattr(r, attr)
        if ph.present:
            out.write(f"  {label} :  {ph.ms:8.2f} ms\n")
    out.write(_SINGLE_SEP + "\n")
    if r.total.present:
        out.write(f"  {'Total         '} :  {r.total.ms:8.2f} ms\n")
    if r.err is not None:
        out.write(f"  ✗ {r.fail_phase}: {r.err}\n")
    if r.detail is not None:
        _print_verbose_block(r.detail, out)


def _print_aggregate(
    agg: Aggregate,
    failed: list[Result],
    fail_flag: bool,
    out: IO,
    detail: VerboseDetail | None = None,
) -> None:
    status = _status_str(agg.status_code, fail_flag)
    header = f"{agg.url}  ({status}, {agg.count} samples"
    if failed:
        header += f", {len(failed)} failed"
    out.write(header + ")\n")

    if agg.total.present:
        out.write(f"  {'':14}    {'min':>9}  {'avg':>9}  {'max':>9}\n")
        for attr, label in _PHASE_LABELS:
            ps: PhaseStats = getattr(agg, attr)
            if ps.present:
                out.write(
                    f"  {label} :  {ps.min_ms:6.2f} ms"
                    f"  {ps.avg_ms:6.2f} ms"
                    f"  {ps.max_ms:6.2f} ms\n"
                )
        out.write(_AGG_SEP + "\n")
        out.write(
            f"  {'Total         '} :  {agg.total.min_ms:6.2f} ms"
            f"  {agg.total.avg_ms:6.2f} ms"
            f"  {agg.total.max_ms:6.2f} ms\n"
        )
    _print_failure_summary(failed, out)
    if detail is not None:
        _print_verbose_block(detail, out)


def _print_all_failed(
    url: str, failed: list[Result], total_count: int, out: IO
) -> None:
    header = f"{url}  (FAILED"
    if total_count > 1:
        header += f", 0/{total_count} succeeded"
    out.write(header + ")\n")

    last = failed[-1]
    any_phase = False
    for attr, label in _PHASE_LABELS:
        ph = getattr(last, attr)
        if ph.present:
            out.write(f"  {label} :  {ph.ms:8.2f} ms\n")
            any_phase = True
    if last.total.present:
        if any_phase:
            out.write(_SINGLE_SEP + "\n")
        out.write(f"  {'Total         '} :  {last.total.ms:8.2f} ms\n")
    _print_failure_summary(failed, out)
    if last.detail is not None:
        _print_verbose_block(last.detail, out)


def _print_failure_summary(failed: list[Result], out: IO) -> None:
    if not failed:
        return
    counts: dict[tuple[str, str], int] = {}
    order: list[tuple[str, str]] = []
    for r in failed:
        key = (r.fail_phase, str(r.err))
        if key not in counts:
            order.append(key)
            counts[key] = 0
        counts[key] += 1
    for phase, msg in order:
        n = counts[(phase, msg)]
        if n == 1:
            out.write(f"  ✗ {phase}: {msg}\n")
        else:
            out.write(f"  ✗ {n} × {phase}: {msg}\n")


def _print_url(
    url: str,
    succeeded: list[Result],
    failed: list[Result],
    total_count: int,
    fail_flag: bool,
    out: IO,
) -> None:
    n_ok   = len(succeeded)
    n_fail = len(failed)
    last_detail = succeeded[-1].detail if succeeded else None

    if n_fail == 0 and total_count == 1:
        _print_single(succeeded[0], fail_flag, out)
    elif n_fail == 0:
        _print_aggregate(summarize(succeeded), [], fail_flag, out, last_detail)
    elif n_ok == 0:
        _print_all_failed(url, failed, total_count, out)
    else:
        _print_aggregate(summarize(succeeded), failed, fail_flag, out, last_detail)


# ── JSON rendering ────────────────────────────────────────────────────────────


def _build_json_entry(
    url: str,
    succeeded: list[Result],
    failed: list[Result],
    detail: VerboseDetail | None = None,
) -> dict:
    entry: dict = {
        "url":       url,
        "status":    0,
        "succeeded": len(succeeded),
        "failed":    len(failed),
    }

    if succeeded:
        agg = summarize(succeeded)
        entry["status"] = agg.status_code
        phases: dict = {}
        for attr in ("dns", "connect", "tls", "ttfb", "transfer", "total"):
            ps: PhaseStats = getattr(agg, attr)
            if ps.present:
                phases[attr] = {
                    "min_ms": ps.min_ms,
                    "avg_ms": ps.avg_ms,
                    "max_ms": ps.max_ms,
                }
        if phases:
            entry["phases"] = phases

    if failed:
        counts: dict[tuple[str, str], int] = {}
        order: list[tuple[str, str]] = []
        for r in failed:
            key = (r.fail_phase, str(r.err))
            if key not in counts:
                order.append(key)
                counts[key] = 0
            counts[key] += 1
        entry["errors"] = [
            {"phase": ph, "count": counts[(ph, msg)], "message": msg}
            for ph, msg in order
        ]

    if detail is not None:
        v: dict = {}
        if detail.resolved_ip:
            v["ip"] = detail.resolved_ip
        if detail.tls_version:
            v["tls_version"] = detail.tls_version
            v["tls_cipher"]  = detail.tls_cipher
            v["tls_bits"]    = detail.tls_bits
        if detail.cert:
            c = detail.cert
            v["cert"] = {
                "cn":        c.cn,
                "sans":      c.sans,
                "expiry":    c.expiry,
                "issuer_cn": c.issuer_cn,
                "verified":  c.verified,
            }
        if detail.headers:
            v["headers"] = dict(detail.headers)
        if v:
            entry["verbose"] = v

    return entry


# ── entry point ───────────────────────────────────────────────────────────────


def run(args: list[str], stdout: IO, stderr: IO) -> int:
    parser = _Parser(
        prog="latprobe",
        description="measure per-phase HTTP request latency",
        out=stdout,
        err=stderr,
    )
    parser.add_argument("urls", nargs="+", metavar="url")
    parser.add_argument(
        "-n", "--count", type=int, default=1, metavar="N",
        help="number of requests per URL (default: 1)",
    )
    parser.add_argument(
        "-c", "--concurrency", type=int, default=0, metavar="N",
        help="max parallel URLs, 0=auto (default: 0)",
    )
    parser.add_argument(
        "--timeout", default="10s", metavar="DURATION",
        help="per-request timeout, e.g. 10s, 500ms (default: 10s)",
    )
    parser.add_argument(
        "--fail", action="store_true",
        help="exit non-zero on HTTP status >= 400",
    )
    parser.add_argument(
        "--json", action="store_true", dest="json_out",
        help="output results as JSON instead of text",
    )
    parser.add_argument(
        "-v", "--verbose", action="store_true",
        help="show resolved IP, TLS version/cipher, certificate, and response headers",
    )

    try:
        ns = parser.parse_args(args)
    except _ArgExit as exc:
        return exc.code

    try:
        timeout_secs = parse_duration(ns.timeout)
    except ValueError:
        stderr.write(f"latprobe: error: invalid timeout: {ns.timeout!r}\n")
        return EXIT_USAGE

    opts  = Options(timeout=timeout_secs, verbose=ns.verbose)
    urls  = ns.urls
    count = ns.count

    workers = ns.concurrency
    if workers <= 0:
        workers = min(len(urls), 8)
    workers = max(1, min(workers, len(urls)))

    def _probe(url: str) -> tuple[list[Result], list[Result]]:
        return _run_samples(url, count, opts)

    with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as ex:
        all_results = list(ex.map(_probe, urls))

    worst      = EXIT_OK
    json_items = []

    for i, (url, (succeeded, failed)) in enumerate(zip(urls, all_results)):
        for r in failed:
            c = _phase_code(r.fail_phase)
            if c > worst:
                worst = c
        if ns.fail:
            for r in succeeded:
                if r.status_code >= 400:
                    worst = max(worst, EXIT_HTTP)

        last_detail = succeeded[-1].detail if succeeded else (
            failed[-1].detail if failed else None
        )

        if ns.json_out:
            json_items.append(_build_json_entry(url, succeeded, failed, last_detail))
            continue

        if i > 0:
            stdout.write("\n")
        _print_url(url, succeeded, failed, count, ns.fail, stdout)

    if ns.json_out:
        stdout.write(json.dumps(json_items, indent=2) + "\n")

    return worst