"""``simurg export bedgraph | finalreport | plink | locus-summary`` (spec CLI-023, OUT-043, OUT-050 … OUT-052,
OUT-071, OUT-072). None of them re-genotypes: they read GTCs (or SimurgArray VCFs) plus the manifest.
"""

from __future__ import annotations

import gzip
import io
import math
import re
import struct
from collections.abc import Iterator, Sequence
from dataclasses import dataclass, field
from pathlib import Path
from typing import IO, Any, cast

import numpy as np
import numpy.typing as npt

from simurg import PRODUCT_NAME, __version__
from simurg.core.atomic import atomic_open, check_outputs_absent
from simurg.core.errors import ExitCode, SimurgError
from simurg.core.log import Reporter
from simurg.core.provenance import RunManifest
from simurg.formats.bpm import Manifest
from simurg.formats.samplesheet import output_stem

_COMP = bytes.maketrans(b"ACGTDI", b"TGCADI")
LARGE_OUTPUT = 100_000_000


# ------------------------------------------------------------------------------------------------ inputs
@dataclass
class GtcSample:
    sample_id: str
    sample_name: str
    path: str
    reported_sex: str = "U"


def gtc_samples(gtcs: Path | None, sheet_path: Path | None, rep: Reporter) -> list[GtcSample]:
    """Sheet order when a sheet is given, else sorted sample IDs (OUT-076)."""
    from simurg.formats.discovery import discover_gtcs
    from simurg.formats.samplesheet import parse_samplesheet

    sheet = parse_samplesheet(sheet_path) if sheet_path else None
    root = gtcs or (sheet.root_folder if sheet else None)
    if root is None:
        raise SimurgError("E_NO_INPUT", "give --gtcs or a sample sheet with a RootFolder")
    found = discover_gtcs(root)
    if sheet is None:
        return [GtcSample(k, k, str(v)) for k, v in sorted(found.items())]
    out = []
    for r in sheet.rows:
        if r.sample_id in found:
            out.append(GtcSample(r.sample_id, r.sample_name or r.sample_id, str(found[r.sample_id]), r.reported_sex))
        else:
            rep.warn("W_QC_UNMATCHED_SAMPLES", f"{r.sample_id}: no GTC found", sample_id=r.sample_id)
    if not out:
        raise SimurgError("E_NO_INPUT", "none of the sheet's samples has a GTC")
    return out


def chrom_rank(chrom: str) -> tuple[int, str]:
    c = chrom[3:] if chrom.lower().startswith("chr") else chrom
    if c.isdigit():
        return (int(c) if c != "0" else 1000, "")
    return ({"X": 23, "XY": 24, "Y": 25, "MT": 26, "M": 26}.get(c, 100), c)


def genome_order(m: Manifest) -> npt.NDArray[np.int64]:
    """Loci by (chromosome, position, name); unmapped (Chr 0) last."""
    keys = [
        ((1000, "") if c in ("0", "") or p <= 0 else chrom_rank(c), int(p), n)
        for c, p, n in zip(m.chroms, m.map_infos.tolist(), m.names, strict=True)
    ]
    return np.array(sorted(range(len(keys)), key=keys.__getitem__), dtype=np.int64)


# ------------------------------------------------------------------------------------------------ bedGraph
def _chrom_label(chrom: str, style: str) -> str:
    c = "X" if chrom == "XY" else chrom
    if style == "ucsc":
        c = "M" if c == "MT" else c
        return c if c.startswith("chr") else f"chr{c}"
    return c


def moving_median(v: npt.NDArray[np.float64], k: int) -> npt.NDArray[np.float64]:
    if k <= 0 or len(v) == 0:
        return v.copy()
    pad = np.pad(v, k, mode="edge")
    win = np.lib.stride_tricks.sliding_window_view(pad, 2 * k + 1)
    return np.nanmedian(win, axis=1)


