"""PGx allele definitions from public sources → normalised VCF-style variants — spec PGX-001 … PGX-005.
Pure functions: the reference sequence comes in as a ``fetch(contig, start0, end0)`` callable.

CPIC tables (CC0, https://api.cpicpgx.org) [DB: CPIC allele definition, allele function, diplotype→phenotype]:
``sequence_location`` gives each defining position as HGVS on the chromosome (``g.``/``m.``); ``allele_location_value``
gives each allele's value there. Values are literal sequences (``A``, ``AA`` for a dup, ``TA(7)`` for a repeat) or
``del…``/``ins…`` forms; the reference allele definition carries the reference value. Every value becomes a
replacement of a reference window, then a left-normalised, anchored VCF (POS, REF, ALT) — the representation the
SNV VCFs use (IN-046), so definitions and array records compare directly.
"""

from __future__ import annotations

import re
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any

Fetch = Callable[[str, int, int], str]
HGVS_RE = re.compile(r"^([gm])\.(\d+)(?:_(\d+))?(.*)$")
REPEAT_RE = re.compile(r"^([ACGT]+)\((\d+)\)$")
IUPAC_WILDCARD = set("RYSWKMBDHVN")
# GRCh38 RefSeq chromosome accessions, for full HGVS (PGX-028 "NC_000022.11:g.42126611C>G")
REFSEQ_GRCH38 = {"1": "NC_000001.11", "2": "NC_000002.12", "3": "NC_000003.12", "4": "NC_000004.12",
                 "5": "NC_000005.10", "6": "NC_000006.12", "7": "NC_000007.14", "8": "NC_000008.11",
                 "9": "NC_000009.12", "10": "NC_000010.11", "11": "NC_000011.10", "12": "NC_000012.12",
                 "13": "NC_000013.11", "14": "NC_000014.9", "15": "NC_000015.10", "16": "NC_000016.10",
                 "17": "NC_000017.11", "18": "NC_000018.10", "19": "NC_000019.10", "20": "NC_000020.11",
                 "21": "NC_000021.9", "22": "NC_000022.11", "X": "NC_000023.11", "Y": "NC_000024.10",
                 "M": "NC_012920.1", "MT": "NC_012920.1"}  # fmt: skip


def full_hgvs(chrom: str, hgvs: str) -> str:
    """``g.…``/``m.…`` → ``NC_…:g.…`` on the GRCh38 chromosome accession (unchanged when unknown)."""
    acc = REFSEQ_GRCH38.get(chrom.removeprefix("chr"))
    return (
        f"{acc}:{hgvs.split(';')[0].strip()}" if acc and not hgvs.startswith("NC_") else hgvs
    )  # CPIC core-allele definitions: the position does not constrain the allele


@dataclass(frozen=True)
class Window:
    """Reference interval [start, end] (1-based, inclusive) replaced by allele values; ``insertion`` windows are the
    single anchor base S of ``g.S_EinsX`` (values are anchor + inserted sequence)."""

    start: int
    end: int
    kind: str  # snv | del | ins | dup | repeat | delins
    unit: str = ""  # repeat unit


def parse_location(hgvs: str) -> Window | None:
    """First HGVS expression of a CPIC ``chromosomelocation`` (several are ``;``-joined for one site)."""
    first = hgvs.split(";")[0].strip()
    m = HGVS_RE.match(first)
    if not m:
        return None
    s = int(m.group(2))
    e = int(m.group(3)) if m.group(3) else s
    rest = m.group(4)
    if re.fullmatch(r"[ACGT]>[ACGT]", rest):
        return Window(s, s, "snv")
    if rest.startswith("delins"):
        return Window(s, e, "delins")
    if rest.startswith("del"):
        return Window(s, e, "del")
    if rest.startswith("ins"):
        return Window(s, s, "ins")
    if rest.startswith("dup"):
        return Window(s, e, "dup")
    rm = re.fullmatch(r"([ACGT]+)\[(\d+)\]", rest)
    if rm:
        return Window(s, s, "repeat", rm.group(1))
    return None


def value_sequence(value: str, w: Window, fetch: Fetch, contig: str) -> tuple[int, int, str] | None:
    """(start, end, sequence) that the value puts in place of the reference window; None = unusable/ambiguous."""
    v = value.strip()
    if " or " in v or ";" in v or not v:
        return None
    if w.kind == "repeat":
        m = REPEAT_RE.match(v)
        unit = w.unit
        n_ref = 0
        while fetch(contig, w.start - 1 + n_ref * len(unit), w.start - 1 + (n_ref + 1) * len(unit)) == unit:
            n_ref += 1
        end = w.start - 1 + n_ref * len(unit)
        if m:
            return w.start, end, m.group(1) * int(m.group(2))
        if re.fullmatch(r"[ACGT]+", v):
            return w.start, end, v
        return None
    if w.kind == "ins":
        anchor = fetch(contig, w.start - 1, w.start)
        if v == "del":
            return w.start, w.start, anchor
        if v.startswith("ins"):
            return w.start, w.start, anchor + v[3:]
        return None
    if w.kind == "delins":
        if v.startswith("del"):
            return w.start, w.end, fetch(contig, w.start - 1, w.end)
        if v.startswith("ins"):
            return w.start, w.end, v[3:]
        return None
    if v.startswith("del"):
        return w.start, w.end, ""
    if re.fullmatch(r"[ACGT]+", v):
        return w.start, w.end, v
    return None


