"""Canonical corpus helpers shared by preparation, training and loss evaluation.

This module deliberately has no torch/transformers dependency. Offsets are
Python character offsets, not UTF-8 byte offsets; byte-level tokens may share
an offset when they encode different bytes of the same Unicode character.
"""
import collections
import hashlib
import json
from pathlib import Path
import random

import gen_data as G


class MaskBoundaryError(ValueError):
    """A tokenizer merged model-written and app-written characters."""

    def __init__(self, issues):
        self.issues = issues
        first = issues[0]
        super().__init__(
            f"unsafe loss boundary: token {first['token_index']} (id {first['token_id']}) "
            f"at characters {first['start']}:{first['end']} encodes "
            f"{first['token_text']!r}; {first['trained_characters']} trained and "
            f"{first['untrained_characters']} untrained characters "
            f"({len(issues)} crossing token(s)). Prepare a structurally split tokenizer; "
            "do not supervise lookup results or repair the canonical spans."
        )


def sha256_file(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def file_hashes(directory):
    """Hash actual top-level checkpoint files, including all weight shards."""
    return {p.name: sha256_file(p) for p in sorted(Path(directory).iterdir()) if p.is_file()}


def save_tokenizer(tok, directory):
    """Save standard HF JSON using the browser-compatible BPE merge encoding."""
    tok.save_pretrained(directory)
    path = Path(directory) / "tokenizer.json"
    payload = json.loads(path.read_text(encoding="utf-8"))
    if payload["model"]["type"] == "BPE":
        payload["model"]["merges"] = [" ".join(m) if isinstance(m, list) else m
                                      for m in payload["model"]["merges"]]
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def evaluation_contract(seed, samples, diaries, cell_samples, semantic=None):
    """Historical suites keep v3 identity; explicit semantic manifests add v4."""
    contract = {"version": 3, "diary_safety": "drop-whole-page", "lookup_counts": "generated-fragments", "seed": seed, "samples": samples, "diaries": diaries,
                "cell_samples": cell_samples,
                "sampling": {"temperature": 0.85, "top_k": 40, "top_p": 0.92},
                "budgets": {"idle": [512, 40], "react": [1024, 64], "diary": [1024, 220]}}
    if semantic is not None:
        contract.update(version=4, semantic=dict(semantic))
    return contract


def rows_digest(rows):
    h = hashlib.sha256()
    for row in rows:
        h.update(json.dumps(row, ensure_ascii=False, sort_keys=True,
                            separators=(",", ":")).encode("utf-8"))
        h.update(b"\n")
    return h.hexdigest()


def load_rows(path, drop_stale=False):
    """Recheck the stored text and mask before augmentation, never rewrite them."""
    rows, stale, examples = [], collections.Counter(), []
    total = 0
    with open(path, encoding="utf-8") as f:
        for line_number, line in enumerate(f, 1):
            total += 1
            try:
                row = json.loads(line)
            except json.JSONDecodeError as e:
                raise ValueError(f"{path}:{line_number}: invalid JSON: {e.msg}") from e
            try:
                if not isinstance(row, dict) or row.get("kind") not in ("chat", "diary"):
                    raise ValueError("row kind must be chat or diary")
                if not isinstance(row.get("text"), str):
                    raise ValueError("row text must be a string")
                G.recheck(row)
            except (ValueError, KeyError, TypeError, IndexError) as e:
                reason = str(e).split(":")[0]
                stale[reason] += 1
                if len(examples) < 12:
                    examples.append({"line": line_number, "reason": str(e)})
            else:
                rows.append(row)
    diagnostics = {"input_rows": total, "accepted_rows": len(rows),
                   "rejected_rows": sum(stale.values()), "reasons": dict(stale),
                   "examples": examples}
    if stale and not drop_stale:
        reasons = "; ".join(f"{n} {why}" for why, n in stale.most_common())
        raise ValueError(f"{path}: {sum(stale.values())} rows fail the current rules: {reasons}; "
                         f"first failure at line {examples[0]['line']}: {examples[0]['reason']}")
    if not rows:
        raise ValueError(f"{path}: no canonical rows")
    return rows, diagnostics


def load_language_rows(path):
    """Explicit unmasked language input, never an escape hatch for pet rows."""
    rows = []
    with open(path, encoding="utf-8") as f:
        for line_number, line in enumerate(f, 1):
            try:
                row = json.loads(line)
            except json.JSONDecodeError as e:
                raise ValueError(f"{path}:{line_number}: invalid JSON: {e.msg}") from e
            if (not isinstance(row, dict) or set(row) - {"text", "source"}
                    or not isinstance(row.get("text"), str) or not row["text"].strip()):
                raise ValueError(f"{path}:{line_number}: pretraining rows require nonempty text "
                                 "and optional source metadata only; pet kind/spans are forbidden")
            rows.append(row)
    if not rows:
        raise ValueError(f"{path}: no language rows")
    return rows


def encode_language(tok, text, max_len):
    """Unmasked causal language supervision, with an explicit document EOS."""
    if max_len is not None and (type(max_len) is not int or max_len < 2):
        raise ValueError("max_len must be at least two tokens, or None for the full row")
    ids = list(tok(text, add_special_tokens=False)["input_ids"])
    if type(tok.eos_token_id) is not int or tok.eos_token_id < 0:
        raise ValueError("language pretraining requires a real EOS token ID")
    if tok.decode(ids, skip_special_tokens=False, clean_up_tokenization_spaces=False) != text:
        raise ValueError("language tokenizer does not preserve the exact input text")
    ids.append(tok.eos_token_id)
    return ids[:max_len], ids[:max_len]


def heldout_split(rows, seed):
    """The same seeded split for every tokenizer, trainer and quantization eval.

    Normal corpora hold out at least 64 rows. Tiny smoke corpora still reserve
    a disjoint evaluation row and a training row rather than training on eval.
    """
    if len(rows) < 2:
        raise ValueError("training/held-out loss requires at least two canonical rows")
    shuffled = list(rows)
    random.Random(seed).shuffle(shuffled)
    n_eval = min(len(shuffled) - 1, max(64, len(shuffled) // 50))
    return shuffled[n_eval:], shuffled[:n_eval]


def inspect_row(tok, row):
    """Return token/mask diagnostics; crossing tokens stay masked, never trained.

    Callers must reject mixed_tokens before using the labels for model loss.
    Coverage checks also reject normalization/offset trimming that hides text.
    """
    text = row["text"]
    spans = G._span_coverage(row["spans"], len(text))
    trained = bytearray(len(text))
    for a, b in spans:
        trained[a:b] = b"\x01" * (b - a)
    prefix = [0]
    for value in trained:
        prefix.append(prefix[-1] + value)
    enc = tok(text, return_offsets_mapping=True, add_special_tokens=False)
    ids, offsets = list(enc["input_ids"]), list(enc["offset_mapping"])
    if len(ids) != len(offsets):
        raise ValueError("tokenizer IDs and character offsets have different lengths")
    labels, issues = [], []
    covered = bytearray(len(text))
    for index, (token_id, offset) in enumerate(zip(ids, offsets)):
        if len(offset) != 2:
            raise ValueError(f"token {index}: malformed character offsets {offset!r}")
        a, b = offset
        if type(a) is not int or type(b) is not int or not 0 <= a < b <= len(text):
            raise ValueError(f"token {index}: invalid character offsets {offset!r} for "
                             f"{len(text)} characters; disable automatic tokens and offset trimming")
        covered[a:b] = b"\x01" * (b - a)
        count = prefix[b] - prefix[a]
        if 0 < count < b - a:
            issues.append({"token_index": index, "token_id": token_id,
                           "start": a, "end": b, "token_text": text[a:b],
                           "trained_characters": count, "untrained_characters": b - a - count,
                           "boundary_positions": [p for p in range(a + 1, b)
                                                  if trained[p] != trained[p - 1]]})
        labels.append(token_id if count == b - a else -100)
    if covered and 0 in covered:
        at = covered.index(0)
        raise ValueError(f"tokenizer offsets leave character {at} ({text[at:at + 16]!r}) "
                         "uncovered; use identity normalization and untrimmed offsets")
    if row["kind"] == "diary":
        eos = tok.eos_token_id
        if type(eos) is not int or eos < 0:
            raise ValueError("diary training requires a real EOS token ID")
        ids.append(eos)
        labels.append(eos)
        offsets.append((len(text), len(text)))
    return {"ids": ids, "labels": labels, "offsets": offsets, "mixed_tokens": issues}


def tokenize_row(tok, row):
    inspected = inspect_row(tok, row)
    if inspected["mixed_tokens"]:
        raise MaskBoundaryError(inspected["mixed_tokens"])
    return inspected["ids"], inspected["labels"], inspected["offsets"]


def encode(tok, row, max_len):
    """Encode canonical character spans; reject even crossings beyond the cap."""
    if max_len is not None and (type(max_len) is not int or max_len < 2):
        raise ValueError("max_len must be at least two tokens, or None for the full row")
    ids, labels, _ = tokenize_row(tok, row)
    return ids[:max_len], labels[:max_len]