#!/usr/bin/env python3
"""
token_estimator.py — deterministic, local token estimation for doc-hygiene.

Produces the `raw_tokens` count that feeds each report entry's
`token_estimate` field (see the frozen report schema:
`openspec/specs/report-schema/spec.md`). This is the deterministic seam
required by invariant #6 — NO model, NO network/API call, ever.

Accuracy goal
-------------
This is BALLPARK relative ranking of documentation bloat, not billing-grade
tokenization. A swappable backend lets a heavier (more accurate) tokenizer be
used when available, while a zero-dependency heuristic guarantees the seam
always works.

Backends (pluggable via the `TokenEstimator` ABC)
-------------------------------------------------
- `HeuristicEstimator` — ALWAYS available. Pure-Python ~4-chars/token rule
  (`ceil(len(text) / 4)`). Zero third-party dependencies.
- `TiktokenEstimator` — OPTIONAL accuracy upgrade. Uses tiktoken's
  `o200k_base` encoding. Only selected when tiktoken is importable AND its
  vocab is vendored in the local cache (no cold-cache network download).

`default_estimator()` performs the selection and NEVER raises — it always
returns a working estimator, falling back to the heuristic when tiktoken or
its vendored vocab is unavailable. The active backend is introspectable via
`.name` so the `check` skill can record which tokenizer produced the counts.

Offline / determinism guarantee (invariant #6)
----------------------------------------------
tiktoken has no "offline only" flag: calling `get_encoding("o200k_base")`
with a cold cache silently fetches the vocab over the network. We therefore
DETECT BEFORE FETCH — we set `TIKTOKEN_CACHE_DIR` to a path inside the plugin
and check that the expected cache file already exists *before* importing the
encoding. If it is absent we fall back to the heuristic rather than triggering
a download.

CACHE-FILENAME GOTCHA: tiktoken names its cache file by
`sha1(blobpath_url).hexdigest()` — NOT by the sha256 content hash and NOT by a
human-readable name. For `o200k_base` the URL is
`https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken`,
which hashes to the filename:

    fb374d419588a4632f3f557e76b4b70aebbca790

A vendored vocab MUST live at `<cache_dir>/fb374d419588a4632f3f557e76b4b70aebbca790`
or the cache never hits and tiktoken falls back to a network fetch.

Pre-warming the vendored vocab (one-time, online, optional)
-----------------------------------------------------------
    TIKTOKEN_CACHE_DIR=scripts/tiktoken_cache \
        python -c "import tiktoken; tiktoken.get_encoding('o200k_base')"

The ~2 MB vocab blob is intentionally NOT committed to git. When it is absent
the estimator runs on the heuristic backend; this keeps the repo lean and the
runtime offline-safe.
"""

from __future__ import annotations

import hashlib
import json
import math
import os
import sys
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Optional


# ---------------------------------------------------------------------------
# Constants — tiktoken vendored-cache resolution
# ---------------------------------------------------------------------------

# The encoding we standardise on (current OpenAI tokenizer; close enough for a
# relative ballpark of Claude-context bloat — accuracy goal is ranking, not
# billing).
_O200K_BASE_NAME = "o200k_base"
_O200K_BASE_BLOBPATH = (
    "https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken"
)

# tiktoken's on-disk cache file is named sha1(blobpath_url).hexdigest().
# Pre-computed here so the offline pre-check needs no tiktoken import.
_O200K_BASE_CACHE_FILENAME = hashlib.sha1(
    _O200K_BASE_BLOBPATH.encode()
).hexdigest()  # == "fb374d419588a4632f3f557e76b4b70aebbca790"

# Default vendored-cache directory inside the plugin (sibling of this file).
_DEFAULT_CACHE_DIR = Path(__file__).resolve().parent / "tiktoken_cache"


# ---------------------------------------------------------------------------
# Estimator abstraction
# ---------------------------------------------------------------------------

class TokenEstimator(ABC):
    """Pluggable, deterministic token-count backend.

    Subclasses implement `estimate(text)`; the file convenience and report
    wrapper are shared. Every backend exposes `.name` so callers can record
    which tokenizer produced a count.
    """

    #: Short, stable backend identifier (e.g. "heuristic", "tiktoken:o200k_base").
    name: str = "abstract"

    @abstractmethod
    def estimate(self, text: str) -> int:
        """Return the token count for *text* (>= 0). Deterministic, no I/O."""
        raise NotImplementedError

    def estimate_file(self, path: Path) -> int:
        """Return the token count of *path*'s UTF-8 contents.

        Decode errors are tolerated (``errors="replace"``) to match the
        scanner's read convention, so arbitrary/binary-ish docs never raise.
        A missing/unreadable file yields 0.
        """
        try:
            text = Path(path).read_text(encoding="utf-8", errors="replace")
        except OSError:
            return 0
        return self.estimate(text)

    def estimate_for_report(self, text: str) -> dict:
        """Return a schema-shaped ``token_estimate`` object for *text*.

        v1 populates only `raw_tokens` (the required field). The
        injection-frequency weighting fields are the v2 bonus and are emitted
        as `null` here so the shape is explicit for the `check` skill.
        See `openspec/specs/report-schema/spec.md` (Per-Entry Token Estimate).
        """
        return {
            "raw_tokens": self.estimate(text),
            "injection_frequency": None,
            "weighted_tokens": None,
        }


