"""BPE tokenizer training and chat formatting for Delta Ultra Mini."""

from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import Any

from tokenizers import Tokenizer
from tokenizers.decoders import ByteLevel as ByteLevelDecoder
from tokenizers.models import BPE
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.processors import TemplateProcessing
from tokenizers.trainers import BpeTrainer

logging.basicConfig(level=os.getenv("DELTA_LOG_LEVEL", "INFO").upper())
logger = logging.getLogger(__name__)

SPECIAL_TOKENS: list[str] = ["[PAD]", "[UNK]", "[BOS]", "[EOS]", "[SYS]", "[USR]", "[ASS]", "[SEP]"]
DEFAULT_SYSTEM_PROMPT = (
    "Você é Delta, assistente criada pela Flame Corporation. "
    "Responda de forma clara, útil e amigável."
)


def train_tokenizer(corpus_files: list[str] | list[Path], output_path: str | Path) -> None:
    """Train a BPE tokenizer from raw text files.

    Args:
        corpus_files: Paths to corpus files.
        output_path: Destination tokenizer JSON path.
    """

    tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
    tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)
    tokenizer.decoder = ByteLevelDecoder()
    trainer = BpeTrainer(vocab_size=32000, special_tokens=SPECIAL_TOKENS, show_progress=True)
    tokenizer.train([str(path) for path in corpus_files], trainer)
    tokenizer.post_processor = TemplateProcessing(
        single="[BOS] $A [EOS]",
        pair="[BOS] $A [SEP] $B [EOS]",
        special_tokens=[
            ("[BOS]", tokenizer.token_to_id("[BOS]")),
            ("[EOS]", tokenizer.token_to_id("[EOS]")),
            ("[SEP]", tokenizer.token_to_id("[SEP]")),
        ],
    )
    output = Path(output_path)
    output.parent.mkdir(parents=True, exist_ok=True)
    tokenizer.save(str(output))
    logger.info("Tokenizer saved to %s", output)


def load_tokenizer(path: str | Path) -> "DeltaTokenizer":
    """Load a DeltaTokenizer from disk."""

    return DeltaTokenizer(path)


class DeltaTokenizer:
    """Thin wrapper around HuggingFace tokenizers.Tokenizer."""

    def __init__(self, path: str | Path) -> None:
        self.path = Path(path)
        self.tokenizer = Tokenizer.from_file(str(self.path))
        self.pad_token_id = self.tokenizer.token_to_id("[PAD]")
        self.unk_token_id = self.tokenizer.token_to_id("[UNK]")
        self.bos_token_id = self.tokenizer.token_to_id("[BOS]")
        self.eos_token_id = self.tokenizer.token_to_id("[EOS]")
        self.sys_token_id = self.tokenizer.token_to_id("[SYS]")
        self.usr_token_id = self.tokenizer.token_to_id("[USR]")
        self.ass_token_id = self.tokenizer.token_to_id("[ASS]")
        self.sep_token_id = self.tokenizer.token_to_id("[SEP]")
        self.special_tokens = set(SPECIAL_TOKENS)

    @property
    def chat_stop_token_ids(self) -> set[int]:
        """Token ids that should end an assistant completion."""

        return {
            token_id
            for token_id in (
                self.eos_token_id,
                self.sep_token_id,
                self.sys_token_id,
                self.usr_token_id,
                self.ass_token_id,
                self.bos_token_id,
                self.pad_token_id,
            )
            if token_id is not None
        }

    def encode(self, text: str, add_special_tokens: bool = True) -> list[int]:
        """Encode a string into token ids."""

        return self.tokenizer.encode(text, add_special_tokens=add_special_tokens).ids

    def decode(self, ids: list[int], skip_special_tokens: bool = True) -> str:
        """Decode token ids into text."""

        return self.tokenizer.decode(ids, skip_special_tokens=skip_special_tokens)

    def batch_encode(self, texts: list[str], add_special_tokens: bool = True) -> list[list[int]]:
        """Encode a batch of strings."""

        return [encoding.ids for encoding in self.tokenizer.encode_batch(texts, add_special_tokens=add_special_tokens)]

    def format_chat(self, messages: list[dict[str, Any]], persona: str | None = None) -> str:
        """Format a multi-turn conversation for Delta.

        Args:
            messages: Conversation turns with role and content.
            persona: Optional system prompt.

        Returns:
            Prompt text ending with an assistant tag for continuation.
        """

        system = persona or DEFAULT_SYSTEM_PROMPT
        parts = [f"[SYS] {system} [SEP]"]
        for message in messages:
            role = str(message.get("role", "")).lower()
            content = str(message.get("content", "")).strip()
            if role == "user":
                parts.append(f"[USR] {content} [SEP]")
            elif role == "assistant":
                parts.append(f"[ASS] {content} [SEP]")
            elif role == "system":
                parts[0] = f"[SYS] {content} [SEP]"
        if not parts[-1].startswith("[ASS]"):
            parts.append("[ASS]")
        else:
            parts.append("[USR]  [SEP]")
            parts.append("[ASS]")
        return "\n".join(parts)