import os
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Union
from tokenizers import Regex, Tokenizer
from tokenizers.decoders import ByteLevel as ByteLevelDecoder
from tokenizers.models import BPE
from tokenizers.normalizers import NFKC, Sequence as NormalizerSequence
from tokenizers.pre_tokenizers import (
    ByteLevel,
    Digits,
    Sequence as PreTokenizerSequence,
    Split,
)
from tokenizers.trainers import BpeTrainer
from transformers import PreTrainedTokenizerFast

SPECIAL_TOKENS = ["<s>", "<pad>", "</s>", "<unk>", "[EOD]", "<|eod|>"]

EMOJIS = [
    "\U0001F602", "\U0001F62D", "\u2728", "\U0001F680", "\U0001F44D",
    "\U0001F64F", "\U0001F525", "\U0001F60A", "\u2764\ufe0f", "\U0001F914",
    "\U0001F923", "\U0001F60D", "\U0001F480", "\U0001F4AF", "\u26a0\ufe0f",
    "\u2705", "\u274c", "\U0001F4CA", "\U0001F4BB", "\U0001F4F1",
    "\U0001F623", "\U0001F970", "\U0001F605", "\U0001F606", "\U0001F979",
    "\U0001F61A", "\U0001F917", "\U0001F61D", "\U0001F440",
]

EMOTICONS = [
    ":-)", ":)", ":D", ":(", ";)", "XD", "OwO", "UwU", "T_T", "QAQ", "¯\\_(ツ)_/¯",
]

MATH_LATEX = [
    "\\alpha", "\\beta", "\\gamma", "\\theta", "\\pi", "\\sigma", "\\omega",
    "\\sum", "\\int", "\\approx", "\\neq", "\\le", "\\ge", "\\infty",
    "\\partial", "\\nabla", "\\forall", "\\exists", "\\in", "\\notin",
    "\\rightarrow", "\\Rightarrow", "\\Leftrightarrow",
]

CODE_OPERATORS = [
    "==", "!=", "<=", ">=", "+=", "-=", "*=", "/=",
    "=>", "->", "&&", "||", "async", "await", "lambda",
]

CUSTOM_TOKENS = [
    "<|im_start|>",
    "<|im_end|>",
    "<|system|>",
    "<|user|>",
    "<|assistant|>",
    "<think>",
    "</think>",
] + (EMOJIS + EMOTICONS + MATH_LATEX + CODE_OPERATORS)

ALL_TOKENS = SPECIAL_TOKENS + CUSTOM_TOKENS

LLM_SPLIT_REGEX = (
    r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}|"""
    r""" ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
)

def build_xonelm_tokenizer(
    corpus: Optional[Union[str, List[str]]] = None,
    vocab_size: int = 32000,
    save_path: Optional[str] = None,
) -> PreTrainedTokenizerFast:
    bpe_model = BPE(unk_token="<unk>")
    tokenizer_raw = Tokenizer(bpe_model)

    tokenizer_raw.normalizer = NormalizerSequence([NFKC()])

    tokenizer_raw.pre_tokenizer = PreTokenizerSequence([
        Split(pattern=Regex(LLM_SPLIT_REGEX), behavior="isolated", invert=False),
        Digits(individual_digits=True),
        ByteLevel(add_prefix_space=False, use_regex=False),
    ])

    tokenizer_raw.decoder = ByteLevelDecoder()

    trainer = BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=ALL_TOKENS,
        initial_alphabet=ByteLevel.alphabet(),
        show_progress=False,
    )

    if corpus is not None:
        if isinstance(corpus, str) and os.path.isfile(corpus):
            tokenizer_raw.train([corpus], trainer)
        elif isinstance(corpus, list) and len(corpus) > 0 and os.path.isfile(corpus[0]):
            tokenizer_raw.train(corpus, trainer)
        else:
            iterator = [corpus] if isinstance(corpus, str) else corpus
            tokenizer_raw.train_from_iterator(iterator, trainer)
    else:
        tokenizer_raw.train_from_iterator(["Hello world 123 \\alpha \\beta == async await"], trainer)

    if save_path is not None:
        tokenizer_raw.save(save_path)

    hf_tokenizer = PreTrainedTokenizerFast(
        tokenizer_object=tokenizer_raw,
        bos_token="<s>",
        eos_token="</s>",
        pad_token="<pad>",
        unk_token="<unk>",
        additional_special_tokens=ALL_TOKENS,
    )
    return hf_tokenizer

@dataclass
class SpecialTokenConfig:
    pad_token_id: int = 0
    bos_token_id: int = 1
    eos_token_id: int = 2
    unk_token_id: int = 3
    eod_token_id: int = 4
    im_start_id: Optional[int] = None
    im_end_id: Optional[int] = None
    separator_token_id: Optional[int] = None

class MultiTurnConversationFormatter:
    def __init__(self, tokenizer: Any, token_config: Optional[SpecialTokenConfig] = None):
        self.tokenizer = tokenizer
        self.config = token_config or SpecialTokenConfig()

        def _get_id(token_str: str) -> Optional[int]:
            if hasattr(tokenizer, "token_to_id"):
                return tokenizer.token_to_id(token_str)
            elif hasattr(tokenizer, "convert_tokens_to_ids"):
                res = tokenizer.convert_tokens_to_ids(token_str)
                return res if isinstance(res, int) and res >= 0 else None
            return None

        if self.config.im_start_id is None:
            self.config.im_start_id = _get_id("<|im_start|>")
        if self.config.im_end_id is None:
            self.config.im_end_id = _get_id("<|im_end|>")
        if self.config.eod_token_id is None:
            self.config.eod_token_id = _get_id("[EOD]")

    def format_conversation(
        self, messages: List[Dict[str, str]], max_len: Optional[int] = None
    ) -> Dict[str, List[int]]:
        input_ids = []
        labels = []

        def _encode_text(t: str) -> List[int]:
            if hasattr(self.tokenizer, "encode"):
                res = self.tokenizer.encode(t)
                return res.ids if hasattr(res, "ids") else res
            elif callable(self.tokenizer):
                return self.tokenizer(t)["input_ids"]
            return []

        for msg in messages:
            role = msg["role"]
            content = msg["content"].strip()

            header_text = f"<|im_start|>{role}\n"
            body_text = f"{content}<|im_end|>\n"

            header_ids = _encode_text(header_text)
            body_ids = _encode_text(body_text)

            turn_input_ids = header_ids + body_ids
            input_ids.extend(turn_input_ids)

            if role == "assistant":
                turn_labels = [-100] * len(header_ids) + body_ids
                labels.extend(turn_labels)
            else:
                labels.extend([-100] * len(turn_input_ids))

        if self.config.eod_token_id is not None:
            input_ids.append(self.config.eod_token_id)
            labels.append(self.config.eod_token_id)

        if max_len is not None:
            input_ids = input_ids[:max_len]
            labels = labels[:max_len]

        return {"input_ids": input_ids, "labels": labels}