"""Standalone tokenizer for Needle 3 ONNX models.

Decodes the SentencePiece BPE vocabulary embedded in tokenizer.json
without external sentencepiece or c++ binary dependencies.
"""

import json
import struct
from typing import Any, Dict, List, Optional, Tuple, Union

_SP_META_SPACE = "\u2581"
TK_NORMAL, TK_UNKNOWN, TK_CONTROL, TK_USER_DEFINED, TK_BYTE = 0, 1, 2, 3, 4
_TK_HDR = "<IIIIIBBH"
_TK_REC = "<fBH"

# Special token constants
PAD_ID = 0
EOS_ID = 1
BOS_ID = 2
UNK_ID = 3
IM_START_ID = 4
IM_END_ID = 5
THINK_START_ID = 6
THINK_END_ID = 7
TOOLS_START_ID = 8
TOOLS_END_ID = 9
TOOL_CALL_START_ID = 10
TOOL_CALL_END_ID = 11

IM_START = "<|im_start|>"
IM_END = "<|im_end|>"
THINK_START = "<think>"
THINK_END = "</think>"
TOOLS_START = "<tools>"
TOOLS_END = "</tools>"
TOOL_CALL_START = "<tool_call>"
TOOL_CALL_END = "</tool_call>"


def parse_tokenizer_blob(blob: bytes) -> Dict[str, Any]:
    """Parse raw SentencePiece vocabulary blob from needle3.cact if provided."""
    off = struct.calcsize(_TK_HDR)
    n, pad, eos, bos, unk, add_dummy, byte_fb, _ = struct.unpack_from(_TK_HDR, blob, 0)
    rec = struct.calcsize(_TK_REC)
    pieces, scores, types = [], [], []
    for _ in range(n):
        score, t, ln = struct.unpack_from(_TK_REC, blob, off)
        off += rec
        pieces.append(blob[off : off + ln].decode("utf-8"))
        scores.append(score)
        types.append(t)
        off += ln
    return {
        "pieces": pieces,
        "scores": scores,
        "types": types,
        "pad_id": pad,
        "eos_id": eos,
        "bos_id": bos,
        "unk_id": unk,
        "add_dummy_prefix": bool(add_dummy),
        "byte_fallback": bool(byte_fb),
    }


class NeedleTokenizer:
    """Tokenizer for Needle 3 models."""

    def __init__(self, meta: Dict[str, Any]):
        self.pieces = meta["pieces"]
        self.scores = meta["scores"]
        self.types = meta["types"]
        self.add_dummy = meta["add_dummy_prefix"]
        self.byte_fallback = meta["byte_fallback"]
        self.unk_id = meta["unk_id"]
        self.bos_id = meta.get("bos_id", BOS_ID)
        self.eos_id = meta.get("eos_id", EOS_ID)
        self.pad_id = meta.get("pad_id", PAD_ID)

        self.p2id = {p: i for i, p in enumerate(self.pieces)}
        self.byte_id = {
            int(p[3:5], 16): i
            for i, (p, t) in enumerate(zip(self.pieces, self.types))
            if t == TK_BYTE
        }
        self.markers = sorted(
            (p for p, t in zip(self.pieces, self.types) if t == TK_USER_DEFINED),
            key=len,
            reverse=True,
        )

    @classmethod
    def from_json(cls, json_path: str) -> "NeedleTokenizer":
        """Load tokenizer vocabulary from JSON file."""
        with open(json_path, "r", encoding="utf-8") as f:
            meta = json.load(f)
        return cls(meta)

    @classmethod
    def from_cact_blob(cls, blob: bytes) -> "NeedleTokenizer":
        meta = parse_tokenizer_blob(blob)
        return cls(meta)

    def _bpe(self, seg: str) -> List[int]:
        syms = list(seg)
        while len(syms) > 1:
            best_score, best_j = None, -1
            for j in range(len(syms) - 1):
                idx = self.p2id.get(syms[j] + syms[j + 1])
                if idx is not None and (best_score is None or self.scores[idx] > best_score):
                    best_score, best_j = self.scores[idx], j
            if best_j < 0:
                break
            syms[best_j : best_j + 2] = [syms[best_j] + syms[best_j + 1]]
        ids = []
        for s in syms:
            idx = self.p2id.get(s)
            if idx is not None:
                ids.append(idx)
            elif self.byte_fallback:
                ids.extend(self.byte_id[b] for b in s.encode("utf-8"))
            else:
                ids.append(self.unk_id)
        return ids

    def encode(self, text: str) -> List[int]:
        if not text:
            return []
        esc = text.replace(" ", _SP_META_SPACE)
        if self.add_dummy:
            esc = _SP_META_SPACE + esc
        ids, buf, i, n = [], [], 0, len(esc)
        while i < n:
            marker = next((m for m in self.markers if esc.startswith(m, i)), None)
            if marker is not None:
                ids += self._bpe("".join(buf))
                buf = []
                ids.append(self.p2id[marker])
                i += len(marker)
            else:
                buf.append(esc[i])
                i += 1
        ids += self._bpe("".join(buf))
        return ids

    def decode(self, ids: List[int]) -> str:
        buf = bytearray()
        for i in ids:
            t = self.types[i]
            if t == TK_BYTE:
                buf.append(int(self.pieces[i][3:5], 16))
            elif t in (TK_CONTROL, TK_UNKNOWN):
                continue
            else:
                buf += self.pieces[i].encode("utf-8")
        text = buf.decode("utf-8", "replace").replace(_SP_META_SPACE, " ")
        if self.add_dummy and text.startswith(" "):
            text = text[1:]
        return text

    def build_prompt(
        self,
        query: str,
        tools: Optional[List[Dict[str, Any]]] = None,
        system: Optional[str] = None,
    ) -> Tuple[str, List[int]]:
        """Construct prompt and tokenize with BOS token."""
        prefix = ""
        if system:
            prefix = f"{IM_START}system\n{system}{IM_END}\n"
        tools_str = ""
        if tools is not None:
            tools_json = json.dumps(tools, separators=(",", ":"), ensure_ascii=False)
            tools_str = f"{TOOLS_START}{tools_json}{TOOLS_END}\n"
        prompt = f"{prefix}{IM_START}user\n{tools_str}{query}{IM_END}\n{IM_START}assistant\n"
        ids = [self.bos_id] + self.encode(prompt)
        return prompt, ids

    def parse_tool_calls(self, output_ids: List[int]) -> Tuple[str, List[Dict[str, Any]]]:
        """Extract reasoning and tool call JSON from decoded tokens."""
        pieces = [self.pieces[i] for i in output_ids]
        raw = "".join(pieces).replace(_SP_META_SPACE, " ")

        reasoning = ""
        if THINK_START in raw and THINK_END in raw:
            reasoning = raw.split(THINK_START, 1)[1].split(THINK_END, 1)[0].strip()

        calls = []
        if TOOL_CALL_START in raw:
            body = raw.split(TOOL_CALL_START, 1)[1].split(TOOL_CALL_END, 1)[0].strip()
            if body.startswith("[]name") or body.startswith('[]"name'):
                body = '[{"' + body[2:]
            try:
                parsed = json.loads(body)
                if isinstance(parsed, list):
                    calls = parsed
                elif isinstance(parsed, dict):
                    calls = [parsed]
            except Exception:
                calls = []
        return reasoning, calls