# ─────────────────────────────────────────────────────────
# EGen Labs | ErebusTN
# AXIOM-1 — EGen-Core/AXIOM-1
# File: axiom1/tokenizer/tokenizer.py
# ─────────────────────────────────────────────────────────

"""
AXIOM-1 Custom Byte-Level BPE Tokenizer.

Built from scratch for the EGen Core architecture.
Vocabulary: exactly 132,000 tokens (including special tokens).
Compatible with LLaMA-3/Qwen-style byte-level BPE.

Special tokens (fixed IDs):
  <pad>  → 0
  <s>    → 1  (BOS)
  </s>   → 2  (EOS)
  <unk>  → 0  (reuses <pad> to avoid unknown tokens in loss)
  <bot>  → 3  (future SPIRAL latent thought)
  <eot>  → 4  (future SPIRAL latent thought)
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Union

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


# ── Fixed special token IDs ─────────────────────────────
SPECIAL_TOKENS = {
    "<pad>": 0,
    "<s>":   1,
    "</s>":  2,
    "<bot>": 3,
    "<eot>": 4,
}

VOCAB_SIZE = 132_000


class AXIOM1Tokenizer:
    """
    Custom byte-level BPE tokenizer for AXIOM-1.

    Features:
      - Exactly 132,000 tokens in vocabulary
      - Byte-level BPE (handles any UTF-8 input)
      - Fixed special token IDs matching AXIOM1Config
      - encode() / decode() with padding and truncation
      - Save/load to disk for reproducibility

    Usage:
        # Train from data:
        tok = AXIOM1Tokenizer.train(files=["corpus.txt"], save_dir="axiom1/tokenizer")

        # Load from saved:
        tok = AXIOM1Tokenizer.load("axiom1/tokenizer")

        # Encode:
        ids = tok.encode("Hello world", add_special_tokens=True)

        # Decode:
        text = tok.decode(ids)
    """

    def __init__(self, tokenizer: Tokenizer) -> None:
        self._tokenizer = tokenizer

    # ── Properties ──────────────────────────────────────
    @property
    def vocab_size(self) -> int:
        """Return the vocabulary size (always 132,000)."""
        return self._tokenizer.get_vocab_size()

    @property
    def pad_token_id(self) -> int:
        return SPECIAL_TOKENS["<pad>"]

    @property
    def bos_token_id(self) -> int:
        return SPECIAL_TOKENS["<s>"]

    @property
    def eos_token_id(self) -> int:
        return SPECIAL_TOKENS["</s>"]

    @property
    def unk_token_id(self) -> int:
        return SPECIAL_TOKENS["<pad>"]  # reuse pad to avoid unknowns in loss

    @property
    def bot_token_id(self) -> int:
        return SPECIAL_TOKENS["<bot>"]

    @property
    def eot_token_id(self) -> int:
        return SPECIAL_TOKENS["<eot>"]

    # ── Encode ──────────────────────────────────────────
    def encode(
        self,
        text: str,
        add_special_tokens: bool = True,
        padding: bool = False,
        truncation: bool = False,
        max_length: int | None = None,
    ) -> list[int]:
        """
        Encode text to a list of token IDs.

        Args:
            text: Input string
            add_special_tokens: If True, prepend <s> and append </s>
            padding: If True, pad to max_length with <pad>
            truncation: If True, truncate from the right to max_length
            max_length: Target sequence length for padding/truncation

        Returns:
            List of integer token IDs
        """
        encoding = self._tokenizer.encode(text)
        ids = encoding.ids

        if add_special_tokens:
            ids = [self.bos_token_id] + ids + [self.eos_token_id]

        if truncation and max_length is not None and len(ids) > max_length:
            ids = ids[:max_length]

        if padding and max_length is not None and len(ids) < max_length:
            ids = ids + [self.pad_token_id] * (max_length - len(ids))

        return ids

    def encode_batch(
        self,
        texts: list[str],
        add_special_tokens: bool = True,
        padding: bool = False,
        truncation: bool = False,
        max_length: int | None = None,
    ) -> list[list[int]]:
        """Encode a batch of texts."""
        return [
            self.encode(t, add_special_tokens, padding, truncation, max_length)
            for t in texts
        ]

    # ── Decode ──────────────────────────────────────────
    def decode(self, ids: list[int], skip_special_tokens: bool = True) -> str:
        """
        Decode token IDs back to a string.

        Args:
            ids: List of token IDs
            skip_special_tokens: If True, remove special tokens from output

        Returns:
            Decoded string
        """
        if skip_special_tokens:
            special_ids = set(SPECIAL_TOKENS.values())
            ids = [i for i in ids if i not in special_ids]

        return self._tokenizer.decode(ids)

    # ── Save / Load ─────────────────────────────────────
    def save(self, save_dir: str | Path) -> None:
        """
        Save tokenizer to disk.

        Creates:
          - tokenizer.json  (HuggingFace tokenizers format)
          - config.json     (AXIOM-1 metadata)

        Args:
            save_dir: Directory to save into
        """
        save_dir = Path(save_dir)
        save_dir.mkdir(parents=True, exist_ok=True)

        self._tokenizer.save(str(save_dir / "tokenizer.json"))

        config = {
            "model": "AXIOM-1",
            "lab": "EGen Labs",
            "author": "ErebusTN",
            "vocab_size": self.vocab_size,
            "special_tokens": SPECIAL_TOKENS,
            "type": "byte-level-bpe",
        }
        with open(save_dir / "config.json", "w") as f:
            json.dump(config, f, indent=2)

    @classmethod
    def load(cls, load_dir: str | Path) -> "AXIOM1Tokenizer":
        """
        Load a saved AXIOM1Tokenizer from disk.

        Args:
            load_dir: Directory containing tokenizer.json and config.json

        Returns:
            AXIOM1Tokenizer instance
        """
        load_dir = Path(load_dir)
        tokenizer = Tokenizer.from_file(str(load_dir / "tokenizer.json"))
        return cls(tokenizer)

    # ── Train ───────────────────────────────────────────
    @classmethod
    def train(
        cls,
        files: list[str | Path] | None = None,
        texts: list[str] | None = None,
        save_dir: str | Path = "axiom1/tokenizer",
        vocab_size: int = VOCAB_SIZE,
    ) -> "AXIOM1Tokenizer":
        """
        Train a new byte-level BPE tokenizer from scratch.

        Either `files` (paths to .txt files) or `texts` (list of strings)
        must be provided. The trained tokenizer is saved to `save_dir`.

        Args:
            files: List of text file paths to train on
            texts: List of raw text strings to train on
            save_dir: Where to save the trained tokenizer
            vocab_size: Target vocabulary size (default: 32,000)

        Returns:
            Trained AXIOM1Tokenizer instance
        """
        special_tokens_list = list(SPECIAL_TOKENS.keys())

        # Build a byte-level BPE tokenizer
        tokenizer = Tokenizer(models.BPE(unk_token="<pad>"))
        tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
        tokenizer.decoder = decoders.ByteLevel()

        trainer = trainers.BpeTrainer(
            vocab_size=vocab_size,
            special_tokens=special_tokens_list,
            show_progress=True,
            min_frequency=2,
        )

        if files:
            tokenizer.train([str(f) for f in files], trainer)
        elif texts:
            tokenizer.train_from_iterator(texts, trainer)
        else:
            raise ValueError("Provide either `files` or `texts` to train on.")

        # NOTE: We do NOT set a TemplateProcessing post-processor here.
        # BOS/EOS are added manually in encode() to give full control
        # over when special tokens appear (avoids double BOS/EOS).

        instance = cls(tokenizer)
        instance.save(save_dir)
        print(f"[AXIOM-1 Tokenizer] Trained with {instance.vocab_size} tokens.")
        print(f"[AXIOM-1 Tokenizer] Saved to {save_dir}/")
        return instance
