"""BPM (binary bead-pool manifest) reader — plan §6.2, spec IN-040.

Layout (little-endian, .NET strings)::

    char[3] "BPM"; uint8 1
    int32  version            (flag 0x1000 is XOR-ed off; 3..5 supported)
    string manifest_name
    string control_config     (version > 1)
    int32  num_loci
    int32[num_loci] indexes
    string[num_loci] names
    uint8[num_loci]  norm_ids (< 100)
    num_loci × LocusEntry
    [int32 m; m × string]     trailer header lines (optional)

LocusEntry field order follows gtc2vcf ``locusentry_read`` (MIT) and BeadArrayFiles ``LocusEntry`` (BSD-2).
Loci are returned in **names order**, which is the order of every per-locus array in a GTC file.

Normalization-bin rule (plan §6.2)::

    norm_id' = (norm_id + 100 * assay_type) mod 256   # mimics the vendor's byte wrap (Omni5)
    lookup   = rank of norm_id' among the sorted distinct values
"""

from __future__ import annotations

import struct
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Any

import numpy as np
import numpy.typing as npt

from simurg.core.errors import FormatError
from simurg.formats.binutil import BinaryReader, load_bytes

ASSAY_INF_II = 0
ASSAY_INF_I_AT = 1  # red channel
ASSAY_INF_I_GC = 2  # green channel

_I32 = struct.Struct("<i")
_F32X4 = struct.Struct("<4f")


@dataclass
class Manifest:
    """Columnar manifest; every list/array has one entry per locus in names order."""

    path: str
    format: str  # "bpm" or "csv"
    version: int
    manifest_name: str
    control_config: str
    names: list[str]
    ilmn_ids: list[str]
    snps: list[str]
    chroms: list[str]
    map_infos: npt.NDArray[np.int64]
    ilmn_strands: list[str]
    source_strands: list[str]
    ref_strands: list[str]
    ploidies: list[str]
    species: list[str]
    genome_builds: list[str]
    sources: list[str]
    source_versions: list[str]
    customer_strands: list[str]
    source_seqs: list[str]
    top_genomic_seqs: list[str]
    allele_a_probe_seqs: list[str]
    allele_b_probe_seqs: list[str]
    address_a: npt.NDArray[np.int32]
    address_b: npt.NDArray[np.int32]
    assay_types: npt.NDArray[np.uint8]
    exp_clusters: npt.NDArray[np.int8]
    intensity_only: npt.NDArray[np.bool_]
    norm_ids: npt.NDArray[np.uint8]  # raw per-locus IDs (< 100)
    entry_versions: npt.NDArray[np.int8]
    base_fractions: npt.NDArray[np.float32] | None = None  # (n, 4): A, C, T, G (entry v >= 7)
    header_lines: list[str] = field(default_factory=list)
    controls: list[dict[str, str]] = field(default_factory=list)
    heading: dict[str, str] = field(default_factory=dict)

    @property
    def num_loci(self) -> int:
        return len(self.names)

    @property
    def has_ref_strand(self) -> bool:
        return any(s for s in self.ref_strands)

    def norm_bin_ids(self) -> npt.NDArray[np.uint8]:
        """``(norm_id + 100 * assay_type) mod 256`` per locus."""
        return ((self.norm_ids.astype(np.int64) + 100 * self.assay_types.astype(np.int64)) % 256).astype(np.uint8)

    def normalization_lookups(self) -> npt.NDArray[np.int32]:
        """Index into the GTC's normalization-transform list for every locus."""
        ids = self.norm_bin_ids()
        uniq = np.unique(ids)  # sorted
        return np.searchsorted(uniq, ids).astype(np.int32)

    def name_index(self) -> dict[str, int]:
        return {n: i for i, n in enumerate(self.names)}

    def lite(self) -> Manifest:
        """Copy keeping only what intensity assembly needs (cheap to pickle to worker processes)."""
        n = self.num_loci
        empty = [""] * 0
        return replace(
            self,
            names=list(self.names),
            ilmn_ids=empty,
            snps=empty,
            chroms=empty,
            ilmn_strands=empty,
            source_strands=empty,
            ref_strands=empty,
            ploidies=empty,
            species=empty,
            genome_builds=empty,
            sources=empty,
            source_versions=empty,
            customer_strands=empty,
            source_seqs=empty,
            top_genomic_seqs=empty,
            allele_a_probe_seqs=empty,
            allele_b_probe_seqs=empty,
            map_infos=np.zeros(0, np.int64),
            base_fractions=None,
            header_lines=[],
            heading={},
            exp_clusters=self.exp_clusters[:n],
            intensity_only=self.intensity_only[:n],
        )

    def summary(self) -> dict[str, Any]:
        at = self.assay_types
        return {
            "path": self.path,
            "format": self.format,
            "version": self.version,
            "manifest_name": self.manifest_name,
            "num_loci": self.num_loci,
            "entry_versions": sorted({int(v) for v in np.unique(self.entry_versions)}),
            "infinium_ii": int((at == ASSAY_INF_II).sum()),
            "infinium_i_at": int((at == ASSAY_INF_I_AT).sum()),
            "infinium_i_gc": int((at == ASSAY_INF_I_GC).sum()),
            "intensity_only": int(self.intensity_only.sum()),
            "normalization_bins": len(np.unique(self.norm_bin_ids())),
            "has_ref_strand": self.has_ref_strand,
            "n_controls": len(self.controls),
            "genome_builds": sorted(set(self.genome_builds) - {""}),
            "chromosomes": sorted(set(self.chroms), key=_chrom_key),
        }