@dataclass
class BedgraphLayout:
    """Sample-independent part of OUT-050: loci in output order with their labels (computed once per manifest)."""

    order: npt.NDArray[np.int64]  # mapped loci, genome order
    chrom: list[str]  # output chromosome label per ordered locus
    pos: npt.NDArray[np.int64]


def bedgraph_layout(m: Manifest, style: str) -> BedgraphLayout:
    order = genome_order(m)
    mapped = np.array([c not in ("0", "") for c in m.chroms]) & (m.map_infos > 0)
    order = order[mapped[order]]
    return BedgraphLayout(order, [_chrom_label(m.chroms[i], style) for i in order.tolist()], m.map_infos[order])


def bedgraph_rows(m: Manifest, g: Any, layout: BedgraphLayout, which: str) -> list[tuple[str, int, float]]:
    """OUT-050 rows (chrom, pos, value) for one track, in output order, same-position probes averaged."""
    vals = g.b_allele_freqs if which == "baf" else g.logr_ratios
    if len(vals) != m.num_loci:
        raise SimurgError("E_GTC_UNSUITABLE", "GTC has no BAF/LRR (GTC v3?)")
    o = layout.order
    v = vals[o].astype(np.float64)
    keep = ~((g.raw_x[o] == 0) & (g.raw_y[o] == 0)) & np.isfinite(v)
    if which == "baf":
        keep &= ~m.intensity_only[o]
    idx = np.flatnonzero(keep)
    if not len(idx):
        return []
    chrom = [layout.chrom[i] for i in idx.tolist()]
    pos = layout.pos[idx]
    new = np.ones(len(idx), np.bool_)
    new[1:] = (pos[1:] != pos[:-1]) | np.array([chrom[k] != chrom[k - 1] for k in range(1, len(chrom))])
    starts = np.flatnonzero(new)
    sums = np.add.reduceat(v[idx], starts)
    counts = np.diff(np.append(starts, len(idx)))
    means = sums / counts
    return [(chrom[k], int(pos[k]), float(x)) for k, x in zip(starts.tolist(), means.tolist(), strict=True)]


def write_bedgraph(path: Path, name: str, which: str, rows: Sequence[tuple[str, int, float]], gz: bool) -> None:
    limits = "0:1" if which == "baf" else "-2:2"
    head = (
        f'track type=bedGraph name="{name} {which.upper()}" description="{PRODUCT_NAME} {__version__}" '
        f"visibility=full autoScale=off viewLimits={limits}\n"
    )
    body = "".join(f"{c}\t{p - 1}\t{p}\t{v:.6g}\n" for c, p, v in rows)
    data = (head + body).encode()
    with atomic_open(path, "wb") as fh:
        cast(IO[bytes], fh).write(gzip.compress(data, mtime=0) if gz else data)


@dataclass
class BedgraphOptions:
    out: Path
    bpm: Path
    gtcs: Path | None = None
    sheet: Path | None = None
    gzip: bool = True
    chrom_style: str = "manifest"
    smoothed_lrr: int = 0
    force: bool = False
    argv: list[str] = field(default_factory=list)


