#!/usr/bin/env python3
"""HuggingFace slow-tokenizer wrapper for our morphological tokenizer, for the BabyLM leaderboard submission.

The official strict harness (`sentence_zero_shot/dataset.py`) calls `tokenizer(text, return_offsets_mapping=True)`
and locates the scored completion by CHARACTER offsets. A pure-Python tokenizer normally can't return offsets;
this subclass overrides the encode path to emit word-granularity `offset_mapping` (via OurTokenizer.encode_with_spans),
which is sufficient because completions always begin at a whitespace boundary.

Ship this file + our_vocab.json (+ the analyzer resources) with the model and load via `trust_remote_code=True`.
Special tokens are NOT auto-added (add_special_tokens is a no-op) to match the harness's no-BOS default for
from-scratch tokenizers.
"""
from __future__ import annotations

import json
from pathlib import Path

from transformers import PreTrainedTokenizer
from transformers.tokenization_utils_base import BatchEncoding

try:
    # NOTE: HF's local-dir dynamic-module loader copies only THIS file's one level of relative
    # imports (no recursion), so the analyzer must be imported here even though only
    # our_tokenizer uses it.
    from .morph_tokenizer import load_resources as _hf_copy_probe  # noqa: F401
    from .our_tokenizer import BOS, EOS, OurTokenizer, PAD, SPECIALS, UNK  # noqa: F401
except ImportError:  # direct use (importlib evades HF check_imports regex)
    import importlib as _il
    _m = _il.import_module("our_tokenizer")
    BOS, EOS, OurTokenizer, PAD, SPECIALS, UNK = _m.BOS, _m.EOS, _m.OurTokenizer, _m.PAD, _m.SPECIALS, _m.UNK


class OurMorphTokenizer(PreTrainedTokenizer):
    """Slow HF tokenizer backed by OurTokenizer. Emits offset_mapping for the strict harness."""

    vocab_files_names = {"vocab_file": "our_vocab.json"}
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(self, vocab_file=None, char_backoff=True, **kwargs):
        self._vocab = json.loads(Path(vocab_file).read_text(encoding="utf-8")) if vocab_file else {s: i for i, s in enumerate(SPECIALS)}
        self._ids_to_tokens = {i: t for t, i in self._vocab.items()}
        self._ot = OurTokenizer.load(vocab_file, res=None) if vocab_file else None
        if self._ot is not None:
            self._ot.char_backoff = char_backoff
        super().__init__(pad_token=PAD, unk_token=UNK, bos_token=BOS, eos_token=EOS,
                         char_backoff=char_backoff, **kwargs)

    @property
    def vocab_size(self) -> int:
        return len(self._vocab)

    def get_vocab(self) -> dict:
        return dict(self._vocab)

    def _tokenize(self, text, **kwargs):
        return self._ot.tokenize(text)

    def _convert_token_to_id(self, token):
        return self._vocab.get(token, self._vocab[UNK])

    def _convert_id_to_token(self, index):
        return self._ids_to_tokens.get(index, UNK)

    def convert_tokens_to_string(self, tokens):
        return self._ot.decode([self._convert_token_to_id(t) for t in tokens])

    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
        # No auto special tokens (match the harness's no-BOS default); pairs are concatenated.
        return token_ids_0 if token_ids_1 is None else token_ids_0 + token_ids_1

    def save_vocabulary(self, save_directory, filename_prefix=None):
        out = Path(save_directory) / ((filename_prefix + "-" if filename_prefix else "") + "our_vocab.json")
        out.write_text(json.dumps(self._vocab, ensure_ascii=False), encoding="utf-8")
        return (str(out),)

    # -- offset-aware encoding (the reason this wrapper exists) -------------------------------------
    def _one(self, text, return_offsets_mapping):
        ids, spans = self._ot.encode_with_spans(text)
        d = {"input_ids": ids, "attention_mask": [1] * len(ids)}
        if return_offsets_mapping:
            d["offset_mapping"] = spans
        return d

    def _encode_plus(self, text, text_pair=None, add_special_tokens=True, return_offsets_mapping=False,
                     return_attention_mask=None, return_tensors=None, **kwargs):
        d = self._one(text, return_offsets_mapping)
        if text_pair is not None:
            d2 = self._one(text_pair, return_offsets_mapping)
            d["input_ids"] += d2["input_ids"]
            d["attention_mask"] += d2["attention_mask"]
            if return_offsets_mapping:
                d["offset_mapping"] += d2["offset_mapping"]
        return BatchEncoding(d, tensor_type=return_tensors, prepend_batch_axis=True)

    def _batch_encode_plus(self, batch_text_or_pairs, add_special_tokens=True, return_offsets_mapping=False,
                           return_attention_mask=None, return_tensors=None, padding_strategy=None, **kwargs):
        rows = []
        for t in batch_text_or_pairs:
            text, pair = t if isinstance(t, tuple) else (t, None)
            r = self._one(text, return_offsets_mapping)
            if pair is not None:
                r2 = self._one(pair, return_offsets_mapping)
                r["input_ids"] += r2["input_ids"]; r["attention_mask"] += r2["attention_mask"]
                if return_offsets_mapping:
                    r["offset_mapping"] += r2["offset_mapping"]
            rows.append(r)
        # left-as-lists (harness scores one at a time); pad only if a tensor type is requested
        out = {k: [r[k] for r in rows] for k in rows[0]}
        if return_tensors is not None:
            mx = max(len(r["input_ids"]) for r in rows)
            pad_id = self._vocab[PAD]
            out["input_ids"] = [r["input_ids"] + [pad_id] * (mx - len(r["input_ids"])) for r in rows]
            out["attention_mask"] = [r["attention_mask"] + [0] * (mx - len(r["attention_mask"])) for r in rows]
            if return_offsets_mapping:
                out["offset_mapping"] = [r["offset_mapping"] + [(0, 0)] * (mx - len(r["offset_mapping"])) for r in rows]
        return BatchEncoding(out, tensor_type=return_tensors)