"""Reference genome access (spec IN-060 … IN-064).

* Plain or BGZF-compressed FASTA with a samtools-style ``.fai`` (and ``.gzi`` for BGZF). Missing indexes are built
  into the cache directory, never next to the input.
* ``Genome.fetch`` returns upper-case bases for a 0-based half-open interval. With the genome cache (default) each
  needed contig is stored once as a flat uint8 file under ``$SIMURG_CACHE_DIR/genomes/<fasta-sha256>/`` and
  memory-mapped; without it, bases are read through the index. Both modes return identical sequences (EDGE-010).

[DOC: samtools faidx / htslib BGZF and .gzi formats — layouts only]
"""

from __future__ import annotations

import hashlib
import json
import os
import struct
import zlib
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import numpy.typing as npt

from simurg.core.errors import SimurgError

GRCH37_CHR1 = 249_250_621
GRCH38_CHR1 = 248_956_422
_BGZF_MAGIC = b"\x1f\x8b\x08\x04"


@dataclass(frozen=True)
class FaiEntry:
    name: str
    length: int
    offset: int  # uncompressed byte offset of the first base
    line_bases: int
    line_width: int


def cache_root() -> Path:
    return Path(os.environ.get("SIMURG_CACHE_DIR", Path.home() / ".cache" / "simurg"))


def is_bgzf(path: Path) -> bool:
    with path.open("rb") as fh:
        head = fh.read(16)
    return head[:4] == _BGZF_MAGIC and head[12:14] == b"BC"


