from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Iterable

BASE_VOCAB_SIZE = 256
ATOM_RE = re.compile(
    r"<!--[\s\S]*?-->|<!doctype\s+html>|</?[A-Za-z][^<>]*>|"
    r"\b(?:display|place-items|align-items|justify-content|min-height|margin|font|padding|border|border-radius|background)\s*:\s*[^;{}]+;?|"
    r"\s+|[A-Za-z_$][A-Za-z0-9_$-]*|[0-9]+(?:\.[0-9]+)?(?:%|px|rem|em|vh|vw)?|[^\x00-\x7F]|[^\w\s]",
    re.IGNORECASE | re.ASCII,
)


class ConstrainedBPE:
    """Byte-fallback BPE matching DivCenterer's syntax-aware Node tokenizer."""

    def __init__(self, model: dict):
        if model.get("format") != "divcenterer-constrained-byte-bpe-v1":
            raise ValueError("Unrecognized tokenizer model format")
        self.special_tokens: list[str] = model["special_tokens"]
        self.special_to_id = {
            token: BASE_VOCAB_SIZE + index
            for index, token in enumerate(self.special_tokens)
        }
        self.special_by_id = {value: key for key, value in self.special_to_id.items()}
        self.merge_ranks: dict[tuple[int, int], int] = {}
        self.token_bytes: list[bytes | None] = [bytes([index]) for index in range(BASE_VOCAB_SIZE)]
        self.token_bytes.extend([None] * len(self.special_tokens))
        for rank, (left, right) in enumerate(model["merges"]):
            if self.token_bytes[left] is None or self.token_bytes[right] is None:
                raise ValueError(f"Merge {rank} refers to a special or missing token")
            token_id = BASE_VOCAB_SIZE + len(self.special_tokens) + rank
            self.merge_ranks[(left, right)] = rank
            self.token_bytes.append(self.token_bytes[left] + self.token_bytes[right])
        self._special_re = re.compile("(" + "|".join(re.escape(t) for t in self.special_tokens) + ")")
        self._atom_cache: dict[str, list[int]] = {}

    @classmethod
    def from_file(cls, filename: str | Path) -> "ConstrainedBPE":
        with Path(filename).open("r", encoding="utf-8") as stream:
            return cls(json.load(stream))

    def _atoms(self, text: str) -> list[str]:
        atoms: list[str] = []
        cursor = 0
        for match in ATOM_RE.finditer(text):
            if match.start() != cursor:
                raise ValueError(f"Pretokenizer left a gap at character {cursor}")
            atoms.append(match.group(0))
            cursor = match.end()
        if cursor != len(text):
            raise ValueError(f"Pretokenizer left a gap at character {cursor}")
        return atoms

    @staticmethod
    def _apply_merge(ids: list[int], left: int, right: int, merged: int) -> list[int]:
        result: list[int] = []
        index = 0
        while index < len(ids):
            if index + 1 < len(ids) and ids[index] == left and ids[index + 1] == right:
                result.append(merged)
                index += 2
            else:
                result.append(ids[index])
                index += 1
        return result

    def _encode_atom(self, atom: str) -> list[int]:
        cached = self._atom_cache.get(atom)
        if cached is not None:
            return cached.copy()
        ids = list(atom.encode("utf-8"))
        while len(ids) > 1:
            best_rank: int | None = None
            best_pair: tuple[int, int] | None = None
            for left, right in zip(ids, ids[1:]):
                rank = self.merge_ranks.get((left, right))
                if rank is not None and (best_rank is None or rank < best_rank):
                    best_rank = rank
                    best_pair = (left, right)
            if best_pair is None or best_rank is None:
                break
            merged_id = BASE_VOCAB_SIZE + len(self.special_tokens) + best_rank
            ids = self._apply_merge(ids, best_pair[0], best_pair[1], merged_id)
        self._atom_cache[atom] = ids.copy()
        return ids

    def encode(self, text: str) -> list[int]:
        result: list[int] = []
        for part in self._special_re.split(text):
            if not part:
                continue
            special_id = self.special_to_id.get(part)
            if special_id is not None:
                result.append(special_id)
            else:
                for atom in self._atoms(part):
                    result.extend(self._encode_atom(atom))
        return result

    def decode(self, ids: Iterable[int]) -> str:
        output: list[str] = []
        byte_buffer = bytearray()

        def flush_bytes() -> None:
            if byte_buffer:
                output.append(bytes(byte_buffer).decode("utf-8", errors="replace"))
                byte_buffer.clear()

        for token_id in ids:
            special = self.special_by_id.get(int(token_id))
            if special is not None:
                flush_bytes()
                output.append(special)
            elif 0 <= int(token_id) < len(self.token_bytes) and self.token_bytes[int(token_id)] is not None:
                byte_buffer.extend(self.token_bytes[int(token_id)])
            else:
                raise ValueError(f"Unknown token id: {token_id}")
        flush_bytes()
        return "".join(output)

    def special_id(self, token: str) -> int:
        return self.special_to_id[token]

    @property
    def vocab_size(self) -> int:
        return len(self.token_bytes)


def conversation_prefix(prompt: str) -> str:
    return f"<|user|>\n{prompt}\n<|assistant|>\n"


def conversation_text(prompt: str, answer: str) -> str:
    return f"{conversation_prefix(prompt)}{answer}\n<|end|>"