#!/usr/bin/env python3
"""Evaluation harness for aave-compress: measures REAL compression ratios.

Why this script is paranoid about its own numbers
-------------------------------------------------
A competing tool ("caveman") marketed 65-75% token savings. Independent
re-benchmarking found the real number was 14-21%, because the original claim
was measured against an artificially verbose, non-representative baseline
corpus. This harness exists specifically to not repeat that mistake:

* The corpus (tests/fixtures/real_prompts.jsonl) is 20 realistic prompts in
  the register people actually type to an assistant. Several intentionally
  contain ZERO applicable AAVE markers — that is an honest, expected outcome
  and it is reported, not hidden.
* Claude token counts come ONLY from Anthropic's count_tokens API
  (requires ANTHROPIC_API_KEY). tiktoken is NEVER used as a stand-in for a
  Claude model — it undercounts Claude tokens by 15-20%+.
* --tiktoken-model provides an explicitly-labeled APPROXIMATE fallback for
  GPT-family models only.
* If no tokenizer is available at all, the run still completes and reports
  marker/word-level results, and says plainly that no token counts exist —
  it never invents a percentage.

Usage:
    python eval/run_eval.py                         # Claude counts (needs ANTHROPIC_API_KEY)
    python eval/run_eval.py --model claude-opus-4-7 # different Claude target
    python eval/run_eval.py --tiktoken-model gpt-4o # approximate GPT-family fallback

Exit code 0 on any successful run (even with no token counts available);
exit code 1 only on an actual crash/bug.
"""

from __future__ import annotations

import argparse
import json
import sys
from datetime import datetime, timezone
from pathlib import Path
from statistics import mean

REPO_ROOT = Path(__file__).resolve().parent.parent
SRC_DIR = REPO_ROOT / "src"
if str(SRC_DIR) not in sys.path:
    sys.path.insert(0, str(SRC_DIR))

from aave_compress.engine import compress  # noqa: E402
from aave_compress.tokenizers import (  # noqa: E402
    count_tokens,
    count_tokens_tiktoken,
)

FIXTURES_PATH = REPO_ROOT / "tests" / "fixtures" / "real_prompts.jsonl"
RESULTS_DIR = REPO_ROOT / "eval" / "results"
LEVELS = ("light", "medium", "full")

NO_TOKENS_BANNER = (
    "NO TOKEN COUNTS AVAILABLE — install/configure a tokenizer to get real numbers"
)


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Measure real aave-compress compression ratios on a realistic prompt corpus."
    )
    parser.add_argument(
        "--model",
        default="claude-sonnet-5",
        help="Target model for token counting (default: claude-sonnet-5). "
        "Claude models require ANTHROPIC_API_KEY.",
    )
    parser.add_argument(
        "--tiktoken-model",
        default=None,
        metavar="GPT_MODEL",
        help="APPROXIMATE fallback: count tokens with tiktoken for this "
        "GPT-family model (e.g. gpt-4o) if the primary counter is "
        "unavailable. NEVER valid for Claude models.",
    )
    parser.add_argument(
        "--strip-filler",
        action="store_true",
        help="Run the generic filler-strip pass (filters.strip_fillers) "
        "before the AAVE markers, so the two layers can be compared.",
    )
    args = parser.parse_args(argv)
    if args.tiktoken_model and args.tiktoken_model.lower().startswith("claude"):
        parser.error(
            "--tiktoken-model must be a GPT-family model: tiktoken counts are "
            "never valid for Claude models (they undercount by 15-20%+). "
            "Set ANTHROPIC_API_KEY and use --model instead."
        )
    return args