def run_export_bedgraph(opts: BedgraphOptions, rep: Reporter) -> int:
    from simurg.formats.bpm import read_bpm
    from simurg.formats.gtc import read_gtc

    if opts.chrom_style not in ("manifest", "ucsc"):
        raise SimurgError("E_CONFIG_KEY", "--chrom-style must be manifest or ucsc")
    rm = RunManifest(command="export bedgraph", argv=opts.argv)
    samples = gtc_samples(opts.gtcs, opts.sheet, rep)
    ext = ".bedgraph.gz" if opts.gzip else ".bedgraph"
    outs = []
    for s in samples:
        stem = output_stem(s.sample_name)
        outs += [opts.out / f"{stem}.baf{ext}", opts.out / f"{stem}.lrr{ext}"]
        if opts.smoothed_lrr:
            outs.append(opts.out / f"{stem}.lrr_smoothed{ext}")
    check_outputs_absent(outs, force=opts.force)
    m = read_bpm(opts.bpm)
    layout = bedgraph_layout(m, opts.chrom_style)
    opts.out.mkdir(parents=True, exist_ok=True)
    failed = 0
    for s in samples:
        try:
            g = read_gtc(s.path)
            stem = output_stem(s.sample_name)
            for which in ("baf", "lrr"):
                rows = bedgraph_rows(m, g, layout, which)
                write_bedgraph(opts.out / f"{stem}.{which}{ext}", s.sample_name, which, rows, opts.gzip)
                if which == "lrr" and opts.smoothed_lrr:
                    sm: list[tuple[str, int, float]] = []
                    for chrom in dict.fromkeys(c for c, _, _ in rows):
                        sub = [r for r in rows if r[0] == chrom]
                        med = moving_median(np.array([r[2] for r in sub]), opts.smoothed_lrr)
                        sm += [(c, p, float(v)) for (c, p, _), v in zip(sub, med, strict=True)]
                    write_bedgraph(opts.out / f"{stem}.lrr_smoothed{ext}", s.sample_name, "lrr", sm, opts.gzip)
        except SimurgError as e:
            failed += 1
            rep.error(e.code, f"{s.sample_id}: {e.message}", sample_id=s.sample_id)
    rm.status = "completed"
    rm.summary = {"samples": len(samples), "failed": failed}
    rm.write(opts.out / "run_manifest.export_bedgraph.json")
    return ExitCode.OK if not failed else (ExitCode.PARTIAL if failed < len(samples) else ExitCode.ALL_FAILED)


# ------------------------------------------------------------------------------------------------ final report
FR_COLUMNS = ["SNP Name", "SNP", "Sample ID", "Allele1 - Top", "Allele2 - Top", "Allele1 - Forward",
              "Allele2 - Forward", "Allele1 - Plus", "Allele2 - Plus", "GC Score", "GT Score", "Log R Ratio",
              "B Allele Freq", "Chr", "Position"]  # fmt: skip


def allele_letters(m: Manifest) -> dict[str, list[tuple[str, str]]]:
    """Per locus (A, B) letters on the TOP, Forward (source) and Plus (reference) strands."""
    top, fwd, plus = [], [], []
    for snp, ilmn, src, ref in zip(m.snps, m.ilmn_strands, m.source_strands, m.ref_strands, strict=True):
        mm = re.match(r"^\[([^/\]]+)/([^\]]+)\]$", snp.strip())
        a, b = (mm.group(1), mm.group(2)) if mm else ("-", "-")

        def comp(x: str) -> str:
            return x.encode().translate(_COMP).decode() if len(x) == 1 else x

        d = (ilmn or "")[:1].upper()
        top.append((a, b) if d in ("T", "P", "") or len(a) > 1 else (comp(a), comp(b)))
        same_src = (src or "")[:1].upper() == d or not src
        fwd.append((a, b) if same_src else (comp(a), comp(b)))
        plus.append((a, b) if (ref or "+").strip() == "+" else (comp(a), comp(b)))
    return {"top": top, "fwd": fwd, "plus": plus}


def _pair(call: int, ab: tuple[str, str]) -> tuple[str, str]:
    a, b = ab
    return {1: (a, a), 2: (a, b), 3: (b, b)}.get(call, ("-", "-"))


@dataclass
class FinalReportOptions:
    out: Path
    bpm: Path
    gtcs: Path | None = None
    sheet: Path | None = None
    csv: Path | None = None
    egt: Path | None = None
    matrix: bool = False
    force: bool = False
    argv: list[str] = field(default_factory=list)


