from __future__ import annotations

import json
from pathlib import Path
from typing import Iterable


class CharacterTokenizer:
    PAD = "[PAD]"
    UNK = "[UNK]"
    BOS = "[BOS]"
    EOS = "[EOS]"

    def __init__(self, vocab: dict[str, int], num_masks: int):
        self.vocab = dict(vocab)
        self.id_to_token = {index: token for token, index in self.vocab.items()}
        self.num_masks = num_masks
        self.pad_id = self.vocab[self.PAD]
        self.unk_id = self.vocab[self.UNK]
        self.bos_id = self.vocab[self.BOS]
        self.eos_id = self.vocab[self.EOS]
        self.mask_token_start = 4
        self.clean_token_start = 4 + num_masks

    @classmethod
    def build(cls, texts: Iterable[str], num_masks: int = 8) -> "CharacterTokenizer":
        specials = [cls.PAD, cls.UNK, cls.BOS, cls.EOS]
        masks = [f"[M{i}]" for i in range(num_masks)]
        characters = sorted(set("".join(texts)))
        vocab = {
            token: index
            for index, token in enumerate(specials + masks + characters)
        }
        return cls(vocab, num_masks)

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

    @property
    def clean_vocab_size(self) -> int:
        return self.vocab_size - self.clean_token_start

    def encode(self, text: str, seq_len: int | None = None) -> list[int]:
        ids = [self.vocab.get(character, self.unk_id) for character in text]
        if seq_len is not None:
            ids = ids[:seq_len]
            ids += [self.pad_id] * (seq_len - len(ids))
        return ids

    def decode(self, ids: Iterable[int], skip_pad: bool = True) -> str:
        pieces: list[str] = []
        for index in ids:
            token = self.id_to_token.get(int(index), self.UNK)
            if skip_pad and token == self.PAD:
                continue
            pieces.append(token)
        return "".join(pieces)

    def save(self, path: str | Path) -> None:
        payload = {"vocab": self.vocab, "num_masks": self.num_masks}
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open("w", encoding="utf-8") as handle:
            json.dump(payload, handle, indent=2, ensure_ascii=False)

    @classmethod
    def load(cls, path: str | Path) -> "CharacterTokenizer":
        with Path(path).open("r", encoding="utf-8") as handle:
            payload = json.load(handle)
        return cls(payload["vocab"], payload["num_masks"])