"""Main evaluation runner for toolsets."""

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

# Add project root to path
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

try:
    import logfire
except ImportError:
    logfire = None

from pydantic_ai_toolsets.evals.base import ResultCollector
from pydantic_ai_toolsets.evals.config import EvaluationConfig, default_config
from pydantic_ai_toolsets.evals.categories.multi_agent.compare_multi_agent import (
    compare_multi_agent_toolsets,
)
from pydantic_ai_toolsets.evals.categories.reflection.compare_reflection import (
    compare_reflection_toolsets,
)
from pydantic_ai_toolsets.evals.categories.thinking_cognition.compare_thinking import (
    compare_thinking_toolsets,
)
from pydantic_ai_toolsets.evals.categories.uniques.search_eval import evaluate_search_toolset
from pydantic_ai_toolsets.evals.categories.uniques.todo_eval import evaluate_todo_toolset


async def run_uniques(config: EvaluationConfig, sequential: bool = True) -> dict[str, Any]:
    """Run evaluations for unique toolsets using pydantic-evals Dataset API.
    
    Args:
        config: Evaluation configuration.
        sequential: If True, run toolsets one by one. If False, run in parallel.
    
    Returns:
        Dictionary of results.
    """
    print("Running evaluations for unique toolsets...")
    if sequential:
        print("Running toolsets sequentially (one by one)...\n")
    
    results = {}
    
    if sequential:
        # Run todo first
        print("-" * 60)
        print("TOOLSET 1/2: TODO")
        print("-" * 60)
        try:
            print("  Running todo toolset evaluation...")
            todo_report = await evaluate_todo_toolset(config)
            
            # Check what attributes the report has
            # EvaluationReport might have 'results' or 'cases' instead of 'case_results'
            if hasattr(todo_report, "case_results"):
                case_results = todo_report.case_results
            elif hasattr(todo_report, "results"):
                case_results = todo_report.results
            elif hasattr(todo_report, "cases"):
                case_results = todo_report.cases
            else:
                # Try to get cases from the report
                case_results = getattr(todo_report, "case_results", [])
            
            num_cases = len(case_results) if case_results else 0
            passed = sum(1 for r in case_results if getattr(r, "passed", False)) if case_results else 0
            failed = num_cases - passed
            
            print(f"  ✓ Todo: {num_cases} cases")
            print(f"    Passed: {passed}")
            print(f"    Failed: {failed}")
            results["todo"] = {
                "report": todo_report,
                "num_cases": num_cases,
                "passed": passed,
                "failed": failed,
            }
        except Exception as e:
            import traceback
            print(f"  ✗ Todo: {e}")
            traceback.print_exc()
            results["todo"] = {"error": str(e)}
        
        print("\n")
        
        # Then run search
        print("-" * 60)
        print("TOOLSET 2/2: SEARCH")
        print("-" * 60)
        try:
            print("  Running search toolset evaluation...")
            search_report = await evaluate_search_toolset(config)
            
            # Check what attributes the report has
            # EvaluationReport might have 'results' or 'cases' instead of 'case_results'
            if hasattr(search_report, "case_results"):
                case_results = search_report.case_results
            elif hasattr(search_report, "results"):
                case_results = search_report.results
            elif hasattr(search_report, "cases"):
                case_results = search_report.cases
            else:
                # Try to get cases from the report
                case_results = getattr(search_report, "case_results", [])
            
            num_cases = len(case_results) if case_results else 0
            passed = sum(1 for r in case_results if getattr(r, "passed", False)) if case_results else 0
            failed = num_cases - passed
            
            print(f"  ✓ Search: {num_cases} cases")
            print(f"    Passed: {passed}")
            print(f"    Failed: {failed}")
            results["search"] = {
                "report": search_report,
                "num_cases": num_cases,
                "passed": passed,
                "failed": failed,
            }
        except Exception as e:
            import traceback
            print(f"  ✗ Search: {e}")
            traceback.print_exc()
            results["search"] = {"error": str(e)}
    else:
        # Run in parallel (original behavior)
        import asyncio
        todo_task = asyncio.create_task(evaluate_todo_toolset(config))
        search_task = asyncio.create_task(evaluate_search_toolset(config))
        
        try:
            todo_report = await todo_task
            # Check what attributes the report has
            if hasattr(todo_report, "case_results"):
                case_results = todo_report.case_results
            elif hasattr(todo_report, "results"):
                case_results = todo_report.results
            elif hasattr(todo_report, "cases"):
                case_results = todo_report.cases
            else:
                case_results = getattr(todo_report, "case_results", [])
            
            num_cases = len(case_results) if case_results else 0
            passed = sum(1 for r in case_results if getattr(r, "passed", False)) if case_results else 0
            failed = num_cases - passed
            
            print(f"  ✓ Todo: {num_cases} cases")
            results["todo"] = {
                "report": todo_report,
                "num_cases": num_cases,
                "passed": passed,
                "failed": failed,
            }
        except Exception as e:
            print(f"  ✗ Todo: {e}")
            results["todo"] = {"error": str(e)}
        
        try:
            search_report = await search_task
            # Check what attributes the report has
            if hasattr(search_report, "case_results"):
                case_results = search_report.case_results
            elif hasattr(search_report, "results"):
                case_results = search_report.results
            elif hasattr(search_report, "cases"):
                case_results = search_report.cases
            else:
                case_results = getattr(search_report, "case_results", [])
            
            num_cases = len(case_results) if case_results else 0
            passed = sum(1 for r in case_results if getattr(r, "passed", False)) if case_results else 0
            failed = num_cases - passed
            
            print(f"  ✓ Search: {num_cases} cases")
            results["search"] = {
                "report": search_report,
                "num_cases": num_cases,
                "passed": passed,
                "failed": failed,
            }
        except Exception as e:
            print(f"  ✗ Search: {e}")
            results["search"] = {"error": str(e)}

    return {
        "summary": {
            "total": results.get("todo", {}).get("num_cases", 0) + results.get("search", {}).get("num_cases", 0),
        },
        "by_toolset": results,
    }