class HeuristicEstimator(TokenEstimator):
    """Zero-dependency ~4-chars/token estimator: ``ceil(len(text) / 4)``.

    Boundaries: ``""`` -> 0, 4 chars -> 1, 5 chars -> 2. Always available;
    this is the guaranteed fallback that keeps the deterministic seam working
    with no third-party packages installed.
    """

    name = "heuristic"

    #: Average characters per token for the heuristic. ~4 is the long-standing
    #: rule-of-thumb for English prose / markdown.
    CHARS_PER_TOKEN = 4

    def estimate(self, text: str) -> int:
        if not text:
            return 0
        return math.ceil(len(text) / self.CHARS_PER_TOKEN)


class TiktokenEstimator(TokenEstimator):
    """Accuracy-upgrade backend wrapping a tiktoken encoding.

    Constructed with an already-loaded encoding object so this class performs
    no import or cache resolution itself (that lives in `default_estimator`,
    which guarantees the encoding came from the vendored cache, not a network
    fetch). `disallowed_special=()` ensures arbitrary markdown containing
    special-token-looking substrings (e.g. ``<|endoftext|>``) never raises.
    """

    def __init__(self, encoding, encoding_name: str = _O200K_BASE_NAME) -> None:
        self._encoding = encoding
        self.name = f"tiktoken:{encoding_name}"

    def estimate(self, text: str) -> int:
        if not text:
            return 0
        return len(self._encoding.encode(text, disallowed_special=()))


# ---------------------------------------------------------------------------
# Backend selection (never raises)
# ---------------------------------------------------------------------------

def _vendored_vocab_present(cache_dir: Path) -> bool:
    """Return True iff the o200k_base vocab is already in *cache_dir*.

    Pure filesystem check by the sha1-of-blobpath filename — see the
    CACHE-FILENAME GOTCHA in the module docstring. This is the detect-before-
    fetch gate that keeps the seam offline (invariant #6).
    """
    return (cache_dir / _O200K_BASE_CACHE_FILENAME).is_file()


def default_estimator(cache_dir: Optional[Path] = None) -> TokenEstimator:
    """Return the best available working estimator. NEVER raises.

    Selection order:
      1. `TiktokenEstimator` — only if tiktoken is importable AND the
         o200k_base vocab is vendored in *cache_dir* (so loading it triggers
         no network fetch).
      2. `HeuristicEstimator` — the always-available fallback otherwise.

    Parameters
    ----------
    cache_dir:
        Directory to resolve the vendored tiktoken vocab from. Defaults to
        the plugin's `scripts/tiktoken_cache/`. Injectable so tests can point
        at an empty dir to force (and assert) the heuristic fallback without
        any network access.
    """
    resolved_cache = Path(cache_dir) if cache_dir is not None else _DEFAULT_CACHE_DIR

    # Detect-before-fetch: bail to heuristic unless the vocab is already local.
    if not _vendored_vocab_present(resolved_cache):
        return HeuristicEstimator()

    try:
        # Point tiktoken at the vendored cache BEFORE importing/using it.
        os.environ["TIKTOKEN_CACHE_DIR"] = str(resolved_cache)
        import tiktoken  # noqa: WPS433 (optional dependency, imported lazily)

        encoding = tiktoken.get_encoding(_O200K_BASE_NAME)
        return TiktokenEstimator(encoding, _O200K_BASE_NAME)
    except Exception:
        # Import error, corrupt cache, or any other failure — fall back.
        # The seam must always yield a working estimator (invariant #6).
        return HeuristicEstimator()


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------

def main(argv: Optional[list] = None) -> int:
    import argparse

    parser = argparse.ArgumentParser(
        description=(
            "doc-hygiene token estimator — emits a JSON token estimate for a "
            "file (deterministic, no model, no network)."
        )
    )
    parser.add_argument("path", help="Path to the file to estimate.")
    parser.add_argument(
        "--cache-dir",
        default=None,
        help="Vendored tiktoken cache dir (default: scripts/tiktoken_cache/).",
    )
    args = parser.parse_args(argv)

    target = Path(args.path)
    if not target.is_file():
        print(
            json.dumps({"error": f"not a file: {args.path}"}, indent=2),
            file=sys.stderr,
        )
        return 1

    estimator = default_estimator(
        cache_dir=Path(args.cache_dir) if args.cache_dir else None
    )
    raw_tokens = estimator.estimate_file(target)

    output = {
        "path": str(target),
        "backend": estimator.name,
        "token_estimate": {
            "raw_tokens": raw_tokens,
            "injection_frequency": None,
            "weighted_tokens": None,
        },
    }
    print(json.dumps(output, indent=2))
    return 0


if __name__ == "__main__":
    sys.exit(main())
