"""AbyssTokenizer: byte-level BPE encode/decode.

Source of truth on disk is merges.txt (integer id pairs, in priority
order). The vocab (id -> bytes) is rebuilt by replaying those merges
over the 256 base bytes."""

import os
import json
import regex as re

from .regex_split import split
from .special_tokens import SPECIAL_TOKENS, special_token_ids


def merge_seq(ids, pair, new_id):
    """Replace every adjacent occurence of 'pair' in 'ids' with new_id."""
    out = []
    i = 0
    while i < len(ids):
        if i < len(ids) - 1 and ids[i] == pair[0] and ids[i + 1] == pair[1]:
            out.append(new_id)
            i += 2
        else:
            out.append(ids[i])
            i += 1
    return out


class AbyssTokenizer:
    def __init__(self, merges, vocab_size):
        # merges: ordered list of (a, b) integer id pairs
        self.vocab_size = vocab_size
        self.merges = merges


        # rank[(a, b)] = priority index (lower = merged earlier)
        self.ranks = {pair: i for i, pair in enumerate(merges)}

        # rebuild id -> bytes
        self.vocab = {i: bytes([i]) for i in range(256)}
        for i, (a, b) in enumerate(merges):
            self.vocab[256 + i] = self.vocab[a] + self.vocab[b]

        # special tokens
        self.special = special_token_ids(vocab_size)
        self.inv_special = {v: k for k, v in self.special.items()}
        self._special_pat = re.compile(
            "(" + "|".join(re.escape(t) for t in SPECIAL_TOKENS) + ")"
        )

    # ---- Encode ----
    def _encode_chunk(self, ids):
        while len(ids) >= 2:
            pairs = set(zip(ids, ids[1:]))
            best = min(pairs, key=lambda p: self.ranks.get(p, float("inf")))
            if best not in self.ranks:
                break
            new_id = 256 + self.ranks[best]
            ids = merge_seq(ids, best, new_id)
        return ids
    
    def encode(self, text):
        out = []
        for piece in self._special_pat.split(text):
            if piece == "":
                continue
            if piece in self.special:
                out.append(self.special[piece])
            else:
                for chunk in split(piece):
                    ids = list(chunk.encode("utf-8"))
                    out.extend(self._encode_chunk(ids))
        return out
    

    # ---- Decode ----
    def decode(self, ids):
        parts = []
        for i in ids:
            if i in self.inv_special:
                parts.append(self.inv_special[i].encode("utf-8"))
            else:
                parts.append(self.vocab[i])
        return b"".join(parts).decode("utf-8", errors="replace")
    
    # ---- load ----
    @classmethod
    def load(cls, in_dir):
        with open(os.path.join(in_dir, "config.json")) as f:
            vocab_size = json.load(f)["vocab_size"]
        merges = []
        with open(os.path.join(in_dir, "merges.txt")) as f:
            for line in f:
                a, b = line.split()
                merges.append((int(a), int(b)))
        return cls(merges, vocab_size)