async def run_thinking(config: EvaluationConfig) -> dict[str, Any]:
    """Run evaluations for thinking/cognition toolsets."""
    print("Running evaluations for thinking/cognition toolsets...")
    return await compare_thinking_toolsets(config)


async def run_multi_agent(config: EvaluationConfig) -> dict[str, Any]:
    """Run evaluations for multi-agent toolsets."""
    print("Running evaluations for multi-agent toolsets...")
    return await compare_multi_agent_toolsets(config)


async def run_reflection(config: EvaluationConfig) -> dict[str, Any]:
    """Run evaluations for reflection toolsets."""
    print("Running evaluations for reflection toolsets...")
    return await compare_reflection_toolsets(config)


async def run_all(config: EvaluationConfig, sequential: bool = True) -> dict[str, Any]:
    """Run all evaluations.
    
    Args:
        config: Evaluation configuration.
        sequential: If True, run evaluations one by one. If False, run in parallel.
    
    Returns:
        Dictionary of results by category.
    """
    print("Running all evaluations...")
    if sequential:
        print("Running evaluations sequentially (one by one)...\n")
    else:
        print("Running evaluations in parallel...\n")
    
    results = {}

    if sequential:
        # Run one by one with clear separation
        print("=" * 60)
        print("CATEGORY 1/4: UNIQUES")
        print("=" * 60)
        results["uniques"] = await run_uniques(config)
        print("\n")

        print("=" * 60)
        print("CATEGORY 2/4: THINKING/COGNITION")
        print("=" * 60)
        results["thinking"] = await run_thinking(config)
        print("\n")

        print("=" * 60)
        print("CATEGORY 3/4: MULTI-AGENT")
        print("=" * 60)
        results["multi_agent"] = await run_multi_agent(config)
        print("\n")

        print("=" * 60)
        print("CATEGORY 4/4: REFLECTION")
        print("=" * 60)
        results["reflection"] = await run_reflection(config)
    else:
        # Run in parallel (original behavior)
        import asyncio
        uniques_task = asyncio.create_task(run_uniques(config))
        thinking_task = asyncio.create_task(run_thinking(config))
        multi_agent_task = asyncio.create_task(run_multi_agent(config))
        reflection_task = asyncio.create_task(run_reflection(config))
        
        results["uniques"] = await uniques_task
        results["thinking"] = await thinking_task
        results["multi_agent"] = await multi_agent_task
        results["reflection"] = await reflection_task

    return results