def normalize(fetch: Fetch, contig: str, start: int, end: int, alt: str) -> tuple[int, str, str]:
    """Replacement of reference [start, end] (1-based; end = start − 1 for a pure insertion before start) by ``alt``
    → left-normalised, anchored (POS, REF, ALT) [LIT: Tan et al. 2015 Bioinformatics 31:2202 (vt normalize)]."""
    ref = fetch(contig, start - 1, end)
    pos = start
    # trim the common suffix, then the common prefix
    while ref and alt and ref[-1] == alt[-1]:
        ref, alt = ref[:-1], alt[:-1]
    while len(ref) > 0 and len(alt) > 0 and ref[0] == alt[0]:
        ref, alt = ref[1:], alt[1:]
        pos += 1
    if ref and alt and len(ref) == len(alt):
        return pos, ref, alt  # SNV / MNV
    # left shift the indel while the base before it equals its last base
    while pos > 1 and (ref or alt):
        prev = fetch(contig, pos - 2, pos - 1)
        body = ref or alt
        if prev != body[-1] or (ref and alt):
            break
        if ref:
            ref = prev + ref[:-1]
        else:
            alt = prev + alt[:-1]
        pos -= 1
    anchor = fetch(contig, pos - 2, pos - 1)
    return pos - 1, anchor + ref, anchor + alt


@dataclass
class VariantSite:
    location_id: int
    gene: str
    chrom: str  # contig name as in the FASTA
    hgvs: str
    rsid: str | None
    ref_value: str | None
    alleles: dict[str, tuple[int, str, str]] = field(default_factory=dict)  # value → normalised (pos, ref, alt)


@dataclass
class AlleleDef:
    gene: str
    name: str
    function: str | None
    activity: str | None
    reference: bool
    structural: bool
    variants: list[tuple[int, str]] = field(default_factory=list)  # (location_id, value) where ≠ reference
    frequencies: dict[str, float] = field(default_factory=dict)
    wildcards: list[int] = field(default_factory=list)  # location ids with an IUPAC value (either base)

    @property
    def undefined(self) -> bool:
        """A non-reference allele without any defining variant cannot be called (e.g. CPIC NUDT15 *2)."""
        return not self.reference and not self.variants


def cpic_definitions(
    t: dict[str, list[dict[str, Any]]], fetch: Fetch, contig_of: Callable[[str], str | None]
) -> tuple[dict[str, list[AlleleDef]], dict[int, VariantSite], list[str]]:
    """CPIC tables → (gene → allele definitions, location id → site, problems)."""
    problems: list[str] = []
    gene_chr = {g["symbol"]: g.get("chr") for g in t["gene"]}
    locs = {loc["id"]: loc for loc in t["sequence_location"]}
    values: dict[int, list[tuple[int, str]]] = {}
    for v in t["allele_location_value"]:
        values.setdefault(v["alleledefinitionid"], []).append((v["locationid"], v["variantallele"]))
    func = {(a["genesymbol"], a["name"]): a for a in t["allele"]}
    defs: dict[str, list[AlleleDef]] = {}
    sites: dict[int, VariantSite] = {}
    ref_value: dict[int, str] = {}
    for d in sorted(t["allele_definition"], key=lambda d: (d["genesymbol"], d["name"])):
        if d.get("matchesreferencesequence"):
            for lid, val in values.get(d["id"], []):
                ref_value[lid] = val
    for d in sorted(t["allele_definition"], key=lambda d: (d["genesymbol"], d["name"])):
        g = d["genesymbol"]
        a = func.get((g, d["name"]), {})
        freq = {k: float(v) for k, v in (a.get("frequency") or {}).items() if isinstance(v, int | float)}
        ref_def, sv = bool(d.get("matchesreferencesequence")), bool(d.get("structuralvariation"))
        ad = AlleleDef(g, d["name"], a.get("clinicalfunctionalstatus"), a.get("activityvalue"), ref_def, sv, [], freq)
        for lid, val in sorted(values.get(d["id"], [])):
            loc = locs.get(lid)
            if loc is None:
                continue
            if lid not in sites:
                hg = loc.get("chromosomelocation") or ""
                chrom = "chrM" if hg.startswith("m.") else (gene_chr.get(g) or "")
                contig = contig_of(chrom) if chrom else None
                if contig is None:
                    problems.append(f"{g} {loc['name']}: no contig for {chrom!r}")
                    continue
                sites[lid] = VariantSite(lid, g, contig, hg, loc.get("dbsnpid"), ref_value.get(lid))
            site = sites[lid]
            if len(val) == 1 and val in IUPAC_WILDCARD:
                ad.wildcards.append(lid)
            elif val != site.ref_value:
                ad.variants.append((lid, val))
        defs.setdefault(g, []).append(ad)
    for lid, site in sites.items():
        w = parse_location(site.hgvs)
        if w is None:
            problems.append(f"{site.gene} {site.hgvs}: unparsed location")
            continue
        vals = {val for d in defs.get(site.gene, []) for (l2, val) in d.variants if l2 == lid}
        for val in sorted(vals):
            seq = value_sequence(val, w, fetch, site.chrom)
            if seq is None:
                problems.append(f"{site.gene} {site.hgvs}: value {val!r} not representable")
                continue
            s, e, alt = seq
            if w.kind == "ins":
                s, e, alt = s + 1, s, alt[1:]
            site.alleles[val] = normalize(fetch, site.chrom, s, e, alt)
    return defs, sites, problems