def resolve_counter(args: argparse.Namespace):
    """Probe tokenizer availability ONCE, up front.

    Returns (counter_fn | None, label, method | None). counter_fn takes a
    string and returns an int token count. Errors are reported here a single
    time instead of once per (prompt, level) pair.
    """
    probe = "tokenizer availability probe"

    try:
        probed = count_tokens(probe, model=args.model)
        if probed.method == "anthropic_api":
            label = f"real Anthropic count_tokens API for {args.model}"
        else:
            label = (
                f"APPROXIMATE tiktoken counts for {args.model} "
                "(GPT-family only — NOT valid for any Claude model)"
            )
        return (
            lambda text: count_tokens(text, model=args.model).tokens,
            label,
            probed.method,
        )
    except RuntimeError as exc:
        # Expected failure mode: missing ANTHROPIC_API_KEY (or missing
        # optional package). Reported once, here, not 60 times below.
        print("=" * 78)
        print(f"Token counting with {args.model!r} is UNAVAILABLE:")
        print(f"  {exc}")
        print()
        print(
            "To get real Claude token numbers, set the ANTHROPIC_API_KEY "
            "environment variable and re-run."
        )
    except Exception as exc:  # unexpected API failure — degrade, don't crash
        print("=" * 78)
        print(f"Token counting with {args.model!r} failed unexpectedly:")
        print(f"  {type(exc).__name__}: {exc}")

    if args.tiktoken_model:
        try:
            count_tokens_tiktoken(probe, model=args.tiktoken_model)
        except RuntimeError as exc:
            print()
            print(f"tiktoken fallback also unavailable: {exc}")
            print("=" * 78)
            return None, NO_TOKENS_BANNER, None
        label = (
            f"APPROXIMATE tiktoken counts for {args.tiktoken_model} "
            "(GPT-family only — NOT valid for any Claude model)"
        )
        print()
        print(f"Falling back to: {label}")
        print(
            "These numbers describe how a GPT-family tokenizer sees the text. "
            f"Do NOT quote them as {args.model} savings."
        )
        print("=" * 78)
        return (
            lambda text: count_tokens_tiktoken(text, model=args.tiktoken_model).tokens,
            label,
            "tiktoken_approx",
        )

    print()
    print(
        "No fallback tokenizer requested. Pass --tiktoken-model gpt-4o for an "
        "explicitly-approximate GPT-family count, or set ANTHROPIC_API_KEY "
        "for real Claude numbers."
    )
    print("=" * 78)
    return None, NO_TOKENS_BANNER, None


class SafeCounter:
    """Wraps a counter so one mid-run failure doesn't spam or crash the eval."""

    def __init__(self, fn):
        self._fn = fn
        self.failed = False

    def count(self, text: str) -> int | None:
        if self._fn is None or self.failed:
            return None
        try:
            return self._fn(text)
        except Exception as exc:
            self.failed = True
            print(
                f"WARNING: token counting failed mid-run "
                f"({type(exc).__name__}: {exc}); remaining token counts "
                "recorded as null."
            )
            return None


def load_corpus() -> list[dict]:
    prompts = []
    with open(FIXTURES_PATH, encoding="utf-8") as fh:
        for line_no, line in enumerate(fh, start=1):
            line = line.strip()
            if not line:
                continue
            record = json.loads(line)
            for key in ("id", "text", "domain"):
                if key not in record:
                    raise ValueError(
                        f"{FIXTURES_PATH}:{line_no} missing required key {key!r}"
                    )
            prompts.append(record)
    return prompts


def evaluate(
    prompts: list[dict], counter: SafeCounter, strip_filler: bool = False
) -> list[dict]:
    rows = []
    for prompt in prompts:
        text = prompt["text"]
        original_tokens = counter.count(text)
        original_words = len(text.split())
        record = {
            "id": prompt["id"],
            "domain": prompt["domain"],
            "original": text,
            "original_tokens": original_tokens,
            "original_words": original_words,
            "levels": {},
        }
        for level in LEVELS:
            result = compress(text, level, strip_filler=strip_filler)
            if result.compressed == text:
                compressed_tokens = original_tokens  # no change -> no extra call
            else:
                compressed_tokens = counter.count(result.compressed)
            compressed_words = len(result.compressed.split())
            token_ratio = (
                compressed_tokens / original_tokens
                if compressed_tokens is not None
                and original_tokens is not None
                and original_tokens > 0
                else None
            )
            record["levels"][level] = {
                "compressed": result.compressed,
                "markers_applied": result.markers_applied,
                "compressed_tokens": compressed_tokens,
                "token_ratio": token_ratio,
                "compressed_words": compressed_words,
                "word_ratio": (
                    compressed_words / original_words if original_words else None
                ),
            }
        rows.append(record)
    return rows


def summarize(rows: list[dict]) -> dict:
    summary: dict[str, dict] = {}
    for level in LEVELS:
        token_ratios = [
            r["levels"][level]["token_ratio"]
            for r in rows
            if r["levels"][level]["token_ratio"] is not None
        ]
        word_ratios = [
            r["levels"][level]["word_ratio"]
            for r in rows
            if r["levels"][level]["word_ratio"] is not None
        ]
        zero_marker = sum(
            1 for r in rows if not r["levels"][level]["markers_applied"]
        )
        orig_token_sum = sum(
            r["original_tokens"]
            for r in rows
            if r["original_tokens"] is not None
            and r["levels"][level]["compressed_tokens"] is not None
        )
        comp_token_sum = sum(
            r["levels"][level]["compressed_tokens"]
            for r in rows
            if r["original_tokens"] is not None
            and r["levels"][level]["compressed_tokens"] is not None
        )
        summary[level] = {
            "prompts": len(rows),
            "zero_marker_prompts": zero_marker,
            "prompts_with_token_counts": len(token_ratios),
            "mean_token_ratio": mean(token_ratios) if token_ratios else None,
            "min_token_ratio": min(token_ratios) if token_ratios else None,
            "max_token_ratio": max(token_ratios) if token_ratios else None,
            "corpus_token_ratio": (
                comp_token_sum / orig_token_sum if orig_token_sum else None
            ),
            "mean_word_ratio": mean(word_ratios) if word_ratios else None,
            "min_word_ratio": min(word_ratios) if word_ratios else None,
            "max_word_ratio": max(word_ratios) if word_ratios else None,
        }
    return summary


