#!/usr/bin/env python3

import os
import json
import logging
from typing import Dict, List, Any
from datetime import datetime

try:
    from openai import OpenAI
except ImportError:
    print("Error: openai package not found. Install with: pip install openai")
    exit(1)

from env.environment import CodeGuardianEnv
from env.models import AgentState
from inference import setup_openai_client, run_episode

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def run_benchmark(num_episodes: int = 10) -> Dict[str, Any]:
    """Run benchmarking across multiple episodes and return results."""
    try:
        # Setup
        client, model_name = setup_openai_client()
        agent_state = AgentState()
        episode_results = []
        
        logger.info(f"Starting benchmark with {num_episodes} episodes using model: {model_name}")
        print(f"\n{'='*60}")
        print(f"CODEGUARDIAN BENCHMARK - {model_name}")
        print(f"{'='*60}")
        print(f"{'Episode':<10} {'Score':<10} {'Steps':<10} {'Performance'}")
        print("-" * 50)
        
        for episode in range(num_episodes):
            env = CodeGuardianEnv()
            episode_reward, episode_history = run_episode(env, client, model_name, agent_state)
            
            # Calculate performance metrics
            steps = len(episode_history)
            performance = "Excellent" if episode_reward >= 0.8 else "Good" if episode_reward >= 0.5 else "Poor"
            
            episode_results.append({
                "episode": episode + 1,
                "reward": episode_reward,
                "steps": steps,
                "history": episode_history
            })
            
            # Print table row
            print(f"{episode + 1:<10} {episode_reward:<10.2f} {steps:<10} {performance}")
            
            # Update agent state for learning
            agent_state.episode += 1
            agent_state.improvement_history.append(episode_reward)
            if episode_reward > agent_state.best_score:
                agent_state.best_score = episode_reward
        
        # Calculate aggregate metrics
        rewards = [ep["reward"] for ep in episode_results]
        avg_score = sum(rewards) / len(rewards)
        best_score = max(rewards)
        worst_score = min(rewards)
        
        print("-" * 50)
        print(f"\n{'='*60}")
        print("BENCHMARK SUMMARY")
        print(f"{'='*60}")
        print(f"Model: {model_name}")
        print(f"Episodes: {num_episodes}")
        print(f"Average Score: {avg_score:.3f}")
        print(f"Best Score: {best_score:.3f}")
        print(f"Worst Score: {worst_score:.3f}")
        print(f"Score Std Dev: {(sum((r - avg_score)**2 for r in rewards) / len(rewards))**0.5:.3f}")
        print(f"Improvement Trend: {'Upward' if rewards[-1] > rewards[0] else 'Downward' if rewards[-1] < rewards[0] else 'Stable'}")
        
        # Create benchmark result object
        benchmark_result = {
            "model": model_name,
            "timestamp": datetime.now().isoformat(),
            "episodes": num_episodes,
            "avg_score": avg_score,
            "best_score": best_score,
            "worst_score": worst_score,
            "score_std": (sum((r - avg_score)**2 for r in rewards) / len(rewards))**0.5,
            "improvement_trend": "upward" if rewards[-1] > rewards[0] else "downward" if rewards[-1] < rewards[0] else "stable",
            "episode_results": episode_results
        }
        
        # Save benchmark results
        results_file = f"benchmark_results_{model_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        with open(results_file, 'w') as f:
            json.dump(benchmark_result, f, indent=2)
        
        logger.info(f"Benchmark results saved to: {results_file}")
        print(f"\nDetailed results saved to: {results_file}")
        
        return benchmark_result
        
    except Exception as e:
        logger.error(f"Benchmark failed: {e}")
        raise

def compare_benchmarks(results_files: List[str]) -> None:
    """Compare multiple benchmark result files."""
    print(f"\n{'='*60}")
    print("BENCHMARK COMPARISON")
    print(f"{'='*60}")
    
    benchmarks = []
    for file in results_files:
        try:
            with open(file, 'r') as f:
                benchmarks.append(json.load(f))
        except Exception as e:
            logger.warning(f"Failed to load {file}: {e}")
    
    if not benchmarks:
        print("No valid benchmark files found")
        return
    
    print(f"{'Model':<20} {'Avg Score':<12} {'Best Score':<12} {'Episodes':<10}")
    print("-" * 60)
    
    for bench in benchmarks:
        print(f"{bench['model']:<20} {bench['avg_score']:<12.3f} {bench['best_score']:<12.3f} {bench['episodes']:<10}")
    
    # Find best performing model
    best_model = max(benchmarks, key=lambda x: x['avg_score'])
    print(f"\n🏆 Best Performing Model: {best_model['model']} (Avg: {best_model['avg_score']:.3f})")

def main():
    """Main benchmark execution."""
    import sys
    
    # Parse command line arguments
    num_episodes = 10
    if len(sys.argv) > 1:
        try:
            num_episodes = int(sys.argv[1])
        except ValueError:
            print("Usage: python benchmark.py [num_episodes]")
            return 1
    
    try:
        # Run benchmark
        result = run_benchmark(num_episodes)
        
        # Check for comparison mode
        if len(sys.argv) > 2 and sys.argv[2] == "--compare":
            # Find existing benchmark files
            import glob
            benchmark_files = glob.glob("benchmark_results_*.json")
            if benchmark_files:
                compare_benchmarks(benchmark_files)
        
        return 0
        
    except Exception as e:
        logger.error(f"Benchmark execution failed: {e}")
        return 1

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