"""Byte-level BPE tokenizer for the phishing encoder: training, statistics, and batch encoding.

Design (config.tokenizer):
- NFKC normalisation only, case preserved. ALL-CAPS subjects ("URGENT", "VERIFY YOUR ACCOUNT"),
  odd casing in brand names and shouty pharma spam are strong phishing/spam signals; lowercasing
  would throw them away. NFKC folds full-width and ligature look-alikes (a common filter-evasion
  trick) onto their ASCII forms so they share subwords with the plain spelling.
- Byte-level pre-tokenizer with the full 256-symbol initial alphabet: every byte sequence is
  representable, so [UNK] is unreachable in practice. Email is full of URLs, base64 fragments,
  HTML entities and mojibake, which a character-level vocabulary would map to [UNK].
- add_prefix_space=False: inputs are "subject\\n\\nbody", not isolated words.
- 32,768 ids (2^15) to match model.vocab_size; [PAD] is id 0 because the model's embedding uses
  padding_idx=0 and the collate code pads with 0.
- Trained ONLY on pretrain.parquet + finetune_train.parquet. finetune_val and test_holdout
  (TREC-06) never touch the vocabulary: a vocabulary fitted on test text leaks its frequent
  strings (template URLs, spam phrases) into single tokens.
"""

from __future__ import annotations

import argparse
import logging
import sys
from collections.abc import Iterator, Sequence
from dataclasses import asdict, dataclass, field
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

import numpy as np
import pyarrow.parquet as pq
import torch
from tokenizers import Tokenizer, decoders, models, normalizers, pre_tokenizers, processors, trainers

import common

LOG = logging.getLogger("tokenizer")

SPLIT_FILES: dict[str, str] = {
    "pretrain": "pretrain.parquet",
    "finetune_train": "finetune_train.parquet",
    "finetune_val": "finetune_val.parquet",
    "test_holdout": "test_holdout.parquet",
}
FORBIDDEN_TRAIN_FILES = frozenset({"finetune_val.parquet", "test_holdout.parquet"})
SOURCE_COLUMNS = ("source", "dataset_name")
TEXT_COLUMN = "text"
READ_BATCH_ROWS = 2048
STATS_ENCODE_BATCH = 4096
PRETRAIN_STATS_SAMPLE = 20_000

EXAMPLES: list[tuple[str, str]] = [
    ("business", "Re: Q3 budget review\n\nHi Sarah, can we move the budget review to Thursday at 2pm? "
                 "I've attached the revised forecast. Thanks, Mike"),
    ("url_heavy", "Your account has been suspended\n\nVerify now: http://secure-login.paypa1-verify.com/"
                  "account/update?id=8f3a9c2e&session=XyZ123 or https://bit.ly/3xYz9Ab before 24 hours."),
    ("all_caps", "URGENT!!! FINAL NOTICE\n\nYOUR MAILBOX IS FULL. CLICK HERE TO UPGRADE YOUR STORAGE "
                 "IMMEDIATELY OR LOSE ALL MESSAGES."),
    ("unicode_lookalike", "Ｍｉｃｒｏｓｏｆｔ 365 — Password expiry ⚠️\n\nDear user, your pass­word expires "
                          "today. Café, naïve, 東京 office."),
    ("html_money", "<html><body><p>Dear Friend,</p><p>I am Barrister James. Transfer of US$15,500,000.00 "
                   "&amp; 30% for you.</p></body></html>"),
]


@dataclass
class TokenizerPaths:
    data_dir: Path
    tokenizer_json: Path
    report_json: Path
    log_dir: Path


@dataclass
class SplitStats:
    split: str
    file: str
    rows: int
    mean_tokens: float
    median_tokens: float
    p95_tokens: float
    max_tokens: int
    chars_per_token: float
    pct_truncated: float
    unk_rate: float
    mean_tokens_by_label: dict[str, float] = field(default_factory=dict)


def _abs(p: str | Path) -> Path:
    p = Path(p)
    return p if p.is_absolute() else common.project_root() / p


