"""Exact token accounting with the frozen local Qwen tokenizer."""

from __future__ import annotations

from hashlib import sha256
from pathlib import Path
from typing import Sequence

from tokenizers import Tokenizer

from .components import Candidate


DEFAULT_TOKENIZER = (
    Path.home()
    / ".lmstudio/models/lmstudio-community/"
    "Qwen3.6-35B-A3B-MLX-4bit/tokenizer.json"
)


class QwenTokenCounter:
    def __init__(self, path: Path = DEFAULT_TOKENIZER):
        self.path = path.resolve()
        raw = self.path.read_bytes()
        self.sha256 = sha256(raw).hexdigest()
        self.tokenizer = Tokenizer.from_file(str(self.path))

    def count(self, text: str) -> int:
        return len(self.tokenizer.encode(text, add_special_tokens=False).ids)

    def pack_ranked(
        self,
        candidates: Sequence[Candidate],
        budget: int,
    ) -> tuple[str, tuple[Candidate, ...], int]:
        blocks: list[str] = []
        included: list[Candidate] = []
        used = 0
        for rank, candidate in enumerate(candidates, start=1):
            block = (
                f"\n--- Rank {rank}: {candidate.path} "
                f"(lines {candidate.line_start}-{candidate.line_end}; {candidate.source}) ---\n"
                f"{candidate.text}\n"
            )
            tokens = self.count(block)
            if used + tokens > budget:
                continue
            blocks.append(block)
            included.append(candidate)
            used += tokens
        return "".join(blocks), tuple(included), used