def export_results(results: dict[str, Any], output_dir: Path, format: str = "json"):
    """Export results to file.

    Args:
        results: Evaluation results.
        output_dir: Output directory.
        format: Export format ('json' or 'csv').
    """
    output_dir.mkdir(parents=True, exist_ok=True)

    if format == "json":
        output_file = output_dir / "results.json"
        with open(output_file, "w") as f:
            json.dump(results, f, indent=2, default=str)
        print(f"Results exported to {output_file}")
    elif format == "csv":
        # Simple CSV export for summary
        import csv

        output_file = output_dir / "summary.csv"
        with open(output_file, "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(["Category", "Toolset", "Metric", "Value"])

            for category, category_results in results.items():
                if "by_toolset" in category_results:
                    for toolset, toolset_data in category_results["by_toolset"].items():
                        if "avg_scores" in toolset_data:
                            for metric, value in toolset_data["avg_scores"].items():
                                writer.writerow([category, toolset, metric, value])

        print(f"Summary exported to {output_file}")


def print_summary(results: dict[str, Any]):
    """Print summary of results."""
    print("\n" + "=" * 60)
    print("EVALUATION SUMMARY")
    print("=" * 60)

    for category, category_results in results.items():
        print(f"\n{category.upper()}:")
        if "summary" in category_results:
            summary = category_results["summary"]
            print(f"  Total cases: {summary.get('total', 0)}")
            print(f"  Errors: {summary.get('errors', 0)}")
            print(f"  Success rate: {summary.get('success_rate', 0):.2%}")
            print(f"  Avg execution time: {summary.get('avg_execution_time', 0):.2f}s")
            print(f"  Total tokens: {summary.get('total_tokens', 0)}")

        if "by_toolset" in category_results:
            print("\n  By toolset:")
            for toolset, toolset_data in category_results["by_toolset"].items():
                print(f"    {toolset}:")
                
                # Handle error case
                if "error" in toolset_data:
                    print(f"      Error: {toolset_data['error']}")
                    continue
                
                # Handle pydantic-evals EvaluationReport (from compare functions)
                if "report" in toolset_data:
                    report = toolset_data["report"]
                    # Check what attributes the report has
                    if hasattr(report, "case_results"):
                        case_results = report.case_results
                    elif hasattr(report, "results"):
                        case_results = report.results
                    elif hasattr(report, "cases"):
                        case_results = report.cases
                    else:
                        case_results = []
                    
                    # Use direct values if available, otherwise extract from report
                    num_cases = toolset_data.get("num_cases", len(case_results) if case_results else 0)
                    passed = toolset_data.get("passed", sum(1 for r in case_results if getattr(r, "passed", False)) if case_results else 0)
                    failed = toolset_data.get("failed", num_cases - passed)
                    
                    print(f"      Cases: {num_cases}")
                    print(f"      Passed: {passed}")
                    print(f"      Failed: {failed}")
                    
                    # Get timing and token info from toolset_data or report
                    avg_time = toolset_data.get("avg_execution_time", 0.0)
                    total_tokens = toolset_data.get("total_tokens", 0)
                    
                    if avg_time == 0.0 and hasattr(report, "summary_stats") and report.summary_stats:
                        stats = report.summary_stats
                        avg_time = stats.get("avg_duration_ms", 0) / 1000.0
                        total_tokens = stats.get("total_tokens", 0)
                    
                    if avg_time > 0:
                        print(f"      Avg duration: {avg_time:.2f}s")
                    if total_tokens > 0:
                        print(f"      Total tokens: {total_tokens}")
                # Handle dict with aggregated data (legacy format)
                elif isinstance(toolset_data, dict):
                    if "avg_scores" in toolset_data:
                        for metric, value in toolset_data["avg_scores"].items():
                            print(f"      {metric}: {value:.3f}")
                    print(f"      Avg time: {toolset_data.get('avg_execution_time', 0):.2f}s")
                    print(f"      Total tokens: {toolset_data.get('total_tokens', 0)}")
                    print(f"      Errors: {toolset_data.get('errors', 0)}")
                    print(f"      Cases: {toolset_data.get('num_cases', 0)}")
                # Handle list of results
                elif isinstance(toolset_data, list):
                    if toolset_data:
                        avg_time = sum(r.execution_time for r in toolset_data) / len(toolset_data)
                        total_tokens = sum(r.token_usage.get("total_tokens", 0) for r in toolset_data)
                        errors = sum(1 for r in toolset_data if r.error)
                        
                        # Calculate average scores
                        score_keys = set()
                        for result in toolset_data:
                            score_keys.update(result.scores.keys())
                        
                        print(f"      Cases: {len(toolset_data)}")
                        print(f"      Errors: {errors}")
                        print(f"      Avg time: {avg_time:.2f}s")
                        print(f"      Total tokens: {total_tokens}")
                        if score_keys:
                            print(f"      Scores:")
                            for key in sorted(score_keys):
                                scores = [r.scores[key] for r in toolset_data if key in r.scores]
                                if scores:
                                    avg_score = sum(scores) / len(scores)
                                    print(f"        {key}: {avg_score:.3f}")


async def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(description="Run toolset evaluations")
    parser.add_argument(
        "--category",
        choices=["uniques", "thinking", "multi_agent", "reflection", "all"],
        default="all",
        help="Category to evaluate",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("eval_results"),
        help="Output directory for results",
    )
    parser.add_argument(
        "--format",
        choices=["json", "csv"],
        default="json",
        help="Export format",
    )
    parser.add_argument(
        "--api-key",
        type=str,
        help="OpenRouter API key (or set OPENROUTER_API_KEY env var)",
    )
    parser.add_argument(
        "--no-logfire",
        action="store_true",
        help="Disable Logfire integration",
    )
    parser.add_argument(
        "--parallel",
        action="store_true",
        help="Run evaluations in parallel instead of sequentially",
    )

    args = parser.parse_args()

    # Setup configuration
    config = EvaluationConfig(api_key=args.api_key)
    if args.no_logfire:
        config.logfire_enabled = False

    # Setup API key environment variable FIRST (before creating any agents)
    config.setup_environment()

    # Setup Logfire if enabled - MUST be configured BEFORE creating datasets
    # This is critical for experiments to show up in Logfire
    # According to docs: https://ai.pydantic.dev/evals/how-to/logfire-integration/
    if config.logfire_enabled and logfire:
        # Configure Logfire
        # Note: We disable scrubbing for evaluations since:
        # 1. We're not logging actual API keys (they're in env vars, not in logs)
        # 2. Model names and evaluation metadata are safe to log
        # 3. Scrubbing was causing false positives on model strings like "openrouter:x-ai/grok-4.1-fast"
        try:
            logfire.configure(
                send_to_logfire="if-token-present",  # Only send if LOGFIRE_TOKEN is set
                service_name="toolset-evaluations",
                environment="development",
                scrubbing=False,  # Disable scrubbing for evaluation runs
            )
        except TypeError:
            # Fallback if scrubbing parameter not supported in this version
            logfire.configure(
                send_to_logfire="if-token-present",
                service_name="toolset-evaluations",
                environment="development",
            )
            print("⚠ Note: Logfire scrubbing parameter not available - some data may be scrubbed")
        
        # Instrument pydantic-ai for full tracing
        logfire.instrument_pydantic_ai()
        print("✓ Logfire configured - experiments will appear in Logfire dashboard")
        print("  (Scrubbing disabled for evaluation data - API keys are not logged)")
    else:
        print("⚠ Logfire disabled - experiments will not appear in Logfire")

    # Run evaluations
    try:
        sequential = not args.parallel  # Default to sequential unless --parallel flag is set
        
        if args.category == "uniques":
            results = {"uniques": await run_uniques(config)}
        elif args.category == "thinking":
            results = {"thinking": await run_thinking(config)}
        elif args.category == "multi_agent":
            results = {"multi_agent": await run_multi_agent(config)}
        elif args.category == "reflection":
            results = {"reflection": await run_reflection(config)}
        else:
            results = await run_all(config, sequential=sequential)

        # Print summary
        print_summary(results)

        # Export results
        export_results(results, args.output_dir, args.format)

        print("\nEvaluation complete!")

    except KeyboardInterrupt:
        print("\nEvaluation interrupted by user")
        sys.exit(1)
    except Exception as e:
        print(f"\nError during evaluation: {e}")
        import traceback

        traceback.print_exc()
        sys.exit(1)


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

