"""``simurg pgx db build`` — the PGx allele-definition database (spec PGX-001 … PGX-005, PGX-010, IN-053, CLI-035).

I/O and orchestration; the parsing and normalisation live in ``simurg.pgx.db``. Sources are the dated folders written
by ``scripts/pgx/fetch_sources.sh`` (``cpic-YYYYMMDD/`` with ``source.json``). The build is deterministic: rows are
inserted in sorted order into a fresh SQLite file with fixed pragmas, and the file is named by its content hash.
"""

from __future__ import annotations

import hashlib
import json
import sqlite3
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from simurg import __version__
from simurg.core.errors import ExitCode, SimurgError
from simurg.core.log import Reporter

DB_SCHEMA = 1
CPIC_TABLES = ("gene", "allele", "allele_definition", "allele_location_value", "sequence_location", "gene_result",
               "gene_result_lookup")  # fmt: skip
GENE_PAD = 50_000  # probes within this distance of a gene's defining positions are resolved for probe_map

SCHEMA_SQL = """
CREATE TABLE meta(key TEXT PRIMARY KEY, value TEXT NOT NULL);
CREATE TABLE genes(gene TEXT PRIMARY KEY, result_basis TEXT NOT NULL, chrom TEXT);
CREATE TABLE alleles(gene TEXT, allele TEXT, kind TEXT, function TEXT, activity TEXT, structural INTEGER,
                     defined INTEGER, PRIMARY KEY(gene, allele));
CREATE TABLE allele_variants(gene TEXT, allele TEXT, build TEXT, chrom TEXT, pos INTEGER, ref TEXT, alt TEXT,
                             hgvs TEXT, rsid TEXT, location_id INTEGER, value TEXT);
CREATE TABLE allele_wildcards(gene TEXT, allele TEXT, location_id INTEGER);
CREATE TABLE sites(gene TEXT, location_id INTEGER, chrom TEXT, hgvs TEXT, rsid TEXT, ref_value TEXT,
                   PRIMARY KEY(gene, location_id));
CREATE TABLE site_values(gene TEXT, location_id INTEGER, value TEXT, pos INTEGER, ref TEXT, alt TEXT);
CREATE TABLE probe_map(manifest TEXT, probe_name TEXT, chrom TEXT, pos INTEGER, ref TEXT, alt TEXT);
CREATE TABLE masks(manifest TEXT, probe_name TEXT, reason TEXT);
CREATE TABLE limitations(manifest TEXT, gene TEXT, allele TEXT, status TEXT, reported_as TEXT);
CREATE TABLE frequencies(gene TEXT, allele TEXT, population TEXT, freq REAL);
CREATE TABLE phenotypes(guideline TEXT, gene TEXT, lookup_key TEXT, activity_score TEXT, phenotype TEXT);
CREATE INDEX av_pos ON allele_variants(chrom, pos);
CREATE INDEX pm_pos ON probe_map(manifest, chrom, pos);
"""


@dataclass
class PgxDbOptions:
    sources: Path
    fasta: Path
    out: Path
    manifests: list[Path] = field(default_factory=list)
    csvs: list[Path] = field(default_factory=list)
    masks: list[Path] = field(default_factory=list)
    force: bool = False


def _latest(sources: Path, prefix: str) -> Path | None:
    cands = sorted(p for p in sources.glob(f"{prefix}-*") if (p / "source.json").exists())
    return cands[-1] if cands else None


def _load_cpic(d: Path) -> dict[str, list[dict[str, Any]]]:
    t = {}
    for name in CPIC_TABLES:
        p = d / f"{name}.json"
        if not p.exists():
            raise SimurgError("E_PGX_SOURCE", f"{p} missing", hint="run scripts/pgx/fetch_sources.sh")
        t[name] = json.loads(p.read_text(encoding="utf-8"))
    return t


