"""``simurg compare A B``: field-level differential testing of GTCs or VCFs (spec CLI-038; plan §14, §16.3, §18 P7).

A and B are two files or two folders; folders are paired by relative path. Outputs: ``summary.json``,
``discrepancies.tsv.gz`` (every discrepancy, never truncated) and ``dashboard.html``.
"""

from __future__ import annotations

import gzip
import hashlib
import html
import io
import json
import multiprocessing as mp
import os
from collections.abc import Iterator, Sequence
from concurrent.futures import ProcessPoolExecutor
from dataclasses import dataclass, field
from pathlib import Path
from typing import IO, Any, cast

import numpy as np

from simurg import PRODUCT_NAME, __version__
from simurg.compare.core import CATEGORY_NAMES, GTC_TOLERANCES, FieldResult, Tolerance, locus_classes, merge_results
from simurg.compare.records import Comparison, VcfRecord, compare_gtc, compare_vcf_records
from simurg.core.atomic import atomic_open, atomic_write_text, check_outputs_absent
from simurg.core.errors import ExitCode, SimurgError
from simurg.core.log import Reporter

GTC_EXT = (".gtc",)
VCF_EXT = (".vcf", ".vcf.gz")


@dataclass
class CompareOptions:
    a: Path
    b: Path
    out: Path
    bpm: Path | None = None
    tolerances: dict[str, str] = field(default_factory=dict)  # field -> "abs=1e-5" | "ulp=2" | "exact"
    samples: list[str] | None = None
    threads: int = 0
    fail_on_discrepancy: bool = False
    title: str = ""
    force: bool = False


def parse_tolerance(text: str) -> Tolerance:
    text = text.strip().lower()
    if text in ("exact", "0"):
        return Tolerance()
    kv = dict(part.split("=", 1) for part in text.split(","))
    try:
        return Tolerance(abs=float(kv["abs"]) if "abs" in kv else None, ulp=int(kv["ulp"]) if "ulp" in kv else None)
    except (KeyError, ValueError):
        raise SimurgError("E_CONFIG_KEY", f"bad tolerance {text!r}; use exact, abs=X, ulp=N or both") from None


def _kind(p: Path) -> str:
    n = p.name.lower()
    if n.endswith(GTC_EXT):
        return "gtc"
    if n.endswith(VCF_EXT):
        return "vcf"
    raise SimurgError("E_FILE_UNREADABLE", f"{p}: compare accepts .gtc and .vcf(.gz) files (convert BCF with bcftools)")


def pair_inputs(a: Path, b: Path) -> tuple[list[tuple[str, Path, Path]], list[str]]:
    """((label, file_a, file_b)…, unmatched notes). Folders are paired by file name (or relative path)."""
    for p in (a, b):
        if not p.exists():
            raise SimurgError("E_FILE_NOT_FOUND", f"not found: {p}", path=str(p))
    if a.is_file() and b.is_file():
        if _kind(a) != _kind(b):
            raise SimurgError("E_FILE_UNREADABLE", f"cannot compare a {_kind(a)} with a {_kind(b)}")
        return [(a.name, a, b)], []
    if not (a.is_dir() and b.is_dir()):
        raise SimurgError("E_FILE_UNREADABLE", "give two files or two folders")

    def listing(root: Path) -> dict[str, Path]:
        """Key files by name when names are unique in the tree (so nested and flat layouts pair up),
        otherwise by relative path."""
        files = [p for p in sorted(root.rglob("*")) if p.is_file() and p.name.lower().endswith(GTC_EXT + VCF_EXT)]
        names = [p.name for p in files]
        if len(set(names)) == len(names):
            return {p.name: p for p in files}
        return {str(p.relative_to(root)): p for p in files}

    la, lb = listing(a), listing(b)
    pairs = [(k, la[k], lb[k]) for k in sorted(la) if k in lb]
    notes = [f"only in A: {k}" for k in sorted(set(la) - set(lb))] + [
        f"only in B: {k}" for k in sorted(set(lb) - set(la))
    ]
    if not pairs:
        raise SimurgError("E_NO_INPUT", f"no files with the same relative path under {a} and {b}")
    return pairs, notes


# ------------------------------------------------------------------------------------------------ workers
_W: dict[str, Any] = {}


def _init(bpm: str | None, tolerances: dict[str, Tolerance], samples: list[str] | None) -> None:
    _W.update(bpm=bpm, tol=tolerances, samples=samples)