def run_export_finalreport(opts: FinalReportOptions, rep: Reporter) -> int:
    from simurg.formats.bpm import read_bpm
    from simurg.formats.gtc import read_gtc

    rm = RunManifest(command="export finalreport", argv=opts.argv)
    samples = gtc_samples(opts.gtcs, opts.sheet, rep)
    target = opts.out / ("final_report_matrix.tsv.gz" if opts.matrix else "final_report.tsv.gz")
    check_outputs_absent([target], force=opts.force)
    m = read_bpm(opts.bpm)
    if opts.csv is not None:
        from simurg.formats.csv_manifest import read_csv_manifest

        cm = read_csv_manifest(opts.csv, load_sequences=False)
        m.source_strands, m.ref_strands = cm.source_strands, cm.ref_strands
    cells = m.num_loci * len(samples)
    if cells > LARGE_OUTPUT:
        rep.warn("W_LARGE_OUTPUT", f"{cells:,} rows (~{cells * 90 / 1e9:.1f} GB uncompressed); consider --matrix")
    gt_score = np.full(m.num_loci, np.nan)
    if opts.egt is not None:
        from simurg.formats.egt import read_egt

        egt = read_egt(opts.egt)
        idx = egt.name_index()
        for i, nm in enumerate(m.names):
            j = idx.get(nm)
            if j is not None:
                gt_score[i] = float(egt.gentrain_score[j])
    letters = allele_letters(m)
    order = genome_order(m).tolist()
    opts.out.mkdir(parents=True, exist_ok=True)

    def f4(v: float) -> str:
        return "NaN" if math.isnan(v) else f"{v:.4f}"

    with atomic_open(target, "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.finalreport/1\n")
        if opts.matrix:
            fh.write("\t".join(["SNP Name", "Chr", "Position", *[s.sample_name for s in samples]]) + "\n")
            calls = []
            for s in samples:
                calls.append(read_gtc(s.path).genotypes)
            for i in order:
                row = [m.names[i], m.chroms[i], str(int(m.map_infos[i]))]
                row += ["".join(_pair(int(c[i]), letters["top"][i])) for c in calls]
                fh.write("\t".join(row) + "\n")
        else:
            fh.write("\t".join(FR_COLUMNS) + "\n")
            for s in samples:
                g = read_gtc(s.path)
                has = len(g.b_allele_freqs) == m.num_loci
                bc = g.base_calls
                for i in order:
                    c = int(g.genotypes[i])
                    f, p = (_pair(c, letters[k][i]) for k in ("fwd", "plus"))
                    t = (chr(bc[i, 0]), chr(bc[i, 1])) if c else ("-", "-")  # TOP-strand base calls (OUT-071)
                    baf = float(g.b_allele_freqs[i]) if has and not m.intensity_only[i] else math.nan
                    lrr = float(g.logr_ratios[i]) if has else math.nan
                    fh.write(
                        f"{m.names[i]}\t{m.snps[i]}\t{s.sample_id}\t{t[0]}\t{t[1]}\t{f[0]}\t{f[1]}\t{p[0]}\t{p[1]}\t"
                        f"{f4(float(g.genotype_scores[i]))}\t{f4(gt_score[i])}\t{f4(lrr)}\t{f4(baf)}\t"
                        f"{m.chroms[i]}\t{int(m.map_infos[i])}\n"
                    )
        fh.flush()
        fh.detach()
        gz.close()
    rm.status = "completed"
    rm.summary = {"samples": len(samples), "loci": m.num_loci}
    rm.write(opts.out / "run_manifest.export_finalreport.json")
    return ExitCode.OK


# ------------------------------------------------------------------------------------------------ PLINK
PLINK_MAGIC = bytes([0x6C, 0x1B, 0x01])


def parse_vcf_records(path: Path) -> Iterator[list[str]]:
    opener = gzip.open if str(path).endswith(".gz") else open
    with opener(path, "rt") as fh:
        for line in fh:
            if not line.startswith("#"):
                yield line.rstrip("\n").split("\t")