def _chrom_key(c: str) -> tuple[int, str]:
    return (int(c), "") if c.isdigit() else (1000, c)


def parse_control_config(text: str) -> list[dict[str, str]]:
    """Control-config lines ``addr[:addr..],Category,Color,Name`` → dicts."""
    out: list[dict[str, str]] = []
    for line in text.splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) < 4 or not parts[0]:
            continue
        out.append({"address": parts[0], "category": parts[1], "color": parts[2], "name": ",".join(parts[3:])})
    return out


def read_bpm(path: str | Path) -> Manifest:
    """Parse a BPM file. Raises FormatError(E_BPM_FORMAT)."""
    spath = str(path)
    buf = load_bytes(path)
    r = BinaryReader(buf, "E_BPM_FORMAT", spath)
    if r.size < 8 or r.raw(3, "magic") != b"BPM":
        raise r.fail("bad magic (not a BPM file)")
    if (fmt := r.u8("format byte")) != 1:
        raise r.fail(f"unknown BPM format byte {fmt}")
    version = r.i32("version")
    if version & 0x1000:
        version ^= 0x1000
    if not 3 <= version <= 5:
        raise r.fail(f"unsupported BPM version {version}")
    manifest_name = r.string("manifest name")
    control_config = r.string("control config") if version > 1 else ""
    n = r.i32("number of loci")
    if n <= 0 or n > 50_000_000:
        raise r.fail(f"implausible number of loci {n}")
    r.array("<i4", n, "indexes")
    names = [r.string("locus name") for _ in range(n)]
    norm_ids = np.array(r.array("u1", n, "normalization ids"), dtype=np.uint8)
    if (norm_ids >= 100).any():
        raise r.fail("invalid normalization ID (>= 100)")
    name_to_idx: dict[str, int] = {}
    for i, nm in enumerate(names):
        if nm in name_to_idx:
            raise r.fail(f"locus name {nm!r} present more than once")
        name_to_idx[nm] = i

    cols = _parse_locus_entries(r, n, name_to_idx)

    header_lines: list[str] = []
    if r.remaining() >= 4:
        m = r.i32("trailer count")
        if 0 <= m <= 1_000_000:
            header_lines = [r.string("trailer line") for _ in range(m)]

    return Manifest(
        path=spath,
        format="bpm",
        version=version,
        manifest_name=manifest_name,
        control_config=control_config,
        names=names,
        norm_ids=norm_ids,
        header_lines=header_lines,
        controls=parse_control_config(control_config),
        **cols,
    )