def _probe_sites(
    bpm_path: Path,
    csv_path: Path | None,
    genome: Any,
    cmap: Any,
    windows: dict[str, list[tuple[int, int]]],
    rep: Reporter,
) -> list[Any]:
    """Plus-strand, normalised sites of the manifest probes inside the gene windows (the SNV VCF representation)."""
    from simurg.formats.bpm import read_bpm
    from simurg.harmonize.sites import ManifestView, map_locus

    bpm = read_bpm(bpm_path)
    csvm = None
    if csv_path is not None:
        from simurg.formats.csv_manifest import read_csv_manifest

        csvm = read_csv_manifest(csv_path)
    elif any("D" in s for s in bpm.snps):
        rep.warn("W_SKIP_INDEL_CONTEXT", f"{bpm_path.name}: no CSV manifest, indel probes are not mapped")
    src = csvm or bpm
    mv = ManifestView(bpm.names, bpm.snps, bpm.chroms, bpm.map_infos.tolist(), src.ilmn_strands, src.source_strands,
                      src.ref_strands if any(src.ref_strands) else bpm.ref_strands,
                      csvm.source_seqs if csvm else [""] * bpm.num_loci, bpm.assay_types.tolist(),
                      bpm.intensity_only.tolist())  # fmt: skip
    out = []
    for i in range(bpm.num_loci):
        c = cmap.get(bpm.chroms[i])
        if c is None or c not in windows:
            continue
        p = int(bpm.map_infos[i])
        if not any(a - GENE_PAD <= p <= b + GENE_PAD for a, b in windows[c]):
            continue
        site, _ = map_locus(i, mv, genome.fetch, cmap.get, genome.build or "GRCh38")
        if site is not None:
            out.append(site)
    return out


