from __future__ import annotations

from typing import Iterable
from transformers import AutoTokenizer


class HybridByteTokenizer:
    def __init__(self, model_id_or_path: str, **kwargs):
        self.tokenizer = AutoTokenizer.from_pretrained(model_id_or_path, **kwargs)
        vocab = self.tokenizer.get_vocab()

        self.byte_token_ids = []
        for value in range(256):
            token = f"<0x{value:02X}>"
            if token not in vocab:
                raise ValueError(f"Missing byte token: {token}")
            self.byte_token_ids.append(int(vocab[token]))

        self.byte_id_to_value = {
            token_id: value
            for value, token_id in enumerate(self.byte_token_ids)
        }
        self.separator_ids = self.tokenizer.encode(
            "\\n\\n", add_special_tokens=False
        )
        self.quote_ids = self.tokenizer.encode(
            '"', add_special_tokens=False
        )

    def encode_bytes(self, text: str) -> list[int]:
        return [
            self.byte_token_ids[value]
            for value in str(text or "").encode("utf-8")
        ]

    def encode_rewrite(
        self,
        instruction: str,
        text: str,
        output: str = "",
    ) -> list[int]:
        ids = []
        instruction = str(instruction or "").strip()

        if instruction:
            ids.extend(
                self.tokenizer.encode(
                    instruction,
                    add_special_tokens=False,
                )
            )
            ids.extend(self.separator_ids)

        ids.extend(self.quote_ids)
        ids.extend(self.encode_bytes(text))
        ids.extend(self.quote_ids)
        ids.extend(self.separator_ids)
        ids.extend(self.quote_ids)
        ids.extend(self.encode_bytes(output))
        ids.extend(self.quote_ids)
        return ids

    def encode_text(
        self,
        text: str,
        byte_inside_quotes: bool = True,
    ) -> list[int]:
        if not byte_inside_quotes:
            return self.tokenizer.encode(text, add_special_tokens=False)

        ids = []
        normal = []
        quoted = []
        inside = False
        escaped = False

        def flush_normal():
            if normal:
                ids.extend(
                    self.tokenizer.encode(
                        "".join(normal),
                        add_special_tokens=False,
                    )
                )
                normal.clear()

        def flush_quoted():
            if quoted:
                ids.extend(self.encode_bytes("".join(quoted)))
                quoted.clear()

        for character in str(text):
            if inside:
                if character == '"' and not escaped:
                    flush_quoted()
                    ids.extend(self.quote_ids)
                    inside = False
                    escaped = False
                    continue
                quoted.append(character)
                escaped = character == "\\\\" and not escaped
                if character != "\\\\":
                    escaped = False
            else:
                if character == '"':
                    flush_normal()
                    ids.extend(self.quote_ids)
                    inside = True
                else:
                    normal.append(character)

        if inside:
            flush_quoted()
        else:
            flush_normal()
        return ids

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

        def flush_normal():
            if normal:
                output.append(
                    self.tokenizer.decode(
                        normal,
                        skip_special_tokens=False,
                        clean_up_tokenization_spaces=False,
                    )
                )
                normal.clear()

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

        for token_id_raw in token_ids:
            token_id = int(token_id_raw)
            byte_value = self.byte_id_to_value.get(token_id)

            if byte_value is None:
                flush_bytes()
                normal.append(token_id)
            else:
                flush_normal()
                byte_buffer.append(byte_value)

        flush_bytes()
        flush_normal()
        return "".join(output)