#!/usr/bin/env python3
"""Batch evaluation runner for agents."""

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from turnabout.agents.random_agent import RandomAgent, run_benchmark
from turnabout.envs.text_env import TextCourtEnv


def main():
    parser = argparse.ArgumentParser(description="Evaluate agents on Turnabout cases")
    parser.add_argument(
        "case",
        nargs="?",
        default=str(Path(__file__).parent.parent / "turnabout" / "cases" / "stolen_prototype.json"),
        help="Path to case JSON file",
    )
    parser.add_argument("-d", "--difficulty", choices=["easy", "hard"], default="easy")
    parser.add_argument("-n", "--episodes", type=int, default=100)
    parser.add_argument("--agent", choices=["random", "llm"], default="random")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("-v", "--verbose", action="store_true")
    parser.add_argument("-o", "--output", help="Save results to JSON file")
    args = parser.parse_args()

    print(f"Evaluating {args.agent} agent on {Path(args.case).stem}")
    print(f"Difficulty: {args.difficulty}, Episodes: {args.episodes}")
    print()

    if args.agent == "random":
        results = run_benchmark(
            case_path=args.case,
            difficulty=args.difficulty,
            n_episodes=args.episodes,
            seed=args.seed,
            verbose=args.verbose,
        )
    elif args.agent == "llm":
        try:
            from turnabout.agents.llm_agent import LLMAgent
        except ImportError as e:
            print(f"Error: {e}")
            sys.exit(1)

        agent = LLMAgent()
        results_list = []
        for i in range(args.episodes):
            env = TextCourtEnv(case_path=args.case, difficulty=args.difficulty)
            result = agent.run_episode(env, verbose=args.verbose and i == 0)
            results_list.append(result)
            print(f"  Episode {i + 1}: {'WON' if result['won'] else 'LOST'} "
                  f"(score={result['composite_score']:.3f})")

        wins = sum(r["won"] for r in results_list)
        results = {
            "n_episodes": args.episodes,
            "win_rate": wins / args.episodes,
            "avg_reward": sum(r["total_reward"] for r in results_list) / args.episodes,
            "avg_steps": sum(r["steps"] for r in results_list) / args.episodes,
            "avg_composite_score": sum(r["composite_score"] for r in results_list) / args.episodes,
        }

    print("Results:")
    for k, v in results.items():
        if isinstance(v, float):
            print(f"  {k}: {v:.4f}")
        else:
            print(f"  {k}: {v}")

    if args.output:
        with open(args.output, "w") as f:
            json.dump(results, f, indent=2)
        print(f"\nSaved to {args.output}")


if __name__ == "__main__":
    main()