def result_basis(gene: str, names: list[str]) -> str:
    """PGX-002 ``genes.result_basis``."""
    if gene in ("HLA-A", "HLA-B"):
        return "tagging_variant"
    if gene in ("G6PD", "APOE"):
        return "haplotype_named"
    if names and all(n.startswith("*") for n in names):
        return "star_allele"
    return "variant_based"


def pharmvar_core_alleles(gene: str, vcfs: dict[str, str]) -> dict[str, list[tuple[str, int, str, str, str | None]]]:
    """PharmVar per-allele GRCh38 VCFs (file name → text) → core alleles (``*N``; sub-alleles ``*N.001`` skipped) and
    their variants (chrom, pos, ref, alt, rsid) as written by PharmVar [DB: PharmVar haplotype VCFs]."""
    out: dict[str, list[tuple[str, int, str, str, str | None]]] = {}
    pat = re.compile(rf"^{re.escape(gene)}_(\d+)\.vcf$")
    for name, text in vcfs.items():
        m = pat.match(name.rsplit("/", 1)[-1])
        if not m:
            continue
        recs = []
        for ln in text.splitlines():
            if not ln or ln.startswith("#"):
                continue
            f = ln.split("\t")
            rsid = f[2] if f[2].startswith("rs") else None
            for alt in f[4].split(","):
                recs.append((f[0], int(f[1]), f[3].upper(), alt.upper(), rsid))
        out[f"*{m.group(1)}"] = recs
    return out


CLINPGX_LEVELS = ("1A", "1B", "2A", "2B")
# PGX-008 parity list of variant-based genes: level-3 annotations are kept for these too
PARITY_VARIANT_GENES = {"DPYD", "RYR1", "OPRM1", "CFTR", "DRD2", "COMT", "OPRD1", "CACNA1C", "HTR2A", "OPRK1",
                        "MT-RNR1", "TNF", "BDNF", "IL6", "HTR2C", "GRIN2B", "MTHFR", "ITGB3", "FKBP5", "ADH1B",
                        "ANKK1", "HMGCR", "ATM", "KIF6", "F2", "CACNA1S", "LPA", "F13A1", "GRIK1", "GRIK4", "YEATS4",
                        "ABCG2", "F5", "ALDH2", "ANK3", "VKORC1"}  # fmt: skip


def clinpgx_variant_genes(tsv: str, exclude: set[str], levels: tuple[str, ...] = CLINPGX_LEVELS) -> dict[str, set[str]]:
    """ClinPGx ``clinicalVariants.tsv`` → gene → rsIDs with a clinical annotation at the given evidence levels, for
    single-gene rows of genes that no allele-definition source covers (level 3 as well for the PGX-008 parity genes)
    [DB: ClinPGx clinical annotations, CC BY-SA 4.0].
    Positions and alleles are resolved later from the manifest's probes (genotype-only variant-based genes)."""
    out: dict[str, set[str]] = {}
    lines = tsv.splitlines()
    head = lines[0].split("\t")
    iv, ig, il = head.index("variant"), head.index("gene"), head.index("level of evidence")
    for ln in lines[1:]:
        f = ln.split("\t")
        if len(f) <= max(iv, ig, il):
            continue
        gene, var, lev = f[ig].strip(), f[iv].strip(), f[il].strip()
        ok_level = lev in levels or (lev == "3" and gene in PARITY_VARIANT_GENES)
        if "," in gene or not gene or gene in exclude or not ok_level or not re.fullmatch(r"rs\d+", var):
            continue
        out.setdefault(gene, set()).add(var)
    return out