def _manifest() -> tuple[Any, Any]:
    if "classes" not in _W:
        if _W["bpm"]:
            from simurg.formats.bpm import read_bpm

            m = read_bpm(_W["bpm"])
            _W["classes"] = locus_classes(m.chroms, m.snps, m.assay_types, m.intensity_only)
            _W["names"] = list(m.names)
        else:
            _W["classes"], _W["names"] = None, None
    return _W["classes"], _W["names"]


VcfKey = tuple[str, int, str, str, str]


def read_vcf(path: Path, sample: str | None = None) -> tuple[list[str], dict[str, dict[VcfKey, VcfRecord]], int]:
    """Records per sample (all samples, or one), keyed by (CHROM, POS, REF, ALT, ID). Repeated keys are kept under
    a numbered ID so nothing is silently overwritten; their count is returned."""
    opener = gzip.open if str(path).endswith(".gz") else open
    samples: list[str] = []
    out: dict[str, dict[VcfKey, VcfRecord]] = {}
    seen: set[VcfKey] = set()
    dups = 0
    with opener(path, "rt") as fh:
        for line in fh:
            if line.startswith("##"):
                continue
            f = line.rstrip("\n").split("\t")
            if line.startswith("#"):
                samples = f[9:]
                keep = [i for i, s in enumerate(samples) if sample is None or s == sample]
                out = {samples[i]: {} for i in keep}
                continue
            keys = f[8].split(":")
            k = (f[0], int(f[1]), f[3], f[4], f[2])
            if k in seen:
                dups += 1
                k = (f[0], int(f[1]), f[3], f[4], f"{f[2]}#{dups}")
            seen.add(k)
            for i in keep:
                vals = f[9 + i].split(":")
                out[samples[i]][k] = VcfRecord(f[0], int(f[1]), f[2], f[3], f[4], dict(zip(keys, vals, strict=False)))
    return samples, out, dups


def _compare_pair(task: tuple[str, str, str]) -> list[Comparison]:
    label, pa, pb = task
    kind = _kind(Path(pa))
    try:
        if kind == "gtc":
            from simurg.formats.gtc import read_gtc

            ga, gb = read_gtc(pa, check_write_complete=False), read_gtc(pb, check_write_complete=False)
            classes, names = _manifest()
            if classes is None or len(names) != ga.num_snps:
                classes = locus_classes(["1"] * ga.num_snps, ["[A/B]"] * ga.num_snps, np.zeros(ga.num_snps),
                                        np.zeros(ga.num_snps, np.bool_))  # fmt: skip
                classes = {"all": classes["all"]}
                names = [str(i) for i in range(ga.num_snps)]
            tol = {**GTC_TOLERANCES, **_W["tol"]}
            return [compare_gtc(ga, gb, label, classes, names, tol)]
        sa, ra, da = read_vcf(Path(pa))
        sb, rb, db = read_vcf(Path(pb))
        wanted = _W["samples"] or [s for s in sa if s in sb]
        out = []
        for s in wanted:
            if s not in ra or s not in rb:
                c = Comparison(f"{label}:{s}", "vcf")
                c.notes.append(f"sample {s} missing in {'A' if s not in ra else 'B'}")
                out.append(c)
                continue
            ftol = _W["tol"].get("float")
            if ftol is None:
                out.append(compare_vcf_records(ra[s], rb[s], f"{label}:{s}"))
            else:
                out.append(compare_vcf_records(ra[s], rb[s], f"{label}:{s}", ftol))
        if not wanted:
            c = Comparison(label, "vcf")
            c.notes.append("no sample in common")
            out.append(c)
        if da or db:
            out[0].notes.append(f"duplicate record keys: {da} in A, {db} in B")
        return out
    except SimurgError as e:
        c = Comparison(label, kind)
        c.notes.append(f"{e.code}: {e.message}")
        return [c]


# ------------------------------------------------------------------------------------------------ outputs
def _totals(fields: dict[str, FieldResult]) -> dict[str, Any]:
    out = {}
    for f, r in fields.items():
        out[f] = {
            "tolerance": r.tolerance,
            "counts": {cls: dict(zip(CATEGORY_NAMES, c, strict=True)) for cls, c in sorted(r.counts.items())},
            "discrepancies": sum(c[2] for k, c in r.counts.items() if k in ("all", "sample", "bins", "controls")),
            "listed": len(r.discrepancies),
        }
    return out


