"""Corpus pipeline: fetch -> normalize (+filter, dedupe) -> tokenizer sample -> tokenize (+decontaminate, split, shard).

Usage: python -m zenith.data.pipeline {fetch,normalize,sample,tokenize} [--sources a,b] [options]
"""

from __future__ import annotations

import argparse
import hashlib
import json
from collections import defaultdict
from pathlib import Path

from zenith.data.clean import dedupe, keep, normalize_text
from zenith.data.decontam import build_ngram_index, is_contaminated
from zenith.data.shards import ShardWriter, split_of, write_manifest
from zenith.data.sources.base import Record
from zenith.tokenizer.export import encode_fast, load_fast

RAW_DIR = Path("data/raw")
INTERIM_DIR = Path("data/interim")
PROCESSED_DIR = Path("data/processed")
EVAL_DIR = Path("data/eval")
TOKENIZER_HF = Path("zenith/tokenizer/zenith-16k.hf.json")
SHOWCASE = Path("zenith/evals/showcase.txt")
DEFAULT_MIX = {"astro": 0.30, "arxiv": 0.15, "textbook": 0.05, "general": 0.50}
NGRAM = 8  # eval questions are short, so the HADES 13-gram window would miss most overlaps


def _read_interim(path: Path):
    with open(path) as f:
        for line in f:
            yield Record.from_json(line)


def normalize_stage(raw_dir: Path, interim_dir: Path, sources: dict) -> None:
    interim_dir.mkdir(parents=True, exist_ok=True)

    def cleaned():
        for mod in sources.values():
            for r in mod.iter_records(raw_dir):
                r = Record(r.source, r.id, r.license, r.group, normalize_text(r.text))
                if keep(r):
                    yield r

    handles = {name: open(interim_dir / f"{name}.jsonl", "w") for name in sources}
    try:
        for r in dedupe(cleaned()):
            handles[r.source].write(r.to_json() + "\n")
    finally:
        for h in handles.values():
            h.close()


