"""Tokenizer for gLM2.

Wraps a BPE-style fast tokenizer with the upstream `tattabio/gLM2_*` vocabulary
(amino acids, nucleotides, strand markers). On top of the upstream behaviour
this version adds *automatic DNA preparation* so callers can pass plain DNA
sequences (e.g. mRNABench-style input) without manually lowercasing or adding
strand markers.
"""

from __future__ import annotations

import re
from typing import List, Optional, Sequence, Union

from tokenizers import Tokenizer
from tokenizers.models import BPE
from transformers import PreTrainedTokenizerFast


_NUC_CHARS = set("acgtuACGTU")
_DNA_PREFIX_RE = re.compile(r"^\s*(<\+>|<->)")


def _is_pure_dna(seq: str) -> bool:
    """True if `seq` only contains DNA/RNA characters (ATCGU, any case)."""
    return len(seq) > 0 and all(c in _NUC_CHARS for c in seq)


def prepare_dna_sequence(seq: str, strand: str = "+") -> str:
    """Normalize a plain DNA/RNA sequence for gLM2.

    - Lower-cases nucleotides so they are tokenized as DNA (gLM2 uses
      lower-case `a/t/c/g` for nucleotides; upper-case letters are amino
      acids).
    - Replaces `U`/`u` with `t` (gLM2's vocab has no uracil token).
    - Prepends the strand marker (`<+>` or `<->`) if not already present.
    """
    if strand not in ("+", "-"):
        raise ValueError(f"strand must be '+' or '-', got {strand!r}")
    if _DNA_PREFIX_RE.match(seq):
        return seq
    cleaned = seq.lower().replace("u", "t")
    return f"<{strand}>{cleaned}"


class gLM2Tokenizer(PreTrainedTokenizerFast):
    """gLM2 fast tokenizer with optional automatic DNA preparation."""

    VOCAB = [
        "<cls>", "<pad>", "<eos>", "<unk>",
        "L", "A", "G", "V", "S", "E", "R", "T", "I", "D", "P", "K",
        "Q", "N", "F", "Y", "M", "H", "W", "C", "X", "B", "U", "Z",
        "O", "a", "t", "c", "g", "<+>", "<->", "<mask>", "<sep>",
    ]

    def __init__(
        self,
        unk_token: str = "<unk>",
        cls_token: str = "<cls>",
        pad_token: str = "<pad>",
        mask_token: str = "<mask>",
        eos_token: str = "<eos>",
        sep_token: str = "<sep>",
        pos_token: str = "<+>",
        neg_token: str = "<->",
        auto_prepare_dna: bool = True,
        model_max_length: int = 4096,
        **kwargs,
    ):
        all_tokens = self.VOCAB
        token_to_id = {tok: ind for ind, tok in enumerate(all_tokens)}

        bpe = BPE(token_to_id, merges=[], unk_token=str(unk_token))
        tokenizer = Tokenizer(bpe)
        special_tokens = [
            cls_token, pad_token, mask_token, eos_token, sep_token, pos_token, neg_token,
        ]
        tokenizer.add_special_tokens(special_tokens)

        super().__init__(
            tokenizer_object=tokenizer,
            unk_token=unk_token,
            cls_token=cls_token,
            pad_token=pad_token,
            mask_token=mask_token,
            eos_token=eos_token,
            sep_token=sep_token,
            model_max_length=model_max_length,
            **kwargs,
        )
        self.auto_prepare_dna = auto_prepare_dna

    def _maybe_prepare(
        self, text: Union[str, Sequence[str]]
    ) -> Union[str, List[str]]:
        if not self.auto_prepare_dna:
            return text  # type: ignore[return-value]
        if isinstance(text, str):
            return prepare_dna_sequence(text) if _is_pure_dna(text) else text
        out: List[str] = []
        for s in text:
            if isinstance(s, str) and _is_pure_dna(s):
                out.append(prepare_dna_sequence(s))
            else:
                out.append(s)
        return out

    def __call__(
        self,
        text: Union[str, Sequence[str], None] = None,
        text_pair=None,
        text_target=None,
        text_pair_target=None,
        *args,
        **kwargs,
    ):
        if text is not None:
            text = self._maybe_prepare(text)
        return super().__call__(
            text,
            text_pair=text_pair,
            text_target=text_target,
            text_pair_target=text_pair_target,
            *args,
            **kwargs,
        )

    def encode(
        self, text: Union[str, Sequence[str]], *args, **kwargs
    ) -> List[int]:
        text = self._maybe_prepare(text)
        return super().encode(text, *args, **kwargs)