def resolve_paths(cfg: dict, data_dir: str | None) -> TokenizerPaths:
    """--data-dir redirects data, artifacts and logs together, so a self-test never writes into the project."""
    if data_dir is not None:
        d = Path(data_dir).resolve()
        return TokenizerPaths(d, d / "tokenizer.json", d / "tokenizer_report.json", d / "logs")
    d = _abs(cfg["paths"]["data_processed"])
    tok = _abs(cfg["paths"]["tokenizer"])
    return TokenizerPaths(d, tok, tok.with_name("tokenizer_report.json"), _abs(cfg["paths"]["logs"]))


# ---------------------------------------------------------------------------------------- leakage


def assert_trainable(path: Path, holdout_source: str) -> None:
    """Refuse validation/test files by name, and any file whose source column carries the holdout corpus."""
    if path.name in FORBIDDEN_TRAIN_FILES:
        raise RuntimeError(f"refusing to train the tokenizer on {path.name}: val/test text must never shape the vocabulary")
    schema_names = set(pq.read_schema(path).names)
    if TEXT_COLUMN not in schema_names:
        raise RuntimeError(f"{path} has no '{TEXT_COLUMN}' column (columns: {sorted(schema_names)})")
    for col in SOURCE_COLUMNS:
        if col in schema_names:
            values = {str(v) for v in pq.read_table(path, columns=[col]).column(col).unique().to_pylist()}
            hits = {v for v in values if v.strip().lower() == holdout_source.lower()}
            if hits:
                raise RuntimeError(f"{path.name} column '{col}' contains holdout source {hits}: test-set leakage")


def iter_texts(path: Path, limit: int | None = None) -> Iterator[list[str]]:
    """Stream the text column in record batches; the 244k-email pretraining corpus never sits in RAM twice."""
    seen = 0
    for batch in pq.ParquetFile(path).iter_batches(batch_size=READ_BATCH_ROWS, columns=[TEXT_COLUMN]):
        texts = [t for t in batch.column(0).to_pylist() if t]
        if limit is not None:
            texts = texts[: max(0, limit - seen)]
        if texts:
            seen += len(texts)
            yield texts
        if limit is not None and seen >= limit:
            return


def holdout_collisions(train_files: Sequence[Path], holdout: Path, limit: int | None) -> dict[str, int]:
    """Exact normalized-text collisions (verify_data.md §7 k_full) between training text and the holdout.

    The data pipeline already decontaminates; this is an independent tripwire at the point where
    text first becomes model state. Reading the holdout here only builds a hash set.
    """
    held = {common.sha1_hex(common.normalize_text(t)) for batch in iter_texts(holdout) for t in batch}
    out: dict[str, int] = {}
    for f in train_files:
        out[f.name] = sum(
            common.sha1_hex(common.normalize_text(t)) in held for batch in iter_texts(f, limit) for t in batch
        )
    return out


# --------------------------------------------------------------------------------------- building


def build_untrained() -> Tokenizer:
    tok = Tokenizer(models.BPE(unk_token="[UNK]"))
    tok.normalizer = normalizers.NFKC()
    tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tok.decoder = decoders.ByteLevel()
    return tok


def attach_post_processor(tok: Tokenizer) -> None:
    """BERT template; segment ids only matter for pairs, which the classifier does not use but the MLM loader may."""
    cls_id, sep_id = tok.token_to_id("[CLS]"), tok.token_to_id("[SEP]")
    tok.post_processor = processors.TemplateProcessing(
        single="[CLS] $A [SEP]",
        pair="[CLS] $A [SEP] $B:1 [SEP]:1",
        special_tokens=[("[CLS]", cls_id), ("[SEP]", sep_id)],
    )


def check_special_ids(tok: Tokenizer, special_tokens: Sequence[str]) -> dict[str, int]:
    ids = {t: tok.token_to_id(t) for t in special_tokens}
    if ids.get("[PAD]") != 0:
        raise AssertionError(f"[PAD] must be id 0 (model padding_idx), got {ids.get('[PAD]')}")
    if list(ids.values()) != list(range(len(special_tokens))):
        raise AssertionError(f"special tokens must occupy ids 0..{len(special_tokens) - 1} in order, got {ids}")
    return ids