def _parse_locus_entries(r: BinaryReader, n: int, name_to_idx: dict[str, int]) -> dict[str, Any]:
    """Hot loop: parse ``n`` LocusEntry records (optimised pure Python; Rust in phase P9)."""
    view = memoryview(r._buf)
    size = r.size
    pos = r.pos
    intern: dict[bytes, str] = {}
    unpack_i32 = _I32.unpack_from
    unpack_4f = _F32X4.unpack_from

    def fail(msg: str) -> FormatError:
        r.pos = pos
        return r.fail(msg)

    def s_at(p: int) -> tuple[str, int]:
        # .NET 7-bit length-prefixed UTF-8 string, interned
        if p >= size:
            raise fail("unexpected end of file in locus entry")
        b = view[p]
        p += 1
        ln = b & 0x7F
        shift = 7
        while b & 0x80:
            if p >= size or shift > 28:
                raise fail("bad string length in locus entry")
            b = view[p]
            p += 1
            ln |= (b & 0x7F) << shift
            shift += 7
        end = p + ln
        if end > size:
            raise fail("string runs past end of file in locus entry")
        raw = bytes(view[p:end])
        s = intern.get(raw)
        if s is None:
            s = raw.decode("utf-8", errors="replace")
            if ln <= 64:
                intern[raw] = s
        return s, end

    str_cols = [
        "ilmn_ids",
        "snps",
        "chroms",
        "ilmn_strands",
        "source_strands",
        "ref_strands",
        "ploidies",
        "species",
        "genome_builds",
        "sources",
        "source_versions",
        "customer_strands",
        "source_seqs",
        "top_genomic_seqs",
        "allele_a_probe_seqs",
        "allele_b_probe_seqs",
    ]
    out: dict[str, list[str]] = {c: [""] * n for c in str_cols}
    map_infos = np.zeros(n, dtype=np.int64)
    address_a = np.zeros(n, dtype=np.int32)
    address_b = np.zeros(n, dtype=np.int32)
    assay = np.zeros(n, dtype=np.uint8)
    exp_clusters = np.zeros(n, dtype=np.int8)
    intensity_only = np.zeros(n, dtype=np.bool_)
    versions = np.zeros(n, dtype=np.int8)
    fracs = np.zeros((n, 4), dtype=np.float32)
    seen = np.zeros(n, dtype=np.bool_)
    any_v7 = False

    o_ilmn, o_snp, o_chr = out["ilmn_ids"], out["snps"], out["chroms"]
    o_istr, o_sstr, o_rstr = out["ilmn_strands"], out["source_strands"], out["ref_strands"]
    o_ploidy, o_species, o_build = out["ploidies"], out["species"], out["genome_builds"]
    o_source, o_sver, o_cstr = out["sources"], out["source_versions"], out["customer_strands"]
    o_sseq, o_top, o_pa, o_pb = (
        out["source_seqs"],
        out["top_genomic_seqs"],
        out["allele_a_probe_seqs"],
        out["allele_b_probe_seqs"],
    )

    for _ in range(n):
        if pos + 4 > size:
            raise fail("unexpected end of file at locus entry")
        ver = unpack_i32(view, pos)[0]
        pos += 4
        if ver not in (4, 6, 7, 8):
            raise fail(f"unsupported locus entry version {ver}")
        ilmn_id, pos = s_at(pos)
        name, pos = s_at(pos)
        _, pos = s_at(pos)
        _, pos = s_at(pos)
        _, pos = s_at(pos)
        if pos + 4 > size:
            raise fail("unexpected end of file at locus index")
        pos += 4  # index (1-based); BeadArrayFiles maps entries by name, as we do
        _, pos = s_at(pos)
        ilmn_strand, pos = s_at(pos)
        snp, pos = s_at(pos)
        chrom, pos = s_at(pos)
        ploidy, pos = s_at(pos)
        species, pos = s_at(pos)
        map_info, pos = s_at(pos)
        top_seq, pos = s_at(pos)
        cust_strand, pos = s_at(pos)
        if pos + 8 > size:
            raise fail("unexpected end of file at addresses")
        a_addr = unpack_i32(view, pos)[0]
        b_addr = unpack_i32(view, pos + 4)[0]
        pos += 8
        probe_a, pos = s_at(pos)
        probe_b, pos = s_at(pos)
        build, pos = s_at(pos)
        source, pos = s_at(pos)
        source_ver, pos = s_at(pos)
        source_strand, pos = s_at(pos)
        source_seq, pos = s_at(pos)
        at = -1
        exp_cl = 0
        int_only = 0
        if ver >= 6:
            if pos + 4 > size:
                raise fail("unexpected end of file at assay type")
            exp_cl = int.from_bytes(view[pos + 1 : pos + 2], "little", signed=True)
            int_only = view[pos + 2]
            at = view[pos + 3]
            pos += 4
            if at > 2:
                raise fail(f"invalid assay type {at}")
            if (b_addr == 0) != (at == 0):
                raise fail("assay type is inconsistent with address B")
        f4 = None
        if ver >= 7:
            if pos + 16 > size:
                raise fail("unexpected end of file at base fractions")
            f4 = unpack_4f(view, pos)
            pos += 16
        ref_strand = ""
        if ver >= 8:
            ref_strand, pos = s_at(pos)

        idx = name_to_idx.get(name)
        if idx is None:
            raise fail(f"locus entry {name!r} not in the names table")
        if seen[idx]:
            raise fail(f"locus entry {name!r} appears twice")
        seen[idx] = True
        if at < 0:
            at = infer_assay_type(probe_a, probe_b, source_seq)
        o_ilmn[idx] = ilmn_id
        o_snp[idx] = snp
        o_chr[idx] = chrom
        o_istr[idx] = ilmn_strand
        o_sstr[idx] = source_strand
        o_rstr[idx] = ref_strand
        o_ploidy[idx] = ploidy
        o_species[idx] = species
        o_build[idx] = build
        o_source[idx] = source
        o_sver[idx] = source_ver
        o_cstr[idx] = cust_strand
        o_sseq[idx] = source_seq
        o_top[idx] = top_seq
        o_pa[idx] = probe_a
        o_pb[idx] = probe_b
        try:
            map_infos[idx] = int(map_info) if map_info else 0
        except ValueError:
            raise fail(f"non-integer MapInfo {map_info!r} for {name!r}") from None
        address_a[idx] = a_addr
        address_b[idx] = b_addr
        assay[idx] = at
        exp_clusters[idx] = exp_cl
        intensity_only[idx] = bool(int_only)
        versions[idx] = ver
        if f4 is not None:
            fracs[idx] = f4
            any_v7 = True

    if not seen.all():
        raise fail(f"{int((~seen).sum())} loci in the names table have no locus entry")
    r.pos = pos
    return {
        **out,
        "map_infos": map_infos,
        "address_a": address_a,
        "address_b": address_b,
        "assay_types": assay,
        "exp_clusters": exp_clusters,
        "intensity_only": intensity_only,
        "entry_versions": versions,
        "base_fractions": fracs if any_v7 else None,
    }


