"""Simple editable eval for Tinker checkpoints.

Input JSONL format:
  {"messages": [{"role": "user", "content": "..."}], "expected": "...", "match": "contains"}

Match modes:
  - exact: normalized output equals normalized expected
  - contains: normalized expected appears in normalized output
  - prefix: normalized output starts with normalized expected

Examples:
  python eval.py --base-model thinkingmachines/Inkling --effort 0.9 --data data/eval.jsonl --out eval_results/baseline.json
  python eval.py --model-path 'tinker://.../sampler_weights/...' --renderer-model thinkingmachines/Inkling --effort 0.9 --data data/eval.jsonl --out eval_results/step-20.json
"""

from __future__ import annotations

import argparse
import asyncio
import json
import re
from pathlib import Path
from typing import Any

import tinker
from tinker_cookbook import model_info
from tinker_cookbook.renderers import get_renderer, get_text_content
from tinker_cookbook.tokenizer_utils import get_tokenizer


def normalize(text: str) -> str:
    text = text.strip().lower()
    text = re.sub(r"\s+", " ", text)
    text = text.strip(" \t\n\r.,;:!?")
    return text


def score_output(output: str, expected: str, match: str) -> bool:
    out = normalize(output)
    exp = normalize(expected)
    if match == "exact":
        return out == exp
    if match == "prefix":
        return out.startswith(exp)
    if match == "contains":
        return exp in out
    raise ValueError(f"unknown match mode: {match!r}")


def load_rows(path: str) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    with open(path) as f:
        for line_number, line in enumerate(f, start=1):
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            if not isinstance(row.get("messages"), list):
                raise ValueError(f"line {line_number}: missing messages[]")
            if "expected" not in row:
                raise ValueError(f"line {line_number}: missing expected")
            rows.append(row)
    return rows


async def main() -> None:
    parser = argparse.ArgumentParser(description="Evaluate a base model or Tinker sampler checkpoint.")
    target = parser.add_mutually_exclusive_group(required=True)
    target.add_argument("--base-model", help="Tinker base model id for baseline evaluation")
    target.add_argument("--model-path", help="Tinker sampler checkpoint path")
    parser.add_argument("--renderer-model", help="Model id used to choose renderer/tokenizer; defaults to --base-model or sampler base model")
    parser.add_argument("--data", default="data/eval.jsonl")
    parser.add_argument("--out", default="eval_results/result.json")
    parser.add_argument("--max-tokens", type=int, default=128)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--effort", type=float, default=0.9, help="Inkling effort in [0, 1); ignored by other renderers")
    parser.add_argument("--limit", type=int, default=0, help="0 means all rows")
    args = parser.parse_args()

    if not 0.0 <= args.effort < 1.0:
        raise SystemExit("--effort must be in [0, 1)")

    rows = load_rows(args.data)
    if args.limit:
        rows = rows[: args.limit]
    if not rows:
        raise SystemExit("No eval rows found.")

    service = tinker.ServiceClient()
    if args.model_path:
        sampling_client = await service.create_sampling_client_async(model_path=args.model_path)
        target_name = args.model_path
        renderer_model = args.renderer_model or sampling_client.get_base_model()
    else:
        sampling_client = await service.create_sampling_client_async(base_model=args.base_model)
        target_name = args.base_model
        renderer_model = args.renderer_model or args.base_model

    # Not sampling_client.get_tokenizer(): for Inkling models that path needs
    # the non-public tml_tokenizers package. The cookbook adapter works for
    # Inkling and falls back to HuggingFace for other models.
    tokenizer = get_tokenizer(renderer_model)
    renderer_name = model_info.get_recommended_renderer_name(renderer_model)
    renderer = get_renderer(renderer_name, tokenizer)
    stop = renderer.get_stop_sequences()

    results: list[dict[str, Any]] = []
    correct = 0

    for index, row in enumerate(rows, start=1):
        messages = row["messages"]
        expected = str(row["expected"])
        match = str(row.get("match", "contains"))
        prompt_kwargs = {"role": "assistant"}
        if renderer_model.startswith("thinkingmachines/Inkling"):
            prompt_kwargs["effort"] = args.effort
        prompt = renderer.build_generation_prompt(messages, **prompt_kwargs)
        response = await sampling_client.sample_async(
            prompt=prompt,
            num_samples=1,
            sampling_params=tinker.SamplingParams(
                max_tokens=args.max_tokens,
                temperature=args.temperature,
                stop=stop,
            ),
        )
        sequence = response.sequences[0] if hasattr(response, "sequences") else response.samples[0]
        parsed_message, termination = renderer.parse_response(sequence.tokens)
        output = get_text_content(parsed_message)
        is_correct = score_output(output, expected, match)
        correct += int(is_correct)
        results.append({
            "index": index,
            "correct": is_correct,
            "match": match,
            "expected": expected,
            "output": output,
            "termination": getattr(termination, "value", str(termination)),
            "messages": messages,
        })
        print(f"[{index}/{len(rows)}] {'✓' if is_correct else '✗'} expected={expected!r} output={output[:120]!r}")

    summary = {
        "target": target_name,
        "renderer_model": renderer_model,
        "renderer": renderer_name,
        "effort": args.effort if renderer_model.startswith("thinkingmachines/Inkling") else None,
        "data": args.data,
        "num_examples": len(rows),
        "num_correct": correct,
        "accuracy": correct / len(rows),
        "results": results,
    }
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(summary, indent=2, ensure_ascii=False) + "\n")
    print(f"accuracy={summary['accuracy']:.3f} ({correct}/{len(rows)})")
    print(f"wrote {out}")


if __name__ == "__main__":
    asyncio.run(main())