def train(cfg: dict, paths: TokenizerPaths, vocab_size: int | None, limit_docs: int | None) -> Tokenizer:
    tcfg = cfg["tokenizer"]
    vocab_size = vocab_size or int(tcfg["vocab_size"])
    model_vocab = int(cfg["model"]["vocab_size"])
    if vocab_size > model_vocab:
        raise ValueError(f"tokenizer vocab {vocab_size} exceeds model.vocab_size {model_vocab}")
    special = list(tcfg["special_tokens"])
    holdout_source = cfg["data"]["holdout_test_source"]

    files = [paths.data_dir / SPLIT_FILES[name] for name in tcfg["train_on"]]
    for f in files:
        if not f.exists():
            raise FileNotFoundError(f)
        assert_trainable(f, holdout_source)

    holdout = paths.data_dir / SPLIT_FILES["test_holdout"]
    if holdout.exists():
        collisions = holdout_collisions(files, holdout, limit_docs)
        LOG.info("exact holdout collisions per training file: %s", collisions)
        if any(collisions.values()):
            raise RuntimeError(f"training text exactly matches TREC-06 holdout rows: {collisions}")
    else:
        collisions = None
        LOG.warning("%s absent: exact-hash leakage tripwire skipped (name/source guards still applied)", holdout.name)

    usage: dict[str, dict[str, int]] = {}

    def corpus() -> Iterator[list[str]]:
        for f in files:
            rows = chars = 0
            for batch in iter_texts(f, limit_docs):
                rows += len(batch)
                chars += sum(map(len, batch))
                yield batch
            usage[f.name] = {"rows_in_file": pq.ParquetFile(f).metadata.num_rows, "rows_used": rows, "chars": chars}
            LOG.info("fed %s: %d rows, %d chars", f.name, rows, chars)

    tok = build_untrained()
    trainer_kwargs: dict = dict(
        vocab_size=vocab_size,
        min_frequency=int(tcfg["min_frequency"]),
        special_tokens=special,
        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
        show_progress=True,
    )
    if tcfg.get("max_token_length"):
        trainer_kwargs["max_token_length"] = int(tcfg["max_token_length"])
    LOG.info("training byte-level BPE: %s on %s (limit_docs=%s)",
             {k: v for k, v in trainer_kwargs.items() if k != "initial_alphabet"}, [f.name for f in files], limit_docs)
    t0 = common.timer()
    tok.train_from_iterator(corpus(), trainer=trainers.BpeTrainer(**trainer_kwargs))
    elapsed = common.timer() - t0

    special_ids = check_special_ids(tok, special)
    attach_post_processor(tok)
    final_vocab = tok.get_vocab_size()
    if final_vocab > model_vocab:
        raise AssertionError(f"trained vocab {final_vocab} > model.vocab_size {model_vocab}")
    if final_vocab < vocab_size:
        LOG.warning("vocab reached only %d of %d (corpus too small / min_frequency); embedding rows above stay unused",
                    final_vocab, vocab_size)

    paths.tokenizer_json.parent.mkdir(parents=True, exist_ok=True)
    tok.save(str(paths.tokenizer_json))
    LOG.info("saved %s (vocab %d) in %.1fs", paths.tokenizer_json, final_vocab, elapsed)

    report = _read_report(paths)
    report["train"] = {
        "tokenizer_json": str(paths.tokenizer_json),
        "tokenizer_sha1": common.sha1_hex(paths.tokenizer_json.read_text()),
        "requested_vocab_size": vocab_size,
        "final_vocab_size": final_vocab,
        "min_frequency": int(tcfg["min_frequency"]),
        "special_token_ids": special_ids,
        "normalizer": "NFKC (case preserved)",
        "pre_tokenizer": "ByteLevel(add_prefix_space=False)",
        "training_files": usage,
        "limit_docs_per_file": limit_docs,
        "excluded_files": sorted(FORBIDDEN_TRAIN_FILES),
        "holdout_exact_collisions": collisions,
        "train_seconds": round(elapsed, 1),
    }
    common.write_json(report, paths.report_json)
    return tok


# ---------------------------------------------------------------------------------------- loading


def load_tokenizer(path: str | Path) -> Tokenizer:
    """Load tokenizer.json with truncation/padding cleared; callers configure them (or use encode_batch)."""
    tok = Tokenizer.from_file(str(path))
    if tok.token_to_id("[PAD]") != 0:
        raise AssertionError(f"{path}: [PAD] is not id 0")
    tok.no_truncation()
    tok.no_padding()
    return tok


