"""Evaluate a saved Unsloth/TRL GRPO checkpoint on the Red Shift eval rows."""

from __future__ import annotations

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

from train_unsloth_grpo import evaluate_model


def checkpoint_step(path: Path) -> int:
    match = re.search(r"checkpoint-(\d+)$", path.name)
    return int(match.group(1)) if match else -1


def latest_checkpoint(run_dir: Path) -> Path:
    checkpoints = sorted(run_dir.glob("checkpoint-*"), key=checkpoint_step)
    if not checkpoints:
        raise FileNotFoundError(f"No checkpoint-* directories found under {run_dir}")
    return checkpoints[-1]


def load_rows(dataset_path: Path, eval_tasks: int) -> list[dict[str, Any]]:
    rows = [json.loads(line) for line in dataset_path.read_text(encoding="utf-8").splitlines() if line.strip()]
    return rows[: min(eval_tasks, len(rows))]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--run-dir", type=Path, default=Path("training_results/unsloth_grpo_qwen3b_easy"))
    parser.add_argument("--checkpoint", type=Path, default=None)
    parser.add_argument("--model-name", default="unsloth/Qwen2.5-3B-Instruct-bnb-4bit")
    parser.add_argument("--max-seq-length", type=int, default=1280)
    parser.add_argument("--max-completion-length", type=int, default=256)
    parser.add_argument("--eval-tasks", type=int, default=32)
    parser.add_argument("--reward-mode", choices=["hard", "easy"], default="easy")
    parser.add_argument("--out-name", default="checkpoint_generations.json")
    args = parser.parse_args()

    checkpoint = args.checkpoint or latest_checkpoint(args.run_dir)
    dataset_path = args.run_dir / "dataset.jsonl"
    if not dataset_path.exists():
        raise FileNotFoundError(f"Missing {dataset_path}; run training at least until dataset creation.")
    eval_rows = load_rows(dataset_path, args.eval_tasks)

    from peft import PeftModel
    from unsloth import FastLanguageModel

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=args.model_name,
        max_seq_length=args.max_seq_length,
        load_in_4bit=True,
        fast_inference=False,
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = PeftModel.from_pretrained(model, str(checkpoint), is_trainable=False)
    try:
        FastLanguageModel.for_inference(model)
    except Exception:
        pass

    generations_path = args.run_dir / args.out_name
    checkpoint_eval = evaluate_model(
        model,
        tokenizer,
        eval_rows,
        generations_path,
        args.max_completion_length,
        args.reward_mode,
    )

    baseline_path = args.run_dir / "baseline_generations.json"
    baseline = json.loads(baseline_path.read_text(encoding="utf-8")) if baseline_path.exists() else {}
    summary = {
        "status": "checkpoint_eval",
        "run_dir": str(args.run_dir),
        "checkpoint": str(checkpoint),
        "checkpoint_step": checkpoint_step(checkpoint),
        "model_name": args.model_name,
        "num_eval_tasks": len(eval_rows),
        "reward_mode": args.reward_mode,
        "baseline_mean_reward": baseline.get("mean_reward"),
        "checkpoint_mean_reward": checkpoint_eval["mean_reward"],
        "generations_path": str(generations_path),
    }
    (args.run_dir / "checkpoint_eval_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()