"""BPE Tokenizer for SplitBit LLM — trains vocabulary from raw text.

Uncensored: no content filtering. Vocab size auto-adjusts per hardware tier.
Encodes/decodes text <-> token IDs. Save/load to disk.
"""

from __future__ import annotations

import json
import logging
import os
import re
from collections import Counter, defaultdict
from pathlib import Path
from typing import Dict, List, Tuple

import numpy as np

logger = logging.getLogger(__name__)

# Special tokens
PAD_TOKEN = "<pad>"
EOS_TOKEN = "<eos>"
BOS_TOKEN = "<bos>"
UNK_TOKEN = "<unk>"

SPECIAL_TOKENS = [PAD_TOKEN, BOS_TOKEN, EOS_TOKEN, UNK_TOKEN]

PAD_ID = 0
BOS_ID = 1
EOS_ID = 2
UNK_ID = 3


class BPETokenizer:
    """Byte-Pair Encoding tokenizer trained from raw text.

    Training is simple and fast:
    1. Split text into words (whitespace + punctuation)
    2. Start with single characters as tokens
    3. Repeatedly merge most frequent adjacent pairs
    4. Stop when vocab_size reached or no more merges

    No content filtering — uncensored by design.
    """

    def __init__(self, vocab_size: int = 4096) -> None:
        self.target_vocab_size = vocab_size
        self.vocab: Dict[str, int] = {}
        self.inv_vocab: Dict[int, str] = {}
        self.merges: List[Tuple[str, str]] = []
        self._word_cache: Dict[str, List[int]] = {}

        # Initialize with special tokens
        for i, tok in enumerate(SPECIAL_TOKENS):
            self.vocab[tok] = i
            self.inv_vocab[i] = tok

    @property
    def actual_vocab_size(self) -> int:
        return len(self.vocab)

    def _tokenize_words(self, text: str) -> List[str]:
        """Split text into word-level tokens with punctuation attached."""
        return re.findall(r"\S+|\s+", text)

    def _get_pairs(self, word_tokens: List[str]) -> List[Tuple[str, str]]:
        """Get all adjacent pairs in a word's token list."""
        return [(word_tokens[i], word_tokens[i + 1]) for i in range(len(word_tokens) - 1)]

    def train(self, texts: List[str] | str, verbose: bool = False) -> None:
        """Train BPE tokenizer on raw text.

        Args:
            texts: list of text strings or a single string
            verbose: print progress
        """
        if isinstance(texts, str):
            texts = [texts]

        # Build word frequency counter
        word_freq: Counter = Counter()
        for text in texts:
            words = self._tokenize_words(text)
            word_freq.update(words)

        # Initialize each word as a sequence of characters
        word_splits: Dict[str, List[str]] = {}
        for word in word_freq:
            word_splits[word] = list(word)

        # Build initial vocab from characters
        char_set: set[str] = set()
        for word in word_splits:
            char_set.update(word_splits[word])
        for ch in sorted(char_set):
            if ch not in self.vocab:
                idx = len(self.vocab)
                self.vocab[ch] = idx
                self.inv_vocab[idx] = ch

        # BPE merge loop
        num_merges = self.target_vocab_size - len(self.vocab)
        if num_merges <= 0:
            logger.info("Vocab size already at target (%d), no merges needed", len(self.vocab))
            return

        for merge_idx in range(num_merges):
            # Count pair frequencies
            pair_freq: Counter = Counter()
            for word, freq in word_freq.items():
                tokens = word_splits[word]
                for pair in self._get_pairs(tokens):
                    pair_freq[pair] += freq

            if not pair_freq:
                break

            # Find most frequent pair
            best_pair = pair_freq.most_common(1)[0][0]
            best_freq = pair_freq[best_pair]

            if best_freq < 2:
                break

            # Create merged token
            merged = best_pair[0] + best_pair[1]
            if merged in self.vocab:
                break

            idx = len(self.vocab)
            self.vocab[merged] = idx
            self.inv_vocab[idx] = merged
            self.merges.append(best_pair)

            # Apply merge to all words
            for word in word_splits:
                tokens = word_splits[word]
                new_tokens: List[str] = []
                i = 0
                while i < len(tokens):
                    if i < len(tokens) - 1 and tokens[i] == best_pair[0] and tokens[i + 1] == best_pair[1]:
                        new_tokens.append(merged)
                        i += 2
                    else:
                        new_tokens.append(tokens[i])
                        i += 1
                word_splits[word] = new_tokens

            if verbose and (merge_idx + 1) % 100 == 0:
                logger.info("BPE merge %d/%d: '%s' (freq=%d), vocab=%d",
                            merge_idx + 1, num_merges, merged, best_freq, len(self.vocab))

        # Build word cache for fast encoding
        self._rebuild_cache(word_splits)
        logger.info("BPE training complete: %d tokens, %d merges", len(self.vocab), len(self.merges))

    def _rebuild_cache(self, word_splits: Dict[str, List[str]]) -> None:
        """Build encoding cache from trained word splits."""
        self._word_cache.clear()
        for word, tokens in word_splits.items():
            ids = [self.vocab.get(t, UNK_ID) for t in tokens]
            self._word_cache[word] = ids

    def encode(self, text: str, add_bos: bool = False, add_eos: bool = False) -> List[int]:
        """Encode text to token IDs."""
        ids: List[int] = []
        if add_bos:
            ids.append(BOS_ID)

        words = self._tokenize_words(text)
        for word in words:
            if word in self._word_cache:
                ids.extend(self._word_cache[word])
            else:
                # Apply merges greedily
                tokens = list(word)
                changed = True
                while changed and len(tokens) > 1:
                    changed = False
                    best_idx = -1
                    best_rank = len(self.merges)
                    for i in range(len(tokens) - 1):
                        pair = (tokens[i], tokens[i + 1])
                        if pair in self._merge_ranks:
                            rank = self._merge_ranks[pair]
                            if rank < best_rank:
                                best_rank = rank
                                best_idx = i
                    if best_idx >= 0:
                        merged = tokens[best_idx] + tokens[best_idx + 1]
                        tokens = tokens[:best_idx] + [merged] + tokens[best_idx + 2:]
                        changed = True
                ids.extend(self.vocab.get(t, UNK_ID) for t in tokens)
                self._word_cache[word] = [self.vocab.get(t, UNK_ID) for t in tokens]

        if add_eos:
            ids.append(EOS_ID)
        return ids

    @property
    def _merge_ranks(self) -> Dict[Tuple[str, str], int]:
        if not hasattr(self, '_merge_ranks_cache'):
            self._merge_ranks_cache = {pair: i for i, pair in enumerate(self.merges)}
        return self._merge_ranks_cache

    def decode(self, ids: List[int]) -> str:
        """Decode token IDs back to text."""
        tokens = []
        for tid in ids:
            if tid in (PAD_ID, BOS_ID, EOS_ID):
                continue
            tokens.append(self.inv_vocab.get(tid, UNK_TOKEN))
        return "".join(tokens)

    def save(self, path: str) -> None:
        """Save tokenizer to disk."""
        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
        data = {
            "vocab_size": len(self.vocab),
            "target_vocab_size": self.target_vocab_size,
            "vocab": self.vocab,
            "merges": [[a, b] for a, b in self.merges],
        }
        with open(path, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False)
        logger.info("Tokenizer saved to %s (%d tokens)", path, len(self.vocab))

    @classmethod
    def load(cls, path: str) -> "BPETokenizer":
        """Load tokenizer from disk."""
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)

        tok = cls(vocab_size=data["target_vocab_size"])
        tok.vocab = {k: int(v) for k, v in data["vocab"].items()}
        tok.inv_vocab = {v: k for k, v in tok.vocab.items()}
        tok.merges = [(a, b) for a, b in data["merges"]]
        tok._merge_ranks_cache = {pair: i for i, pair in enumerate(tok.merges)}
        logger.info("Tokenizer loaded from %s (%d tokens)", path, len(tok.vocab))
        return tok