def bed_code(gt: str, alt_index: int) -> int:
    """PLINK 1 two-bit code with A1 = the split ALT, A2 = REF: 0 hom A1, 1 missing, 2 het, 3 hom A2."""
    if gt in (".", "./."):
        return 1
    al = gt.replace("|", "/").split("/")
    if "." in al:
        return 1
    n_alt = sum(int(x) == alt_index for x in al)
    if len(al) == 1:
        n_alt *= 2
    return {0: 3, 1: 2, 2: 0}[n_alt]


@dataclass
class PlinkOptions:
    out: Path
    vcfs: Path
    sheet: Path | None = None
    gtcs: Path | None = None
    multiallelic: str = "split"
    format: str = "bed"
    prefix: str = "cohort"
    force: bool = False
    argv: list[str] = field(default_factory=list)


PGEN_FROM_BED = bytes([2, 3, 1, 0])  # PLINK 1 code (A1 = ALT) → PLINK 2 hardcall (0 ref/ref, 1 het, 2 alt/alt, 3 NA)


def _pack(codes: Sequence[int], nbytes: int) -> bytes:
    buf = bytearray(nbytes)
    for j, c in enumerate(codes):
        buf[j >> 2] |= c << (2 * (j & 3))
    return bytes(buf)


def run_export_plink(opts: PlinkOptions, rep: Reporter) -> int:
    """OUT-043 / plan §8.4: SimurgArray per-sample VCFs (identical record order) → PLINK 1 .bed/.bim/.fam and/or
    PLINK 2 .pgen/.pvar/.psam (fixed-width hardcall storage, mode 0x02 [DOC: pgen format]). Multi-allelic sites are
    split (``bcftools norm -m-``) or dropped."""
    if opts.multiallelic not in ("split", "drop"):
        raise SimurgError("E_CONFIG_KEY", "--multiallelic must be split or drop")
    if opts.format not in ("bed", "pgen", "both"):
        raise SimurgError("E_CONFIG_KEY", "--format must be bed, pgen or both")
    rm = RunManifest(command="export plink", argv=opts.argv)
    vcfs = sorted(p for p in opts.vcfs.glob("*.snv.vcf*") if not p.name.endswith((".tbi", ".csi")))
    if not vcfs:
        raise SimurgError("E_NO_INPUT", f"no *.snv.vcf(.gz) files in {opts.vcfs}")
    base = opts.out / opts.prefix
    want_bed, want_pgen = opts.format in ("bed", "both"), opts.format in ("pgen", "both")
    targets = ([".bed", ".bim", ".fam"] if want_bed else []) + ([".pgen", ".pvar", ".psam"] if want_pgen else [])
    check_outputs_absent([base.with_suffix(x) for x in targets], force=opts.force)
    names, sex = [], {}
    for v in vcfs:
        opener = gzip.open if str(v).endswith(".gz") else open
        with opener(v, "rt") as fh:
            for line in fh:
                if line.startswith("#CHROM"):
                    names.append(line.rstrip("\n").split("\t")[9])
                    break
    if opts.gtcs is not None:
        from simurg.formats.discovery import discover_gtcs
        from simurg.formats.gtc import read_gtc

        for key, gpath in discover_gtcs(opts.gtcs).items():
            sex[key] = read_gtc(gpath, check_write_complete=False).gender_char
    if opts.sheet is not None:
        from simurg.formats.samplesheet import parse_samplesheet

        for r in parse_samplesheet(opts.sheet).rows:
            if r.reported_sex in ("M", "F"):
                sex[r.sample_name or r.sample_id] = sex[r.sample_id] = r.reported_sex
    streams = [parse_vcf_records(v) for v in vcfs]
    n = len(vcfs)
    nbytes = (n + 3) // 4
    opts.out.mkdir(parents=True, exist_ok=True)
    n_var = dropped = 0
    pvar_header = "".join(
        ln for ln in _vcf_header(vcfs[0]) if ln.startswith("##contig") or ln.startswith("##fileformat")
    )
    with (
        atomic_open(base.with_suffix(".bed"), "wb") if want_bed else _Null() as bed_raw,
        atomic_open(base.with_suffix(".bim"), "w", "utf-8") if want_bed else _Null() as bim_raw,
        atomic_open(base.with_suffix(".pgen"), "wb") if want_pgen else _Null() as pgen_raw,
        atomic_open(base.with_suffix(".pvar"), "w", "utf-8") if want_pgen else _Null() as pvar_raw,
    ):
        bed, bim = cast(IO[bytes], bed_raw), cast(IO[str], bim_raw)
        pgen, pvar = cast(IO[bytes], pgen_raw), cast(IO[str], pvar_raw)
        if want_bed:
            bed.write(PLINK_MAGIC)
        if want_pgen:
            # header byte 12: nonref-flag storage 1 = every REF is the reference base (from the FASTA)
            pgen.write(b"l\x1b\x02" + struct.pack("<II", 0, n) + b"\x40")  # variant count patched below
            pvar.write(pvar_header + "#CHROM\tPOS\tID\tREF\tALT\n")
        for recs in zip(*streams, strict=True):
            r0 = recs[0]
            if any(r[:5] != r0[:5] for r in recs[1:]):
                raise SimurgError("E_GTC_UNSUITABLE", f"VCF record order differs between samples at {r0[0]}:{r0[1]}")
            alts = [] if r0[4] == "." else r0[4].split(",")
            if len(alts) > 1 and opts.multiallelic == "drop":
                dropped += 1
                continue
            gts = [r[9].split(":", 1)[0] for r in recs]
            for k, alt in enumerate(alts, start=1):
                codes = [bed_code(g, k) for g in gts]
                vid = r0[2] if len(alts) == 1 else f"{r0[2]}:{alt}"
                if want_bed:
                    bed.write(_pack(codes, nbytes))
                    bim.write(f"{r0[0]}\t{vid}\t0\t{r0[1]}\t{alt}\t{r0[3]}\n")
                if want_pgen:
                    pgen.write(_pack([PGEN_FROM_BED[c] for c in codes], nbytes))
                    pvar.write(f"{r0[0]}\t{r0[1]}\t{vid}\t{r0[3]}\t{alt}\n")
                n_var += 1
        if want_pgen:
            pgen.seek(3)
            pgen.write(struct.pack("<I", n_var))
    sex_code = {"M": 1, "F": 2}
    if want_bed:
        fam = "".join(f"{s}\t{s}\t0\t0\t{sex_code.get(sex.get(s, 'U'), 0)}\t-9\n" for s in names)
        with atomic_open(base.with_suffix(".fam"), "w", "utf-8") as fam_raw:
            cast(IO[str], fam_raw).write(fam)
    if want_pgen:
        psam = "#IID\tSEX\n" + "".join(f"{s}\t{ {'M': '1', 'F': '2'}.get(sex.get(s, 'U'), 'NA') }\n" for s in names)
        with atomic_open(base.with_suffix(".psam"), "w", "utf-8") as psam_raw:
            cast(IO[str], psam_raw).write(psam)
    rm.status = "completed"
    rm.summary = {"samples": n, "variants": n_var, "multiallelic_dropped": dropped, "format": opts.format}
    rm.write(opts.out / "run_manifest.export_plink.json")
    rep.info(f"PLINK fileset {base} ({opts.format}): {n} samples, {n_var} variants")
    return ExitCode.OK


class _Null:
    """Context manager standing in for an output that is not requested."""

    def __enter__(self) -> None:
        return None

    def __exit__(self, *a: object) -> None:
        return None


def _vcf_header(path: Path) -> list[str]:
    opener = gzip.open if str(path).endswith(".gz") else open
    out = []
    with opener(path, "rt") as fh:
        for line in fh:
            if not line.startswith("##"):
                break
            out.append(line)
    return out