def write_discrepancies(path: Path, rows: Iterator[dict[str, Any]]) -> int:
    n = 0
    with atomic_open(path, "wb") as raw:
        gz = gzip.GzipFile(fileobj=cast(IO[bytes], raw), mode="wb", mtime=0)
        fh = io.TextIOWrapper(gz, encoding="utf-8", newline="\n")
        fh.write("#simurg_schema=simurg.compare_discrepancies/1\npair\tfield\tlocus\ta\tb\n")
        for r in rows:
            fh.write(
                "\t".join(str(r.get(k, "")).replace("\t", " ") for k in ("pair", "field", "locus", "a", "b")) + "\n"
            )
            n += 1
        fh.flush()
        fh.detach()
        gz.close()
    return n


DASH_CSS = """
:root{--bg:#fafafa;--fg:#222;--muted:#666;--line:#ddd;--ok:#d6e6f5;--near:#e8f2d6;--bad:#f9d3ac;--unk:#e6e6e6}
@media (prefers-color-scheme: dark){:root{--bg:#1d1f21;--fg:#e6e6e6;--muted:#aaa;--line:#444;
  --ok:#23405c;--near:#35451f;--bad:#6b3d10;--unk:#3a3a3a}}
body{margin:0;background:var(--bg);color:var(--fg);font:14px/1.4 "Noto Sans",system-ui,sans-serif}
.sa-main{padding:16px}
table{border-collapse:collapse;margin:6px 0}
td,th{border:1px solid var(--line);padding:3px 8px;text-align:right}
td:first-child,th:first-child{text-align:left}
.cEXACT{background:var(--ok)} .cNEAR{background:var(--near)} .cUNKNOWN{background:var(--unk)}
tr.bad .cDISCREPANCY{background:var(--bad)}
.verdict{font-size:20px;font-weight:600} .muted{color:var(--muted)} .wrap{overflow-x:auto}
"""


def _table(head: Sequence[str], rows: Sequence[Sequence[str]], cls: str = "") -> str:
    th = "".join(f"<th>{h}</th>" for h in head)
    body = "".join("<tr>" + "".join(f"<td>{c}</td>" for c in r) + "</tr>" for r in rows)
    return f'<div class="wrap"><table class="{cls}"><tr>{th}</tr>{body}</table></div>'


def _counts_cells(c: dict[str, int] | list[int]) -> str:
    vals = c if isinstance(c, list) else [c.get(k, 0) for k in CATEGORY_NAMES]
    return "".join(f'<td class="c{k}">{v:,}</td>' for k, v in zip(CATEGORY_NAMES, vals, strict=True))


def render_dashboard(summary: dict[str, Any], examples: dict[str, list[dict[str, Any]]]) -> str:
    esc = html.escape
    title = esc(summary["title"] or "Differential test")
    verdict = summary["verdict"]
    head = "".join(f"<th>{k}</th>" for k in CATEGORY_NAMES)
    rows = []
    for f, d in summary["fields"].items():
        counts = d["counts"]
        main: dict[str, int] = next((counts[k] for k in ("all", "sample", "bins", "controls") if k in counts), {})
        per = "".join(f"<tr><td>{esc(cls)}</td>{_counts_cells(c)}</tr>" for cls, c in counts.items())
        inner = f"<table><tr><th>class</th>{head}</tr>{per}</table>"
        bad = ' class="bad"' if d["discrepancies"] else ""
        rows.append(
            f"<tr{bad}><td><details><summary>{esc(f)}</summary>{inner}</details></td>"
            f"<td>{esc(d['tolerance'])}</td>{_counts_cells(main)}</tr>"
        )
    fields_html = (
        f'<div class="wrap"><table><tr><th>field (click for classes)</th><th>tolerance</th>{head}</tr>'
        + "".join(rows)
        + "</table></div>"
    )
    parts = [
        f'<h1>{title}</h1><div class="muted">{PRODUCT_NAME} {__version__} · {summary["pairs"]} pair(s)</div>',
        f'<div class="muted">A = {esc(summary["a"])}<br>B = {esc(summary["b"])}</div>',
        f'<p class="verdict">{verdict}: {summary["discrepancies"]:,} discrepancies, {summary["unknown"]:,} unknown</p>',
        "<h2>Fields</h2>" + fields_html,
    ]
    meta = [
        [esc(p), esc(k), esc(v[0]), esc(v[1])]
        for p, md in summary["metadata_differences"].items()
        for k, v in md.items()
    ]
    if meta:
        parts.append("<h2>Metadata differences (not counted)</h2>" + _table(["pair", "field", "A", "B"], meta))
    if summary["notes"]:
        parts.append("<h2>Notes</h2><ul>" + "".join(f"<li>{esc(n)}</li>" for n in summary["notes"]) + "</ul>")
    for f, lst in examples.items():
        if lst:
            n = summary["fields"][f]["listed"]
            body = [[esc(str(r.get(k, "")))[:200] for k in ("pair", "locus", "a", "b")] for r in lst]
            parts.append(
                f"<h3>{esc(f)}: first {len(lst)} of {n:,} discrepancies</h3>"
                + _table(["pair", "locus", "A", "B"], body)
            )
    parts.append('<p class="muted">Every discrepancy and every unmatched record is listed in discrepancies.tsv.gz.</p>')
    from simurg.core.brand import BAR_CSS, favicon_link, report_bar

    checksum = hashlib.sha256(json.dumps(summary, sort_keys=True, default=str).encode("utf-8")).hexdigest()
    bar = report_bar("Differential test", __version__, "summary sha256", checksum)
    return (
        '<!doctype html><html lang="en"><head><meta charset="utf-8">'
        '<meta name="viewport" content="width=device-width, initial-scale=1">'
        f"<title>{title} · {PRODUCT_NAME}</title>{favicon_link()}<style>{DASH_CSS}{BAR_CSS}</style></head><body>"
        + bar
        + '<div class="sa-main">'
        + "\n".join(parts)
        + "</div></body></html>\n"
    )


