"""The tokenizer of bard-small."""
import os
import re

import numpy as np
import tokenmonster

VOCAB_FILE = "bard-16k.vocab"
FORMAT_FILE = "format.txt"
EOS = "<|end|>"
NUM_MARKERS = ["<|num|>", "<|snum|>"]
MAX_DIGITS = 8
LEN_SHIFT = 27
VALUE_MASK = (1 << LEN_SHIFT) - 1
NUMBER_RE = re.compile(r"( ?)([0-9]+)")


def pack(digits: str) -> int:
    assert 0 < len(digits) <= MAX_DIGITS
    return int(digits) | (len(digits) << LEN_SHIFT)


def unpack(packed: int) -> str:
    return str(packed & VALUE_MASK).zfill(packed >> LEN_SHIFT)


def split_numbers(text):
    out, pos = [], 0
    for m in NUMBER_RE.finditer(text):
        if m.start() > pos:
            out.append(text[pos:m.start()])
        spaced, digits = bool(m.group(1)), m.group(2)
        for i in range(0, len(digits), MAX_DIGITS):
            out.append((spaced and i == 0, pack(digits[i:i + MAX_DIGITS])))
        pos = m.end()
    if pos < len(text):
        out.append(text[pos:])
    return out


def parse_special_tokens(format_text):
    markers = re.findall(r"<\|[a-z_]+\|>", format_text)
    roles = re.findall(r"<\|start\|>([a-z_]+)<\|content\|>", format_text)
    return list(dict.fromkeys(markers + [f"<|start|>{r}" for r in roles]))


class BardTokenizer:
    def __init__(self, vocab_path, format_path):
        self.vocab = tokenmonster.load(vocab_path)
        self.specials = {t: self.vocab.token_to_id(t)
                         for t in parse_special_tokens(open(format_path, encoding="utf-8").read()) + NUM_MARKERS}
        missing = [t for t, i in self.specials.items() if i is None]
        assert not missing, f"the vocabulary lacks {missing}"
        self.eos_id = self.specials[EOS]
        self.num_id, self.snum_id = (self.specials[t] for t in NUM_MARKERS)
        self._special_text = {i: t for t, i in self.specials.items()}

    @property
    def vocab_size(self):
        return max(self.specials.values()) + 1

    def encode(self, text):
        plan = split_numbers(text)
        segs = [item for item in plan if isinstance(item, str)]
        seg_ids = iter(self.vocab.tokenize(segs) if segs else [])
        ids, vals = [], []
        for item in plan:
            if isinstance(item, str):
                t = next(seg_ids)
                ids.append(t.astype(np.int64, copy=False))
                vals.append(np.zeros(len(t), np.int64))
            else:
                spaced, packed = item
                ids.append(np.array([self.snum_id if spaced else self.num_id], np.int64))
                vals.append(np.array([packed], np.int64))
        if not ids:
            return np.zeros(0, np.int64), np.zeros(0, np.int64)
        return np.concatenate(ids), np.concatenate(vals)

    def decode(self, ids, vals=None):
        ids = np.asarray(ids).tolist()
        vals = np.asarray(vals).tolist() if vals is not None else [0] * len(ids)
        parts, run = [], []
        for i, v in zip(ids, vals):
            if i in (self.num_id, self.snum_id):
                if run:
                    parts.append(self.vocab.decode(run))
                    run = []
                parts.append((" " if i == self.snum_id else "") + unpack(v))
            elif i in self._special_text:

                if run:
                    parts.append(self.vocab.decode(run))
                    run = []
                parts.append(self._special_text[i])
            else:
                run.append(i)
        if run:
            parts.append(self.vocab.decode(run))
        return "".join(parts)

    @classmethod
    def from_dir(cls, path):
        return cls(os.path.join(path, VOCAB_FILE), os.path.join(path, FORMAT_FILE))