def encode_batch(
    tok: Tokenizer, texts: Sequence[str], max_len: int, pad_to_multiple_of: int | None = 8
) -> tuple[torch.Tensor, torch.Tensor]:
    """Encode with [CLS]/[SEP], truncate to max_len, pad to the longest in the batch (dynamic padding).

    pad_to_multiple_of=8 keeps sequence lengths tensor-core friendly; it must divide max_len so a
    full-length row is never padded past the model's position table.
    Returns (input_ids int64 [B, L], attention_mask bool [B, L]).
    """
    if pad_to_multiple_of and max_len % pad_to_multiple_of:
        raise ValueError(f"max_len {max_len} is not a multiple of pad_to_multiple_of {pad_to_multiple_of}")
    if len(texts) == 0:
        return torch.empty((0, 0), dtype=torch.int64), torch.empty((0, 0), dtype=torch.bool)
    # Truncation/padding are configured on the shared Tokenizer object for this call only and
    # cleared afterwards, so a later tok.encode() on the same object returns raw ids again.
    tok.enable_truncation(max_length=max_len)
    tok.enable_padding(pad_id=0, pad_token="[PAD]", pad_to_multiple_of=pad_to_multiple_of)
    try:
        encs = tok.encode_batch_fast(list(texts))
    finally:
        tok.no_truncation()
        tok.no_padding()
    ids = np.asarray([e.ids for e in encs], dtype=np.int64).reshape(len(texts), -1)
    mask = np.asarray([e.attention_mask for e in encs], dtype=np.bool_).reshape(len(texts), -1)
    return torch.from_numpy(ids), torch.from_numpy(mask)


# ------------------------------------------------------------------------------------------ stats


def _read_report(paths: TokenizerPaths) -> dict:
    return common.read_json(paths.report_json) if paths.report_json.exists() else {}


def _load_split(path: Path, sample: int | None, seed: int) -> tuple[list[str], list | None]:
    cols = [TEXT_COLUMN] + (["label"] if "label" in pq.read_schema(path).names else [])
    table = pq.read_table(path, columns=cols)
    if sample is not None and table.num_rows > sample:
        idx = np.sort(np.random.default_rng(seed).choice(table.num_rows, size=sample, replace=False))
        table = table.take(idx)
    texts = [t or "" for t in table.column(TEXT_COLUMN).to_pylist()]
    labels = table.column("label").to_pylist() if "label" in cols else None
    return texts, labels


def split_stats(tok: Tokenizer, split: str, path: Path, max_len: int, sample: int | None, seed: int) -> SplitStats:
    """Content-token counts (no specials); truncation is judged on content + [CLS] + [SEP] > max_len."""
    texts, labels = _load_split(path, sample, seed)
    tok.no_truncation()
    tok.no_padding()
    unk_id = tok.token_to_id("[UNK]")
    lengths = np.empty(len(texts), dtype=np.int64)
    unk = 0
    for i in range(0, len(texts), STATS_ENCODE_BATCH):
        encs = tok.encode_batch_fast(texts[i : i + STATS_ENCODE_BATCH], add_special_tokens=False)
        for j, e in enumerate(encs):
            lengths[i + j] = len(e.ids)
            unk += e.ids.count(unk_id)
    total_tokens = int(lengths.sum())
    by_label: dict[str, float] = {}
    if labels is not None:
        lab = np.asarray(labels)
        by_label = {str(v): round(float(lengths[lab == v].mean()), 1) for v in np.unique(lab)}
    return SplitStats(
        split=split,
        file=path.name,
        rows=len(texts),
        mean_tokens=round(float(lengths.mean()), 1),
        median_tokens=float(np.median(lengths)),
        p95_tokens=float(np.percentile(lengths, 95)),
        max_tokens=int(lengths.max()),
        chars_per_token=round(sum(map(len, texts)) / max(total_tokens, 1), 3),
        pct_truncated=round(100.0 * float((lengths + 2 > max_len).mean()), 2),
        unk_rate=unk / max(total_tokens, 1),
        mean_tokens_by_label=by_label,
    )


