"""Text tokenizers behind one interface (numpy only, no torch):

    tok.encode_batch(texts, max_len, dynamic=False) -> (n, L) int64, padded with tok.pad_id
    tok.encode_choice_sets(choice_sets, max_len)    -> (B, N, L); absent choices are all-pad rows

`ByteTokenizer` is the original MiniJEV byte-level tokenizer (vocab 260, pad 0).
`HFTokenizer` wraps a pretrained encoder's `tokenizer.json` through the Rust
`tokenizers` package (no transformers/torch needed at inference).
"""
from __future__ import annotations

from pathlib import Path
from typing import Sequence

import numpy as np

from minijev import tokenizer as byte_tok


class ByteTokenizer:
    kind = "byte"
    pad_id = byte_tok.PAD
    vocab_size = byte_tok.VOCAB_SIZE

    def encode_batch(self, texts: Sequence[str], max_len: int, dynamic: bool = False) -> np.ndarray:
        out = byte_tok.encode_batch(texts, max_len)
        if dynamic:
            n = int((out != self.pad_id).sum(1).max()) if len(out) else 1
            out = out[:, : max(1, n)]
        return out

    def encode_choice_sets(self, choice_sets, max_len):
        return byte_tok.encode_choice_sets(choice_sets, max_len)


class HFTokenizer:
    kind = "hf"

    def __init__(self, name_or_path: str):
        from tokenizers import Tokenizer

        p = Path(name_or_path)
        if p.is_dir():
            p = p / "tokenizer.json"
        if not p.exists():
            from huggingface_hub import hf_hub_download
            p = Path(hf_hub_download(name_or_path, "tokenizer.json"))
        self.path = str(p)
        self.tok = Tokenizer.from_file(str(p))
        self.tok.no_padding()
        self.tok.no_truncation()
        pad = self.tok.token_to_id("<pad>")
        self.pad_id = int(pad if pad is not None else (self.tok.token_to_id("[PAD]") or 0))
        self.vocab_size = self.tok.get_vocab_size()

    def encode_batch(self, texts: Sequence[str], max_len: int, dynamic: bool = False) -> np.ndarray:
        encs = self.tok.encode_batch(list(texts))
        L = max_len
        if dynamic:
            L = max(1, min(max_len, max((len(e.ids) for e in encs), default=1)))
        out = np.full((len(texts), L), self.pad_id, np.int64)
        for i, e in enumerate(encs):
            ids = e.ids
            if len(ids) > L:  # keep the final special token (</s> / [SEP])
                ids = ids[: L - 1] + ids[-1:]
            out[i, : len(ids)] = ids
        return out

    def encode_choice_sets(self, choice_sets, max_len):
        n = max(len(c) for c in choice_sets)
        flat = [c for cs in choice_sets for c in cs]
        enc = self.encode_batch(flat, max_len, dynamic=True)
        out = np.full((len(choice_sets), n, enc.shape[1]), self.pad_id, np.int64)
        k = 0
        for i, cs in enumerate(choice_sets):
            out[i, : len(cs)] = enc[k:k + len(cs)]
            k += len(cs)
        return out


def make_tokenizer(text_cfg: dict | None):
    text_cfg = text_cfg or {}
    if text_cfg.get("type", "byte") == "byte":
        return ByteTokenizer()
    return HFTokenizer(text_cfg.get("tokenizer") or text_cfg["name"])