def run_db_build(o: PgxDbOptions, rep: Reporter) -> int:
    """``simurg pgx db build`` (CLI-035)."""
    from simurg.formats.bpm import read_bpm
    from simurg.formats.fasta import ContigMap, Genome
    from simurg.pgx.db import cpic_definitions, full_hgvs, result_basis

    cdir = _latest(o.sources, "cpic")
    if cdir is None:
        raise SimurgError("E_PGX_SOURCE", f"no cpic-YYYYMMDD source folder in {o.sources}",
                          hint="run scripts/pgx/fetch_sources.sh")  # fmt: skip
    csrc = json.loads((cdir / "source.json").read_text())
    genome = Genome(o.fasta)
    if genome.build not in ("GRCh38", "unknown", ""):
        raise SimurgError("E_BUILD_MISMATCH", f"the PGx DB is GRCh38 only (FASTA is {genome.build})")  # PCN-015
    cmap = ContigMap([e.name for e in genome.entries], "auto")
    t = _load_cpic(cdir)
    defs, sites, problems = cpic_definitions(t, genome.fetch, cmap.get)
    for p in problems:
        rep.info(f"definition note: {p}")
    rows: dict[str, list[tuple[Any, ...]]] = {k: [] for k in ("genes", "alleles", "allele_variants", "allele_wildcards",
                                                             "sites", "site_values", "frequencies", "phenotypes",
                                                             "probe_map", "masks")}  # fmt: skip
    gene_chr = {g["symbol"]: g.get("chr") for g in t["gene"]}
    for gene, ds in sorted(defs.items()):
        rows["genes"].append((gene, result_basis(gene, [d.name for d in ds]), gene_chr.get(gene)))
        for d in ds:
            kind = "reference" if d.reference else ("star_allele" if d.name.startswith("*") else "named_haplotype")
            rows["alleles"].append(
                (gene, d.name, kind, d.function, d.activity, int(d.structural), int(not d.undefined))
            )
            for lid, val in d.variants:
                s = sites.get(lid)
                if s is None or val not in s.alleles:
                    continue
                pos, ref, alt = s.alleles[val]
                rows["allele_variants"].append(
                    (gene, d.name, "GRCh38", s.chrom, pos, ref, alt, full_hgvs(s.chrom, s.hgvs), s.rsid, lid, val)
                )
            for lid in d.wildcards:
                rows["allele_wildcards"].append((gene, d.name, lid))
            for pop, f in sorted(d.frequencies.items()):
                rows["frequencies"].append((gene, d.name, pop, f))
    for lid, s in sorted(sites.items()):
        rows["sites"].append((s.gene, lid, s.chrom, full_hgvs(s.chrom, s.hgvs), s.rsid, s.ref_value))
        for val, (pos, ref, alt) in sorted(s.alleles.items()):
            rows["site_values"].append((s.gene, lid, val, pos, ref, alt))
    results = {r["id"]: r for r in t["gene_result"]}
    for lk in t["gene_result_lookup"]:
        res = results.get(lk["phenotypeid"])
        if res is None:
            continue
        key = json.dumps(lk.get("lookupkey"), sort_keys=True, ensure_ascii=False)
        rows["phenotypes"].append(("cpic", res["genesymbol"], key, lk.get("totalactivityscore"), res["result"]))
    # probe map (PGX-003): manifest probes whose normalised site carries a DB variant
    windows: dict[str, list[tuple[int, int]]] = {}
    for av in rows["allele_variants"]:
        windows.setdefault(av[3], []).append((av[4], av[4] + len(av[5])))
    var_keys = {(av[3], av[4], av[5], av[6]) for av in rows["allele_variants"]}
    manifests = []
    csvs = dict(zip(o.manifests, o.csvs, strict=False))
    for bpm_path in o.manifests:
        name = read_bpm(bpm_path).manifest_name
        manifests.append(name)
        n = 0
        for ps in _probe_sites(bpm_path, csvs.get(bpm_path), genome, cmap, windows, rep):
            for alt in ps.alleles:
                if alt != ps.ref and (ps.contig, ps.pos, ps.ref, alt) in var_keys:
                    rows["probe_map"].append((name, ps.name, ps.contig, ps.pos, ps.ref, alt))
                    n += 1
        rep.info(f"{name}: {n} probe-to-variant link(s)")
    for mp in o.masks:
        for ln in mp.read_text(encoding="utf-8").splitlines():
            if ln.strip() and not ln.startswith("#"):
                cols = ln.split("\t")
                rows["masks"].append((cols[2] if len(cols) > 2 else mp.stem, cols[0], cols[1] if len(cols) > 1 else ""))
    meta = {"schema": "simurg.pgx_db", "schema_version": str(DB_SCHEMA), "builder": f"SimurgArray {__version__}",
            "genome_build": "GRCh38", "fasta_sha256": genome.sha256, "sqlite_version": sqlite3.sqlite_version,
            "allele_sources": json.dumps([{"name": "CPIC", "version": csrc["version"], "retrieved": csrc["retrieved"],
                                           "license": csrc.get("license")}]),
            "phenotype_sources": json.dumps([{"name": "CPIC", "version": csrc["version"]}]),
            "cpic_files_sha256": json.dumps(csrc.get("files", {}), sort_keys=True),
            "manifests": json.dumps(sorted(manifests)), "definition_notes": json.dumps(problems)}  # fmt: skip
    # content hash over the sorted rows and meta (IN-053: the name and meta carry it; a file cannot hold its own hash)
    for rs in rows.values():
        rs.sort(key=lambda r: tuple(str(x) for x in r))
    sha = hashlib.sha256(json.dumps({"rows": rows, "meta": sorted(meta.items())}, sort_keys=True,
                                    default=str).encode()).hexdigest()  # fmt: skip
    meta["content_sha256"] = sha
    o.out.mkdir(parents=True, exist_ok=True)
    tmp = o.out / ".pgx_db.part.sqlite"
    tmp.unlink(missing_ok=True)
    con = sqlite3.connect(tmp)
    con.execute("PRAGMA page_size=4096")
    con.execute("PRAGMA journal_mode=OFF")
    con.executescript(SCHEMA_SQL)
    for table, rs in rows.items():
        if rs:
            ph = ",".join("?" * len(rs[0]))
            con.executemany(f"INSERT INTO {table} VALUES ({ph})", rs)
    con.executemany("INSERT INTO meta VALUES (?, ?)", sorted(meta.items()))
    con.commit()
    con.execute("VACUUM")
    con.close()
    final = o.out / f"pgx_db-{DB_SCHEMA}-{sha[:12]}.sqlite"
    if final.exists() and not o.force:
        tmp.unlink()
        rep.info(f"identical PGx DB already present: {final}")
        return ExitCode.OK
    tmp.replace(final)
    rep.info(f"PGx DB: {len(defs)} genes, {len(rows['alleles'])} alleles, {len(rows['allele_variants'])} variant "
             f"rows, {len(rows['probe_map'])} probe links → {final} (content sha256 {sha})")  # fmt: skip
    return ExitCode.OK


