#!/usr/bin/env python3
"""
SynapCore Phobos-Scratch v1 — 32,768-Vocab Byte-Level BPE Tokenizer Trainer & Wrapper
(`phobos-1/models/phobos_scratch_v1/bpe_tokenizer.py`)
=====================================================================================
Trains a 100% from-scratch Byte-Level BPE tokenizer (`vocab_size = 32768`) using the
Rust-backed Hugging Face `tokenizers` library on the canonical Phobos training corpus:
  • Zero `<unk>` tokens (full UTF-8 byte fallback for code, math LaTeX, and Hinglish).
  • Dedicated structural & Socratic control tokens (`<|pad|>`, `<|bos|>`, `<|eos|>`,
    `<|user|>`, `<|assistant|>`, `<think>`, `</think>`, `<tool_call>`, `</tool_call>`,
    plus cortical branch tags `<|branch:E0|>..<|branch:E3|>`).
  • Preserves code indentation and multi-digit number splitting (`Digits(individual_digits=True)`).

Strictly <= 200 lines.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Dict, Iterator, List, Optional

from tokenizers import Tokenizer, decoders, models, pre_tokenizers, processors, trainers

TOKENIZER_DIR: Path = (
    Path(__file__).resolve().parent / "tokenizer"
    if "__file__" in globals()
    else Path("/content/phobos_training_data")
)
DEFAULT_TOKENIZER_PATH: Path = TOKENIZER_DIR / "phobos_bpe_32k.json"

SPECIAL_TOKENS: List[str] = [
    "<|pad|>",
    "<|bos|>",
    "<|eos|>",
    "<|user|>",
    "<|assistant|>",
    "<think>",
    "</think>",
    "<tool_call>",
    "</tool_call>",
    "<|branch:E0|>",
    "<|branch:E1|>",
    "<|branch:E2|>",
    "<|branch:E3|>",
]


def format_canonical_record_for_lm(rec: Dict[str, Any]) -> str:
    """Formats a canonical 8-key JSONL record into the unified Phobos causal training string."""
    user_txt = str(rec.get("user") or "").strip()
    think_txt = str(rec.get("think") or "").strip()
    asst_txt = str(rec.get("assistant") or "").strip()
    if think_txt:
        return f"<|bos|><|user|>\n{user_txt}\n<|assistant|>\n<think>\n{think_txt}\n</think>\n{asst_txt}<|eos|>"
    return f"<|bos|><|user|>\n{user_txt}\n<|assistant|>\n{asst_txt}<|eos|>"


def _iter_balanced_corpus(
    data_dir: Path, max_rows_per_phase: int = 60000
) -> Iterator[str]:
    """Streams balanced training texts across Phase 1, Phase 2, and Phase 4 JSONL files."""
    for fname in ("phase1_foundation.jsonl", "phase2_specialists.jsonl", "phase4_socratic_think.jsonl"):
        fpath = data_dir / fname
        if not fpath.exists():
            continue
        count = 0
        with open(fpath, "r", encoding="utf-8", errors="ignore") as f:
            for line in f:
                if count >= max_rows_per_phase:
                    break
                try:
                    rec = json.loads(line)
                    yield format_canonical_record_for_lm(rec)
                    count += 1
                except Exception:
                    continue


def train_phobos_bpe_tokenizer(
    data_dir: Path,
    save_path: Path = DEFAULT_TOKENIZER_PATH,
    vocab_size: int = 32768,
    max_rows_per_phase: int = 60000,
) -> Dict[str, Any]:
    """
    Trains the 32,768-vocab Byte-Level BPE tokenizer on the canonical Phobos JSONL corpus
    and verifies lossless roundtrip + compression ratio across domains.
    """
    tokenizer = Tokenizer(models.BPE(byte_fallback=True))
    tokenizer.pre_tokenizer = pre_tokenizers.Sequence(
        [
            pre_tokenizers.Digits(individual_digits=True),
            pre_tokenizers.ByteLevel(add_prefix_space=False),
        ]
    )
    tokenizer.decoder = decoders.ByteLevel()
    tokenizer.post_processor = processors.ByteLevel(trim_offsets=False)

    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        min_frequency=2,
        show_progress=False,
        special_tokens=SPECIAL_TOKENS,
        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
    )

    tokenizer.train_from_iterator(
        _iter_balanced_corpus(data_dir, max_rows_per_phase=max_rows_per_phase),
        trainer=trainer,
    )

    save_path.parent.mkdir(parents=True, exist_ok=True)
    tokenizer.save(str(save_path))

    # Verify roundtrip and compression across all Phobos domains
    probe_samples = {
        "conversation": "<|bos|><|user|>\nExplain how standard railway gauges were created.\n<|assistant|>\nStandard railway gauges were established to allow trains to move across regions.<|eos|>",
        "python_code": "<|bos|><|user|>\nWrite binary search in Python.\n<|assistant|>\n```python\ndef binary_search(arr: list[int], target: int) -> int:\n    lo, hi = 0, len(arr) - 1\n    return lo\n```<|eos|>",
        "math_cot": "<|bos|><|user|>\nSolve 4x - 5 = 19.\n<|assistant|>\n<think>\nAdd 5 to both sides: 4x = 24. Divide by 4: x = 6.\n</think>\n\\boxed{6}<|eos|>",
        "sql_tool": "<|bos|><|user|>\nQuery users table.\n<|assistant|>\n<tool_call>\n{\"query\": \"SELECT id, name FROM users WHERE active = 1;\"}\n</tool_call><|eos|>",
    }
    probes_out: Dict[str, Any] = {}
    for name, text in probe_samples.items():
        enc = tokenizer.encode(text)
        dec = tokenizer.decode(enc.ids, skip_special_tokens=False)
        probes_out[name] = {
            "char_len": len(text),
            "token_len": len(enc.ids),
            "chars_per_token": round(len(text) / max(1, len(enc.ids)), 2),
            "exact_roundtrip": (dec == text),
        }

    return {
        "vocab_size": tokenizer.get_vocab_size(),
        "save_path": str(save_path),
        "special_token_ids": {tok: tokenizer.token_to_id(tok) for tok in SPECIAL_TOKENS},
        "domain_probes": probes_out,
    }


class PhobosBPETokenizer:
    """Runtime wrapper around `phobos_bpe_32k.json` compatible with `PhobosScratchTrainer`."""

    def __init__(self, tokenizer_path: Path = DEFAULT_TOKENIZER_PATH) -> None:
        self.backend = Tokenizer.from_file(str(tokenizer_path))
        self.PAD_ID: int = int(self.backend.token_to_id("<|pad|>"))
        self.BOS_ID: int = int(self.backend.token_to_id("<|bos|>"))
        self.EOS_ID: int = int(self.backend.token_to_id("<|eos|>"))
        self.USER_ID: int = int(self.backend.token_to_id("<|user|>"))
        self.ASST_ID: int = int(self.backend.token_to_id("<|assistant|>"))
        self.THINK_START_ID: int = int(self.backend.token_to_id("<think>"))
        self.THINK_END_ID: int = int(self.backend.token_to_id("</think>"))
        self.VOCAB_SIZE: int = int(self.backend.get_vocab_size())

    def encode(self, text: str, max_len: int = 512) -> List[int]:
        ids = self.backend.encode(text).ids
        if len(ids) > max_len:
            ids = ids[: max_len - 1] + [self.EOS_ID]
        return ids

    def decode(self, token_ids: List[int], skip_special_tokens: bool = False) -> str:
        clean = [int(t) for t in token_ids if int(t) != self.PAD_ID]
        return self.backend.decode(clean, skip_special_tokens=skip_special_tokens)


if __name__ == "__main__":
    data_root = Path("/content/phobos_training_data")
    out_tok = data_root / "phobos_bpe_32k.json"
    summary = train_phobos_bpe_tokenizer(data_root, save_path=out_tok, vocab_size=32768, max_rows_per_phase=50000)
    print(json.dumps(summary, indent=2))