def sample_stage(interim_dir: Path, out_path: Path, total_bytes: int, mix: dict[str, float]) -> None:
    """Write a tokenizer-training sample with the group mix, taking evenly spaced docs from each group."""
    by_group: dict[str, list[Path]] = defaultdict(list)
    sizes: dict[str, int] = defaultdict(int)
    for path in sorted(interim_dir.glob("*.jsonl")):
        first = next(_read_interim(path), None)
        if first is None:
            continue
        by_group[first.group].append(path)
        sizes[first.group] += path.stat().st_size
    with open(out_path, "w") as out:
        for group, share in mix.items():
            budget = int(total_bytes * share)
            stride = max(1, sizes[group] // max(budget, 1))
            written = 0
            for path in by_group.get(group, []):
                for i, r in enumerate(_read_interim(path)):
                    if written >= budget:
                        break
                    if i % stride == 0:
                        chunk = r.text[: budget - written]
                        out.write(chunk + "\n\n")
                        written += len(chunk.encode()) + 2


def eval_texts_from(eval_dir: Path, showcase_path: Path | None) -> tuple[list[str], str]:
    """Every evaluation question (`question` or `text` field) plus showcase lines, and their combined hash."""
    texts, h = [], hashlib.sha256()
    paths = sorted(Path(eval_dir).glob("*.jsonl")) + ([Path(showcase_path)] if showcase_path and Path(showcase_path).exists() else [])
    for path in paths:
        data = path.read_bytes()
        h.update(data)
        for line in data.decode().splitlines():
            if not line.strip():
                continue
            if path.suffix == ".jsonl":
                row = json.loads(line)
                texts.append(row.get("question") or row["text"])
            else:
                texts.append(line.strip())
    return texts, h.hexdigest()


def tokenize_stage(interim_dir: Path, out_dir: Path, tokenizer_path: Path, eval_dir: Path, sources: dict,
                   showcase_path: Path | None = SHOWCASE, shard_tokens: int = 100_000_000, batch: int = 512,
                   allow_no_decontam: bool = False) -> None:
    fast = load_fast(tokenizer_path)
    bos, eos = fast.token_to_id("<|bos|>"), fast.token_to_id("<|eos|>")
    eval_texts, eval_sha = eval_texts_from(eval_dir, showcase_path)
    if not eval_texts and not allow_no_decontam:
        raise ValueError(
            f"no evaluation texts in {eval_dir}: decontamination would silently do nothing. "
            "Run `python -m zenith.sft.generate` first, or pass --no-decontam deliberately."
        )
    letter_ids = {i for t, i in fast.get_vocab().items() if any(c.isalpha() for c in fast.decode([i]))}
    index = build_ngram_index(encode_fast(fast, eval_texts), n=NGRAM, letter_ids=letter_ids)
    decontamination = {"eval_texts": len(eval_texts), "ngrams": len(index),
                       "eval_sha256": eval_sha if eval_texts else None, "ngram": NGRAM}

    writers: dict[tuple[str, str], ShardWriter] = {}
    source_stats: dict[str, dict] = {}
    for name in sources:
        path = interim_dir / f"{name}.jsonl"
        if not path.exists():
            continue
        st = {"docs": 0, "tokens": 0, "bytes": 0, "contaminated": 0, "group": None, "licenses": set()}
        pending: list[Record] = []

        def flush():
            for r, ids in zip(pending, encode_fast(fast, [r.text for r in pending])):
                if is_contaminated(ids, index, n=NGRAM):
                    st["contaminated"] += 1
                    continue
                split = split_of(r.id)
                key = (r.group, split)
                if key not in writers:
                    writers[key] = ShardWriter(out_dir / r.group, split, bos, eos, shard_tokens)
                n_bytes = len(r.text.encode("utf-8", errors="surrogatepass"))
                writers[key].add(ids, n_bytes)
                st["docs"] += 1
                st["tokens"] += len(ids) + 2
                st["bytes"] += n_bytes
                st["group"] = r.group
                st["licenses"].add(r.license)
            pending.clear()

        for r in _read_interim(path):
            pending.append(r)
            if len(pending) >= batch:
                flush()
        flush()
        st["licenses"] = sorted(st["licenses"])[:20]
        source_stats[name] = st

    groups: dict[str, dict] = defaultdict(dict)
    for (group, split), w in sorted(writers.items()):
        groups[group][split] = {"tokens": w.tokens, "bytes": w.bytes, "shards": w.close()}
    tok_sha = hashlib.sha256(Path(tokenizer_path).read_bytes()).hexdigest()
    write_manifest(out_dir, {"tokenizer_sha256": tok_sha, "groups": dict(groups), "sources": source_stats,
                             "decontamination": decontamination})


def main(argv=None) -> None:
    from zenith.data.sources import SOURCES

    ap = argparse.ArgumentParser(prog="python -m zenith.data.pipeline")
    ap.add_argument("stage", choices=["fetch", "normalize", "sample", "tokenize"])
    ap.add_argument("--sources", help="comma-separated subset of sources")
    ap.add_argument("--bytes", default="200MB", help="tokenizer sample size (sample stage)")
    ap.add_argument("--tokenizer", type=Path, default=TOKENIZER_HF)
    ap.add_argument("--no-decontam", action="store_true", help="tokenize without evaluation texts (not for real runs)")
    args = ap.parse_args(argv)
    sources = {k: v for k, v in SOURCES.items() if not args.sources or k in args.sources.split(",")}

    if args.stage == "fetch":
        for name, mod in sources.items():
            print(f"fetch {name}", flush=True)
            mod.fetch(RAW_DIR)
    elif args.stage == "normalize":
        normalize_stage(RAW_DIR, INTERIM_DIR, sources)
    elif args.stage == "sample":
        units = {"KB": 10**3, "MB": 10**6, "GB": 10**9}
        n = args.bytes.upper()
        total = int(float(n[:-2]) * units[n[-2:]]) if n[-2:] in units else int(n)
        sample_stage(INTERIM_DIR, INTERIM_DIR / "tokenizer_sample.txt", total, DEFAULT_MIX)
    else:
        tokenize_stage(INTERIM_DIR, PROCESSED_DIR, args.tokenizer, EVAL_DIR, sources, allow_no_decontam=args.no_decontam)


if __name__ == "__main__":
    main()