def file_sha256(path: Path) -> str:
    """SHA-256 of a (large) file, memoised by path, size and mtime in the cache directory."""
    st = path.stat()
    memo = cache_root() / "sha256.json"
    key = f"{path.resolve()}|{st.st_size}|{st.st_mtime_ns}"
    try:
        known = json.loads(memo.read_text())
    except (OSError, ValueError):
        known = {}
    if key in known:
        return str(known[key])
    h = hashlib.sha256()
    with path.open("rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 22), b""):
            h.update(chunk)
    known[key] = h.hexdigest()
    memo.parent.mkdir(parents=True, exist_ok=True)
    tmp = memo.with_suffix(".part")
    tmp.write_text(json.dumps(known, indent=0))
    tmp.replace(memo)
    return str(known[key])


# ------------------------------------------------------------------------------------------------ BGZF
def read_gzi(path: Path) -> list[tuple[int, int]]:
    """``.gzi``: (compressed offset, uncompressed offset) of every BGZF block after the first."""
    data = path.read_bytes()
    (n,) = struct.unpack_from("<Q", data, 0)
    return [struct.unpack_from("<QQ", data, 8 + 16 * i) for i in range(n)]


def build_gzi(fasta: Path) -> list[tuple[int, int]]:
    out = []
    coff = uoff = 0
    with fasta.open("rb") as fh:
        while True:
            head = fh.read(18)
            if not head:
                break
            if head[:4] != _BGZF_MAGIC:
                raise SimurgError("E_FASTA_FORMAT", f"{fasta}: not a BGZF file (bad block at {coff})")
            bsize = struct.unpack_from("<H", head, 16)[0] + 1
            body = fh.read(bsize - 18)
            isize = struct.unpack_from("<I", body, len(body) - 4)[0]
            if coff and isize:  # htslib omits empty blocks (the EOF marker)
                out.append((coff, uoff))
            coff += bsize
            uoff += isize
    return out


def write_gzi(path: Path, entries: list[tuple[int, int]]) -> None:
    buf = struct.pack("<Q", len(entries)) + b"".join(struct.pack("<QQ", c, u) for c, u in entries)
    path.write_bytes(buf)


class _BgzfReader:
    """Random access to uncompressed offsets of a BGZF file via its ``.gzi``."""

    def __init__(self, path: Path, gzi: list[tuple[int, int]]) -> None:
        self.fh = path.open("rb")
        self.blocks = [(0, 0), *gzi]
        self.ustarts = np.array([u for _, u in self.blocks], dtype=np.int64)
        self._cache: tuple[int, bytes] | None = None

    def _block(self, i: int) -> bytes:
        if self._cache and self._cache[0] == i:
            return self._cache[1]
        self.fh.seek(self.blocks[i][0])
        head = self.fh.read(18)
        bsize = struct.unpack_from("<H", head, 16)[0] + 1
        body = self.fh.read(bsize - 18)
        data = zlib.decompress(body[:-8], -15)
        self._cache = (i, data)
        return data

    def read(self, uoff: int, n: int) -> bytes:
        out = bytearray()
        i = int(np.searchsorted(self.ustarts, uoff, side="right") - 1)
        pos = uoff - int(self.ustarts[i])
        while len(out) < n and i < len(self.blocks):
            data = self._block(i)
            out += data[pos : pos + n - len(out)]
            i, pos = i + 1, 0
        return bytes(out)


class _PlainReader:
    def __init__(self, path: Path) -> None:
        self.fh = path.open("rb")

    def read(self, uoff: int, n: int) -> bytes:
        self.fh.seek(uoff)
        return self.fh.read(n)


def _stream_uncompressed(path: Path, bgzf: bool) -> object:
    if not bgzf:
        return path.open("rb")
    import gzip

    return gzip.open(path, "rb")


# ------------------------------------------------------------------------------------------------ FAI
def read_fai(path: Path) -> list[FaiEntry]:
    out = []
    for i, line in enumerate(path.read_text().splitlines(), start=1):
        if not line.strip():
            continue
        f = line.split("\t")
        try:
            out.append(FaiEntry(f[0], int(f[1]), int(f[2]), int(f[3]), int(f[4])))
        except (IndexError, ValueError):
            raise SimurgError("E_FAI_MISMATCH", f"{path}: malformed line {i}", path=str(path)) from None
    return out


def build_fai(fasta: Path, bgzf: bool) -> list[FaiEntry]:
    """samtools-compatible index; offsets are uncompressed byte offsets."""
    entries: list[FaiEntry] = []
    fh = _stream_uncompressed(fasta, bgzf)
    off = 0
    name: str | None = None
    length = seq_off = lb = lw = 0
    with fh:  # type: ignore[attr-defined]
        for raw in fh:  # type: ignore[attr-defined]
            if raw.startswith(b">"):
                if name is not None:
                    entries.append(FaiEntry(name, length, seq_off, lb, lw))
                name = raw[1:].split()[0].decode()
                length, lb, lw = 0, 0, 0
                seq_off = off + len(raw)
            else:
                n = len(raw.rstrip(b"\r\n"))
                if lb == 0:
                    lb, lw = n, len(raw)
                length += n
            off += len(raw)
    if name is not None:
        entries.append(FaiEntry(name, length, seq_off, lb, lw))
    if not entries:
        raise SimurgError("E_FASTA_FORMAT", f"{fasta}: no sequences found", path=str(fasta))
    return entries


def write_fai(path: Path, entries: list[FaiEntry]) -> None:
    path.write_text("".join(f"{e.name}\t{e.length}\t{e.offset}\t{e.line_bases}\t{e.line_width}\n" for e in entries))


def detect_build(entries: list[FaiEntry]) -> str:
    """IN-063: chr1 length decides GRCh37 vs GRCh38; '' when unknown."""
    for e in entries:
        if e.name in ("1", "chr1"):
            return {GRCH37_CHR1: "GRCh37", GRCH38_CHR1: "GRCh38"}.get(e.length, "")
    return ""


# ------------------------------------------------------------------------------------------------ genome
class Genome:
    """Indexed reference genome; see module docstring."""

    def __init__(self, fasta: str | Path, *, use_cache: bool = True, cache_dir: Path | None = None) -> None:
        self.path = Path(fasta).resolve()
        if not self.path.exists():
            raise SimurgError("E_FILE_NOT_FOUND", f"FASTA not found: {fasta}", path=str(fasta))
        self.bgzf = is_bgzf(self.path)
        self.use_cache = use_cache
        self._sha: str | None = None
        self._cache_dir = cache_dir
        self.notes: list[str] = []
        fai = Path(str(self.path) + ".fai")
        if fai.exists():
            self.entries = read_fai(fai)
        else:
            built = self.cache_dir / "index.fai"
            if not built.exists():
                built.parent.mkdir(parents=True, exist_ok=True)
                write_fai(built, build_fai(self.path, self.bgzf))
                self.notes.append(f"built FASTA index {built}")
            self.entries = read_fai(built)
        self.index = {e.name: e for e in self.entries}
        if self.bgzf:
            gzi = Path(str(self.path) + ".gzi")
            if not gzi.exists():
                gzi = self.cache_dir / "index.gzi"
                if not gzi.exists():
                    gzi.parent.mkdir(parents=True, exist_ok=True)
                    write_gzi(gzi, build_gzi(self.path))
            self._reader: _BgzfReader | _PlainReader = _BgzfReader(self.path, read_gzi(gzi))
        else:
            self._reader = _PlainReader(self.path)
        self._mm: dict[str, npt.NDArray[np.uint8]] = {}
        self.verify()
        self.build = detect_build(self.entries)

    @property
    def sha256(self) -> str:
        if self._sha is None:
            self._sha = file_sha256(self.path)
        return self._sha

    @property
    def cache_dir(self) -> Path:
        if self._cache_dir is None:
            self._cache_dir = cache_root() / "genomes" / self.sha256
        return self._cache_dir

    def verify(self, n: int = 3) -> None:
        """IN-060: sample up to ``n`` contigs and check header position and line layout."""
        picks = self.entries[:: max(1, len(self.entries) // n)][:n]
        for e in picks:
            name_bytes = e.name.encode()
            before = self._reader.read(max(0, e.offset - len(name_bytes) - 64), min(e.offset, len(name_bytes) + 64))
            first = self._reader.read(e.offset, e.line_width)
            ok = before.endswith(b"\n") and b">" + name_bytes in before
            if e.length >= e.line_bases:
                ok = ok and len(first.rstrip(b"\r\n")) == e.line_bases and first.endswith(b"\n")
            if not ok:
                raise SimurgError(
                    "E_FAI_MISMATCH",
                    f"FASTA index does not match {self.path} at contig {e.name}",
                    hint="delete the stale .fai (it will be rebuilt in the cache) or regenerate it with samtools faidx",
                    path=str(self.path),
                )

    def contig_length(self, name: str) -> int:
        return self.index[name].length

    def _read_raw(self, e: FaiEntry, start: int, end: int) -> bytes:
        lines_before = start // e.line_bases
        uoff = e.offset + lines_before * e.line_width + start % e.line_bases
        last = end - 1
        uend = e.offset + (last // e.line_bases) * e.line_width + last % e.line_bases + 1
        raw = self._reader.read(uoff, uend - uoff)
        return raw.replace(b"\n", b"").replace(b"\r", b"")

    def _contig_array(self, name: str) -> npt.NDArray[np.uint8]:
        if name not in self._mm:
            e = self.index[name]
            f = self.cache_dir / f"{name}.u8"
            if not f.exists() or f.stat().st_size != e.length:
                f.parent.mkdir(parents=True, exist_ok=True)
                tmp = f.with_suffix(".part")
                with tmp.open("wb") as out:
                    step = 1 << 24
                    for s in range(0, e.length, step):
                        out.write(self._read_raw(e, s, min(e.length, s + step)).upper())
                tmp.replace(f)
            self._mm[name] = np.memmap(f, dtype=np.uint8, mode="r")
        return self._mm[name]

    def fetch(self, name: str, start: int, end: int) -> str:
        """Upper-case bases of ``name``[start:end] (0-based, half-open), clipped to the contig."""
        e = self.index[name]
        start, end = max(0, start), min(e.length, end)
        if end <= start:
            return ""
        if self.use_cache:
            return bytes(self._contig_array(name)[start:end]).decode("ascii")
        return self._read_raw(e, start, end).decode("ascii").upper()


# ------------------------------------------------------------------------------------------------ contig names
def load_alias_file(path: Path) -> dict[str, str]:
    out = {}
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.split("#", 1)[0].strip()
        if line:
            a, b = line.split("\t")[:2]
            out[a] = b
    return out


class ContigMap:
    """IN-061: manifest chromosome → FASTA contig (``XY`` → X). Policies: ``strict`` (verbatim), ``auto``
    (``1``↔``chr1``, ``MT``↔``chrM``, ``22_KI270879v1_alt``↔``chr22_KI270879v1_alt``) or an alias TSV path."""

    def __init__(self, genome_contigs: set[str] | list[str], policy: str = "strict") -> None:
        self.contigs = set(genome_contigs)
        self.policy = policy
        self.alias = load_alias_file(Path(policy)) if policy not in ("strict", "auto") else {}
        self._memo: dict[str, str | None] = {}

    def get(self, chrom: str) -> str | None:
        if chrom in self._memo:
            return self._memo[chrom]
        target = "X" if chrom == "XY" else chrom
        if self.alias:
            cands = [self.alias.get(chrom, ""), self.alias.get(target, ""), target]
        elif self.policy == "auto":
            bare = target[3:] if target.lower().startswith("chr") else target
            cands = [target, bare, f"chr{bare}"]
            if bare in ("MT", "M"):
                cands += ["chrM", "MT", "M"]
        else:
            cands = [target]
        hit = next((c for c in cands if c and c in self.contigs), None)
        self._memo[chrom] = hit
        return hit