def _fmt(value: float | None, pattern: str = "{:.3f}") -> str:
    return pattern.format(value) if value is not None else "n/a"


def print_table(title: str, summary: dict, metric_prefix: str) -> None:
    print()
    print(title)
    header = (
        f"{'level':<8} {'prompts':>7} {'zero-marker':>11} "
        f"{'mean ratio':>10} {'min':>7} {'max':>7} {'mean savings':>12}"
    )
    print(header)
    print("-" * len(header))
    for level in LEVELS:
        s = summary[level]
        mean_ratio = s[f"mean_{metric_prefix}_ratio"]
        savings = (
            f"{(1 - mean_ratio) * 100:.1f}%" if mean_ratio is not None else "n/a"
        )
        print(
            f"{level:<8} {s['prompts']:>7} "
            f"{s['zero_marker_prompts']:>4}/{s['prompts']:<6} "
            f"{_fmt(mean_ratio):>10} "
            f"{_fmt(s[f'min_{metric_prefix}_ratio']):>7} "
            f"{_fmt(s[f'max_{metric_prefix}_ratio']):>7} "
            f"{savings:>12}"
        )
    print("(ratio = compressed/original; 1.000 = no change)")


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    prompts = load_corpus()
    print(f"aave-compress eval harness")
    print(f"Corpus: {FIXTURES_PATH.relative_to(REPO_ROOT)} ({len(prompts)} prompts)")

    counter_fn, counter_label, counter_method = resolve_counter(args)
    counter = SafeCounter(counter_fn)
    tokens_available = counter_fn is not None
    if tokens_available:
        print(f"Token counting: {counter_label}")
    if args.strip_filler:
        print("Filler strip: ON (strip_fillers pass runs before AAVE markers)")

    rows = evaluate(prompts, counter, strip_filler=args.strip_filler)
    summary = summarize(rows)

    if tokens_available and not counter.failed:
        print_table(
            f"TOKEN COMPRESSION — {counter_label}", summary, metric_prefix="token"
        )
    elif tokens_available and counter.failed:
        print()
        print(
            "Token counting failed partway through; per-prompt token counts "
            "are incomplete (see nulls in the results file)."
        )
        print_table(
            f"TOKEN COMPRESSION (PARTIAL) — {counter_label}",
            summary,
            metric_prefix="token",
        )
    else:
        print()
        print(NO_TOKENS_BANNER)

    print_table(
        "WORD-LEVEL COMPRESSION (crude, tokenizer-independent; words != tokens)",
        summary,
        metric_prefix="word",
    )

    zero_full = summary["full"]["zero_marker_prompts"]
    print()
    print(
        f"Honesty check: {zero_full}/{len(prompts)} prompts had ZERO applicable "
        "markers even at level=full. That is expected on a realistic corpus — "
        "do not cherry-pick them out when reporting numbers."
    )
    print(
        "Reminder: only quote ratios measured by this harness on this "
        "realistic corpus. Numbers measured against artificially verbose "
        "baselines are how the caveman tool's 65-75% claim became a real 14-21%."
    )

    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    output = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "corpus": str(FIXTURES_PATH.relative_to(REPO_ROOT)),
        "n_prompts": len(prompts),
        "strip_filler": args.strip_filler,
        "token_counting": {
            "available": tokens_available and not counter.failed,
            "partial": tokens_available and counter.failed,
            "label": counter_label,
            "method": counter_method,
            "requested_model": args.model,
            "tiktoken_fallback_model": args.tiktoken_model,
        },
        "summary": summary,
        "prompts": rows,
    }
    results_path = RESULTS_DIR / "latest.json"
    with open(results_path, "w", encoding="utf-8") as fh:
        json.dump(output, fh, indent=2, ensure_ascii=False)
        fh.write("\n")
    print()
    print(f"Full per-prompt results written to {results_path.relative_to(REPO_ROOT)}")
    return 0


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