# ------------------------------------------------------------------------------------------------ coverage
def coverage_rows(db: Path, manifest: str) -> list[dict[str, Any]]:
    """PGX-009 / PGX-035: per allele, callable | not_covered | masked | equivalent_to:<allele> on a manifest."""
    import numpy as np

    from simurg.pipeline.pgx_star import load_models

    con = sqlite3.connect(db)
    links: dict[tuple[str, int, str, str], list[str]] = {}
    for probe, c, p, r, a in con.execute("SELECT probe_name, chrom, pos, ref, alt FROM probe_map WHERE manifest = ?",
                                         (manifest,)):  # fmt: skip
        links.setdefault((c, p, r, a), []).append(probe)
    masked = {p for (p,) in con.execute("SELECT probe_name FROM masks WHERE manifest = ?", (manifest,))}
    basis = dict(con.execute("SELECT gene, result_basis FROM genes").fetchall())
    kinds = {(g, a): (k, d) for g, a, k, d in con.execute("SELECT gene, allele, kind, defined FROM alleles")}
    con.close()
    out = []
    for gene, m in sorted(load_models(db, "global").items()):
        usable = np.array([bool(links.get(v)) and not all(pr in masked for pr in links[v]) for v in m.variants])
        masked_v = np.array([bool(links.get(v)) and all(pr in masked for pr in links[v]) for v in m.variants])
        sig: dict[bytes, list[int]] = {}
        for h in range(len(m.haplotypes)):
            sig.setdefault(np.where(usable, m.carries[h], -9).tobytes(), []).append(h)
        rep = {}
        for members in sig.values():
            best = max(members, key=lambda h: (m.freq[h], -h))
            for h in members:
                rep[h] = best
        for h, name in enumerate(m.haplotypes):
            defining = m.carries[h] == 1
            probes = sorted({pr for k in np.flatnonzero(defining).tolist() for pr in links.get(m.variants[k], [])})
            if defining.any() and not (defining & usable).any():
                status = "masked" if (defining & masked_v).any() else "not_covered"
            elif rep[h] != h:
                status = f"equivalent_to:{m.haplotypes[rep[h]]}"
            else:
                status = "callable"
            out.append({"gene": gene, "allele": name, "allele_kind": kinds.get((gene, name), ("", 1))[0],
                        "result_basis": basis.get(gene, ""), "status": status,
                        "covering_probes": ";".join(probes)})  # fmt: skip
    con = sqlite3.connect(db)
    for gene, allele, kind in con.execute("SELECT gene, allele, kind FROM alleles WHERE defined = 0 ORDER BY 1, 2"):
        out.append({"gene": gene, "allele": allele, "allele_kind": kind, "result_basis": basis.get(gene, ""),
                    "status": "not_covered", "covering_probes": ""})  # fmt: skip
    con.close()
    return out


def run_coverage(bpm: Path, db: Path, out: Path, rep: Reporter) -> int:
    """``simurg pgx coverage --bpm --db`` (PGX-035)."""
    from simurg.formats.bpm import read_bpm
    from simurg.pipeline.pgx_pheno import db_meta, tsv_header, write_meta_tsv

    name = read_bpm(bpm).manifest_name
    rows = coverage_rows(db, name)
    if not any(r["covering_probes"] for r in rows):
        rep.warn("W_PGX_NO_PROBE_MAP", f"the PGx DB has no probe map for {name}; rebuild it with --manifests")
    out.mkdir(parents=True, exist_ok=True)
    cols = ["gene", "allele", "allele_kind", "result_basis", "status", "covering_probes"]
    write_meta_tsv(out / "pgx_coverage.tsv", tsv_header(db, db_meta(db), "none")[:-1], cols, rows)
    counts: dict[str, int] = {}
    for r in rows:
        k = r["status"].split(":")[0]
        counts[k] = counts.get(k, 0) + 1
    rep.info(f"{name}: " + ", ".join(f"{k} {v}" for k, v in sorted(counts.items())))
    return ExitCode.OK