def example_encodings(tok: Tokenizer) -> list[dict]:
    tok.no_truncation()
    tok.no_padding()
    out = []
    for name, text in EXAMPLES:
        enc = tok.encode(text)
        decoded = tok.decode(enc.ids)
        out.append({
            "name": name,
            "text": text,
            "n_tokens_with_specials": len(enc.ids),
            "tokens": enc.tokens,
            "decoded": decoded,
            "round_trip_exact": decoded == text,
            "round_trip_matches_nfkc": decoded == normalizers.NFKC().normalize_str(text),
        })
    return out


def stats(cfg: dict, paths: TokenizerPaths) -> dict:
    max_len = int(cfg["model"]["max_seq_len"])
    tok = load_tokenizer(paths.tokenizer_json)
    rows: list[SplitStats] = []
    for split in ("pretrain", "finetune_train", "finetune_val", "test_holdout"):
        f = paths.data_dir / SPLIT_FILES[split]
        if not f.exists():
            LOG.warning("%s missing, skipped", f)
            continue
        sample = PRETRAIN_STATS_SAMPLE if split == "pretrain" else None
        t0 = common.timer()
        rows.append(split_stats(tok, split, f, max_len, sample, int(cfg["seed"])))
        LOG.info("%s: %d rows in %.1fs", split, rows[-1].rows, common.timer() - t0)
    examples = example_encodings(tok)

    report = _read_report(paths)
    report["stats"] = {
        "max_seq_len": max_len,
        "vocab_size": tok.get_vocab_size(),
        "pretrain_sample_rows": PRETRAIN_STATS_SAMPLE,
        "splits": [asdict(r) for r in rows],
        "examples": examples,
    }
    common.write_json(report, paths.report_json)
    _print_report(rows, examples, max_len, tok.get_vocab_size())
    return report


def _print_report(rows: list[SplitStats], examples: list[dict], max_len: int, vocab: int) -> None:
    print(f"\n### Tokenizer statistics (vocab {vocab:,}, max_seq_len {max_len}; token counts exclude [CLS]/[SEP])\n")
    print(f"| split | rows | mean tok | median | p95 | max | chars/tok | % > {max_len} | [UNK] rate | mean tok by label |")
    print("|---|---:|---:|---:|---:|---:|---:|---:|---:|---|")
    for r in rows:
        lab = ", ".join(f"{k}: {v}" for k, v in r.mean_tokens_by_label.items()) or "-"
        print(f"| {r.split} | {r.rows:,} | {r.mean_tokens:,.1f} | {r.median_tokens:,.0f} | {r.p95_tokens:,.0f} | "
              f"{r.max_tokens:,} | {r.chars_per_token:.2f} | {r.pct_truncated:.1f}% | {r.unk_rate:.2e} | {lab} |")
    print("\n### Example encodings\n")
    for e in examples:
        shown = " ".join(e["tokens"][:48]) + (" ..." if len(e["tokens"]) > 48 else "")
        print(f"**{e['name']}** ({e['n_tokens_with_specials']} tokens; round trip exact={e['round_trip_exact']}, "
              f"matches NFKC(input)={e['round_trip_matches_nfkc']})")
        print(f"- tokens: `{shown}`")
        print(f"- decoded: {e['decoded']!r}\n")


# -------------------------------------------------------------------------------------------- CLI


def main(argv: Sequence[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--config", default="config.yaml")
    ap.add_argument("--data-dir", default=None, help=argparse.SUPPRESS)  # self-test override
    sub = ap.add_subparsers(dest="cmd", required=True)
    tr = sub.add_parser("train", help="train byte-level BPE on pretrain + finetune_train")
    tr.add_argument("--vocab-size", type=int, default=None)
    tr.add_argument("--limit-docs", type=int, default=None, help="dev only: cap rows read per training file")
    sub.add_parser("stats", help="token-length / truncation / [UNK] statistics per split")
    args = ap.parse_args(argv)

    cfg = common.load_config(args.config)
    paths = resolve_paths(cfg, args.data_dir)
    global LOG
    LOG = common.setup_logging("tokenizer", paths.log_dir)
    common.set_seed(int(cfg["seed"]))
    if args.cmd == "train":
        train(cfg, paths, args.vocab_size, args.limit_docs)
    else:
        stats(cfg, paths)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())