"""
BPE Tokenizer — built from scratch.

Implements byte-pair encoding following the original algorithm:
1. Start with character-level vocabulary
2. Iteratively merge most frequent adjacent pairs
3. Build vocab of target size

Special tokens (15): <pad>, <bos>, <eos>, <query>, <vqa>, <detect>, <alert>, <caption>, <track>,
                     <count>, <ocr>, <reason>, <temporal>, <multi_cam>, <system>

Byte fallback: 256 byte tokens <0x00>-<0xFF> for handling unknown characters.
"""

import re
import json
from collections import Counter
from pathlib import Path


class BPETokenizer:
    """
    Byte-Pair Encoding tokenizer built from scratch.

    Usage:
        tokenizer = BPETokenizer(vocab_size=32768)
        tokenizer.train(texts)  # list of strings
        tokens = tokenizer.encode("What color is the car?")
        text = tokenizer.decode(tokens)
    """

    SPECIAL_TOKENS = [
        "<pad>", "<bos>", "<eos>", "<query>",
        "<vqa>", "<detect>", "<alert>", "<caption>", "<track>",
        "<count>", "<ocr>", "<reason>", "<temporal>", "<multi_cam>", "<system>"
    ]

    # 256 byte fallback tokens: indices 15-270
    BYTE_TOKENS = [f"<0x{i:02X}>" for i in range(256)]

    # Regex pattern to detect byte tokens during decoding
    _BYTE_TOKEN_RE = re.compile(r"<0x([0-9A-F]{2})>")

    def __init__(self, vocab_size: int = 32768):
        self.vocab_size = vocab_size
        self.merges: list[tuple[str, str]] = []  # ordered merge rules
        self.vocab: dict[str, int] = {}  # token -> id
        self.inverse_vocab: dict[int, str] = {}  # id -> token

        # Initialize special tokens (indices 0-14)
        for i, tok in enumerate(self.SPECIAL_TOKENS):
            self.vocab[tok] = i
            self.inverse_vocab[i] = tok

        # Initialize byte fallback tokens (indices 15-270)
        byte_start = len(self.SPECIAL_TOKENS)
        for i, tok in enumerate(self.BYTE_TOKENS):
            idx = byte_start + i
            self.vocab[tok] = idx
            self.inverse_vocab[idx] = tok

        self.pad_id = self.vocab["<pad>"]
        self.bos_id = self.vocab["<bos>"]
        self.eos_id = self.vocab["<eos>"]

        # BPE merges start after special + byte tokens
        self._reserved_end = byte_start + len(self.BYTE_TOKENS)  # 271

    def _get_word_freqs(self, texts: list[str]) -> dict[tuple[str, ...], int]:
        """Split texts into words and count frequencies. Each word is a tuple of characters + end marker."""
        word_freqs: dict[tuple[str, ...], int] = Counter()
        for text in texts:
            words = text.strip().split()
            for word in words:
                # Represent word as tuple of characters + end-of-word marker
                char_tuple = tuple(word) + ("</w>",)
                word_freqs[char_tuple] += 1
        return word_freqs

    def _get_pair_freqs(self, word_freqs: dict[tuple[str, ...], int]) -> Counter:
        """Count frequency of adjacent pairs across all words."""
        pairs = Counter()
        for word, freq in word_freqs.items():
            for i in range(len(word) - 1):
                pairs[(word[i], word[i + 1])] += freq
        return pairs

    def _merge_pair(
        self, pair: tuple[str, str], word_freqs: dict[tuple[str, ...], int]
    ) -> dict[tuple[str, ...], int]:
        """Merge all occurrences of pair in word_freqs."""
        new_word_freqs: dict[tuple[str, ...], int] = {}
        merged = pair[0] + pair[1]

        for word, freq in word_freqs.items():
            new_word: list[str] = []
            i = 0
            while i < len(word):
                if i < len(word) - 1 and word[i] == pair[0] and word[i + 1] == pair[1]:
                    new_word.append(merged)
                    i += 2
                else:
                    new_word.append(word[i])
                    i += 1
            new_word_freqs[tuple(new_word)] = freq

        return new_word_freqs

    def train(self, texts: list[str]) -> None:
        """
        Train BPE on a corpus of texts.

        Args:
            texts: List of text strings to learn merges from
        """
        word_freqs = self._get_word_freqs(texts)

        # Build initial character vocabulary
        chars: set[str] = set()
        for word in word_freqs:
            for char in word:
                chars.add(char)

        # Start vocab after reserved tokens (special + byte = 271)
        idx = self._reserved_end
        for char in sorted(chars):
            if char not in self.vocab:
                self.vocab[char] = idx
                self.inverse_vocab[idx] = char
                idx += 1

        # Iteratively merge most frequent pairs
        num_merges = self.vocab_size - idx
        for _ in range(num_merges):
            pair_freqs = self._get_pair_freqs(word_freqs)
            if not pair_freqs:
                break

            best_pair = pair_freqs.most_common(1)[0][0]
            self.merges.append(best_pair)

            # Add merged token to vocab
            merged_token = best_pair[0] + best_pair[1]
            if merged_token not in self.vocab:
                self.vocab[merged_token] = idx
                self.inverse_vocab[idx] = merged_token
                idx += 1

            # Apply merge to all words
            word_freqs = self._merge_pair(best_pair, word_freqs)

    def _encode_as_bytes(self, token: str) -> list[int]:
        """Encode a token as byte-level fallback tokens."""
        ids = []
        for byte in token.encode("utf-8"):
            byte_token = f"<0x{byte:02X}>"
            ids.append(self.vocab[byte_token])
        return ids

    def _tokenize_word(self, word: str) -> list[str]:
        """Apply learned merges to a single word."""
        tokens = list(word) + ["</w>"]

        for pair in self.merges:
            i = 0
            while i < len(tokens) - 1:
                if tokens[i] == pair[0] and tokens[i + 1] == pair[1]:
                    tokens = tokens[:i] + [pair[0] + pair[1]] + tokens[i + 2:]
                else:
                    i += 1

        return tokens

    def encode(self, text: str, add_special: bool = True) -> list[int]:
        """
        Encode text to token IDs.

        Args:
            text: Input string
            add_special: Whether to wrap with <bos>/<eos>
        Returns:
            List of token IDs
        """
        ids: list[int] = []
        if add_special:
            ids.append(self.bos_id)

        words = text.strip().split()
        for word in words:
            tokens = self._tokenize_word(word)
            for token in tokens:
                if token in self.vocab:
                    ids.append(self.vocab[token])
                else:
                    # Byte fallback: encode each byte of the UTF-8 representation
                    ids.extend(self._encode_as_bytes(token))

        if add_special:
            ids.append(self.eos_id)

        return ids

    def decode(self, ids: list[int]) -> str:
        """
        Decode token IDs back to text.

        Args:
            ids: List of token IDs
        Returns:
            Decoded string
        """
        tokens = []
        for id_ in ids:
            if id_ in self.inverse_vocab:
                token = self.inverse_vocab[id_]
                if token in self.SPECIAL_TOKENS:
                    continue
                tokens.append(token)

        # Join tokens, then decode byte fallback sequences
        text = "".join(tokens)

        # Decode consecutive byte tokens back to UTF-8
        def _replace_byte_sequences(s: str) -> str:
            result = []
            i = 0
            byte_buffer = []
            while i < len(s):
                m = self._BYTE_TOKEN_RE.match(s, i)
                if m:
                    byte_buffer.append(int(m.group(1), 16))
                    i = m.end()
                else:
                    if byte_buffer:
                        try:
                            result.append(bytes(byte_buffer).decode("utf-8", errors="replace"))
                        except Exception:
                            result.append(bytes(byte_buffer).decode("utf-8", errors="replace"))
                        byte_buffer = []
                    result.append(s[i])
                    i += 1
            if byte_buffer:
                try:
                    result.append(bytes(byte_buffer).decode("utf-8", errors="replace"))
                except Exception:
                    result.append(bytes(byte_buffer).decode("utf-8", errors="replace"))
            return "".join(result)

        text = _replace_byte_sequences(text)

        # Clean up end-of-word markers
        text = text.replace("</w>", " ").strip()
        return text

    def save(self, path: str) -> None:
        """Save tokenizer to JSON file."""
        data = {
            "vocab_size": self.vocab_size,
            "merges": self.merges,
            "vocab": self.vocab,
            "special_tokens": self.SPECIAL_TOKENS,
            "byte_tokens": self.BYTE_TOKENS,
        }
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            json.dump(data, f, indent=2)

    def load(self, path: str) -> None:
        """Load tokenizer from JSON file."""
        with open(path) as f:
            data = json.load(f)

        self.vocab_size = data["vocab_size"]
        self.merges = [tuple(m) for m in data["merges"]]
        self.vocab = data["vocab"]
        self.inverse_vocab = {v: k for k, v in self.vocab.items()}

    def __len__(self) -> int:
        return len(self.vocab)