_IUPAC = {
    "A": 1,
    "C": 2,
    "G": 4,
    "T": 8,
    "R": 5,
    "Y": 10,
    "S": 6,
    "W": 9,
    "K": 12,
    "M": 3,
    "B": 14,
    "D": 13,
    "H": 11,
    "V": 7,
    "N": 15,
}


def infer_assay_type(probe_a: str, probe_b: str, source_seq: str) -> int:
    """Assay type for locus entries older than v6 (gtc2vcf ``get_assay_type``, MIT).

    Returns 0 for Infinium II, 1 for Infinium I red (A/T), 2 for Infinium I green (C/G), 255 if unknown.
    """
    if not probe_a or not source_seq:
        return 255
    if not probe_b:
        return 0
    left = source_seq.find("[")
    right = source_seq.find("]")
    if left <= 0 or right < 0 or right + 1 >= len(source_seq):
        return 255
    tl = source_seq[left - 1].upper()
    tr = source_seq[right + 1].upper()
    if tl in "AT" and tr in "AT":
        return 1
    if tl in "CG" and tr in "CG":
        return 2
    i = 2
    while i <= min(len(probe_a), len(probe_b)) and not (
        _IUPAC.get(probe_a[-i].upper(), 0) & _IUPAC.get(probe_b[-i].upper(), 0)
    ):
        i += 1
    if i > min(len(probe_a), len(probe_b)):
        return 255
    t = probe_a[-i].upper()
    if t in "CGS":
        return 1
    if t in "ATW":
        return 2
    rules = {
        ("Y", "G"): 1,
        ("Y", "T"): 1,
        ("Y", "A"): 2,
        ("K", "C"): 1,
        ("K", "A"): 2,
        ("M", "G"): 1,
        ("M", "T"): 2,
        ("R", "C"): 1,
        ("R", "T"): 2,
    }
    return rules.get((t, tr), 255)