#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""G4.2 canary — token-count delta measurement (system.md, tiktoken cl100k_base).

Computes baseline and candidate token counts for system.md and emits the delta.
Used to verify Rule A/C/E lift achieves the expected token reduction
(target threshold: candidate - baseline <= -850).

Exit criteria (analyzer):
  exit_criteria_pass == True  iff  token_delta <= -850

Usage (config mode):
    python scripts/canary/g4_2_token_count.py \\
        --config canary_config.json \\
        --baseline-version 16.202605041200 \\
        --candidate-version 17.202605042018 \\
        --output results/token_count.json
"""
from __future__ import annotations

import argparse
import json
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Dict


def iso8601_now() -> str:
    return datetime.utcnow().isoformat() + "Z"


def _fail(msg: str, code: int = 1) -> None:
    print(f"[g4_2_token_count] ERROR: {msg}", file=sys.stderr)
    sys.exit(code)


def load_pinning(args: argparse.Namespace) -> Dict[str, Any]:
    if args.config:
        cfg_path = Path(args.config)
        if not cfg_path.exists():
            _fail(f"--config file not found: {cfg_path}")
        try:
            with cfg_path.open("r", encoding="utf-8") as fp:
                cfg = json.load(fp)
        except Exception as exc:
            _fail(f"--config parse error: {exc}")
        return {
            "pid": cfg.get("pid"),
            "scene_index_list": cfg.get("scene_index_list"),
            "shot_index_list_per_scene": cfg.get("shot_index_list_per_scene"),
            "model_routing": cfg.get("model_routing"),
            "prompt_source_mode": cfg.get("prompt_source_mode"),
            "card_commit_hash": cfg.get("card_commit_hash"),
            "chain_bg_card_snapshot_hash": cfg.get("chain_bg_card_snapshot_hash"),
        }
    try:
        scene_index_list = (
            json.loads(args.scene_index_list) if args.scene_index_list else None
        )
        shot_index_list_per_scene = (
            json.loads(args.shot_index_list_per_scene)
            if args.shot_index_list_per_scene else None
        )
    except Exception as exc:
        _fail(f"failed to parse list/dict JSON args: {exc}")
    return {
        "pid": args.pid,
        "scene_index_list": scene_index_list,
        "shot_index_list_per_scene": shot_index_list_per_scene,
        "model_routing": args.model_routing,
        "prompt_source_mode": args.prompt_source_mode,
        "card_commit_hash": args.card_commit_hash,
        "chain_bg_card_snapshot_hash": args.chain_bg_card_snapshot_hash,
    }


def validate_pinning(pinning: Dict[str, Any]) -> None:
    required = [
        "pid", "scene_index_list", "shot_index_list_per_scene",
        "model_routing", "prompt_source_mode",
        "card_commit_hash", "chain_bg_card_snapshot_hash",
    ]
    missing = [k for k in required if pinning.get(k) in (None, "")]
    if missing:
        _fail(
            f"pinning missing required field(s): {missing}. "
            f"Provide via --config <json> or individual flags. NO silent defaults."
        )


def main() -> int:
    parser = argparse.ArgumentParser(
        description="G4.2 canary: token-count delta (system.md, tiktoken cl100k_base)."
    )
    parser.add_argument("--config", type=str, default=None)
    parser.add_argument("--baseline-version", type=str, required=True,
                        help="Baseline prompt version directory name (e.g. 16.202605041200)")
    parser.add_argument("--candidate-version", type=str, required=True,
                        help="Candidate prompt version directory name (e.g. 17.202605042018)")
    parser.add_argument("--prompt-base-path", type=str,
                        default="prompts/_base/scene_detail",
                        help="Base directory containing version subdirs")
    parser.add_argument("--exit-criteria-threshold", type=int, default=-850,
                        help="token_delta upper bound for pass (default: -850)")
    # 7 pinning fields (echoed for cross-canary consistency).
    parser.add_argument("--pid", type=str, default=None)
    parser.add_argument("--scene-index-list", type=str, default=None)
    parser.add_argument("--shot-index-list-per-scene", type=str, default=None)
    parser.add_argument("--model-routing", type=str, default=None)
    parser.add_argument("--prompt-source-mode", type=str, default=None,
                        choices=[None, "file", "db"])
    parser.add_argument("--card-commit-hash", type=str, default=None)
    parser.add_argument("--chain-bg-card-snapshot-hash", type=str, default=None)
    parser.add_argument("--output", type=str, required=True)
    args = parser.parse_args()

    pinning = load_pinning(args)
    validate_pinning(pinning)

    try:
        import tiktoken  # type: ignore
    except ImportError:
        _fail(
            "tiktoken not installed. Install via `pip install tiktoken` "
            "(or use the backend venv: backend/.venv/bin/pip install tiktoken). "
            "NO silent fallback."
        )

    encoder = tiktoken.get_encoding("cl100k_base")  # type: ignore[name-defined]

    base_dir = Path(args.prompt_base_path)
    baseline_path = base_dir / args.baseline_version / "system.md"
    candidate_path = base_dir / args.candidate_version / "system.md"

    if not baseline_path.exists():
        _fail(f"baseline system.md not found: {baseline_path}")
    if not candidate_path.exists():
        _fail(f"candidate system.md not found: {candidate_path}")

    baseline_text = baseline_path.read_text(encoding="utf-8")
    candidate_text = candidate_path.read_text(encoding="utf-8")

    baseline_token_count = len(encoder.encode(baseline_text))
    candidate_token_count = len(encoder.encode(candidate_text))
    token_delta = candidate_token_count - baseline_token_count
    exit_pass = token_delta <= args.exit_criteria_threshold

    out = {
        "timestamp": iso8601_now(),
        "pinning": pinning,
        "baseline_prompt_version": args.baseline_version,
        "candidate_prompt_version": args.candidate_version,
        "baseline_token_count": baseline_token_count,
        "candidate_token_count": candidate_token_count,
        "token_delta": token_delta,
        "exit_criteria_threshold": args.exit_criteria_threshold,
        "exit_criteria_pass": exit_pass,
        "measurement_tool": "tiktoken cl100k_base",
        "measurement_scripts": {
            "token_count": "scripts/canary/g4_2_token_count.py (system.md, tiktoken cl100k_base)"
        },
    }

    if args.output == "-":
        json.dump(out, sys.stdout, ensure_ascii=False, indent=2, sort_keys=True)
        sys.stdout.write("\n")
    else:
        out_path = Path(args.output)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        with out_path.open("w", encoding="utf-8") as fp:
            json.dump(out, fp, ensure_ascii=False, indent=2, sort_keys=True)

    print(
        f"[g4_2_token_count]\n"
        f"  baseline ({args.baseline_version}): {baseline_token_count} tokens\n"
        f"  candidate ({args.candidate_version}): {candidate_token_count} tokens\n"
        f"  token_delta: {token_delta}\n"
        f"  threshold: {args.exit_criteria_threshold}\n"
        f"  exit_criteria_pass: {exit_pass}",
        file=sys.stderr,
    )
    # Codex 1 fix — non-zero exit when token-reduction gate fails.
    # exit_pass is True iff token_delta <= threshold (default -850).
    # Returning 0 unconditionally would mask gate failure for CI/shell
    # callers — they'd see success while the threshold is violated.
    return 0 if exit_pass else 1


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