def run_compare(opts: CompareOptions, rep: Reporter) -> int:
    targets = [opts.out / "summary.json", opts.out / "discrepancies.tsv.gz", opts.out / "dashboard.html"]
    check_outputs_absent(targets, force=opts.force)
    tol = {k: parse_tolerance(v) for k, v in opts.tolerances.items()}
    unknown = [k for k in tol if k not in GTC_TOLERANCES and k != "float"]
    if unknown:
        valid = ", ".join(sorted(GTC_TOLERANCES))
        raise SimurgError("E_CONFIG_KEY", f"unknown tolerance field(s) {unknown}; GTC fields: {valid}; VCF: float")
    pairs, notes = pair_inputs(opts.a, opts.b)
    tasks = [(label, str(pa), str(pb)) for label, pa, pb in pairs]
    workers = max(1, min(opts.threads or (os.cpu_count() or 1), len(tasks)))
    rep.info(f"comparing {len(tasks)} pair(s) with {workers} worker(s)")
    if workers == 1:
        _init(str(opts.bpm) if opts.bpm else None, tol, opts.samples)
        results = [c for t in tasks for c in _compare_pair(t)]
    else:
        with ProcessPoolExecutor(workers, mp_context=mp.get_context("spawn"), initializer=_init,
                                 initargs=(str(opts.bpm) if opts.bpm else None, tol, opts.samples)) as ex:  # fmt: skip
            results = [c for batch in ex.map(_compare_pair, tasks) for c in batch]
    merged = merge_results([f for c in results for f in c.fields])
    for c in results:
        notes += [f"{c.label}: {n}" for n in c.notes]

    def main_counts(r: FieldResult) -> list[int]:
        for k in ("all", "sample", "bins", "controls"):
            if k in r.counts:
                return r.counts[k]
        return [0, 0, 0, 0]

    total_disc = sum(main_counts(r)[2] for r in merged.values())
    total_unknown = sum(main_counts(r)[3] for r in merged.values())
    summary = {
        "schema": "simurg.compare/1",
        "title": opts.title,
        "a": str(opts.a),
        "b": str(opts.b),
        "pairs": len(results),
        "discrepancies": total_disc,
        "unknown": total_unknown,
        "verdict": "DISCREPANCIES" if total_disc else ("UNKNOWNS" if total_unknown else "PASS"),
        "fields": _totals(merged),
        "metadata_differences": {c.label: c.metadata for c in results if c.metadata},
        "notes": notes,
    }
    opts.out.mkdir(parents=True, exist_ok=True)
    write_discrepancies(opts.out / "discrepancies.tsv.gz", (d for r in merged.values() for d in r.discrepancies))
    atomic_write_text(opts.out / "summary.json", json.dumps(summary, indent=2) + "\n")
    examples = {f: r.discrepancies[:50] for f, r in merged.items()}
    atomic_write_text(opts.out / "dashboard.html", render_dashboard(summary, examples))
    rep.info(f"{summary['verdict']}: {total_disc} discrepancies, {total_unknown} unknown in {len(results)} pair(s); "
             f"report {opts.out / 'dashboard.html'}")  # fmt: skip
    if total_disc and opts.fail_on_discrepancy:
        return ExitCode.PARTIAL
    return ExitCode.OK