import shutil
from datetime import datetime
from pathlib import Path
from typing import Any

import typer
from rich import get_console
from rich.panel import Panel
from rich.table import Table

from bench_af.cli.validate_model import validate_model_organism
from bench_af.util import (
    list_environments,
    list_model_organisms,
    load_environment,
    load_model_organism,
)

console = get_console()


def _get_results_dir(model_name: str, environment_name: str) -> Path:
    """
    Get the results directory for a specific model-environment combination.
    """
    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    path = Path(__file__).parent.parent / "results" / "validate_model" / "all" / timestamp / f"{model_name}_{environment_name}"
    path.mkdir(parents=True, exist_ok=True)
    return path


def _find_compatible_environments(model_organism) -> list[str]:
    """
    Find all environments that are compatible with the given model organism.
    """
    all_environments = list_environments()
    compatible_environments = []
    
    for env_name in all_environments:
        if env_name in model_organism.supported_environments:
            compatible_environments.append(env_name)
    
    return compatible_environments


def _run_validation_for_model_environment(
    model_name: str,
    environment_name: str,
    eval_config: dict[str, Any],
    results_dir: Path,
    head: int = 50,
) -> dict[str, Any]:
    """
    Run validation for a specific model-environment combination.
    """
    try:
        # Load model and environment
        model_organism = load_model_organism(model_name)
        environment = load_environment(environment_name)
        
        # Run validation
        validation_result = validate_model_organism(
            model_organism=model_organism,
            environment=environment,
            eval_config=eval_config,
            results_dir=results_dir,
            raw_eval_logs=True,
            head=head
        )
        
        return {
            "status": "success",
            "result": validation_result,
            "error": None
        }
    except Exception as e:
        return {
            "status": "error",
            "result": None,
            "error": str(e)
        }


def validate_all_models(
    eval_config: dict[str, Any] = None,
    results_base_dir: Path = None,
    head: int = 50,
) -> None:
    """
    🧪 Validate ALL models in ALL their compatible environments.
    
    This function:
    1. Discovers all available models and environments
    2. Finds compatible model-environment pairs
    3. Runs validation for each pair
    4. Reports results in a comprehensive summary
    """
    if eval_config is None:
        eval_config = {}
    
    if results_base_dir is None:
        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        results_base_dir = Path(__file__).parent.parent / "results" / "validate_model" / "all" / timestamp
        results_base_dir.mkdir(parents=True, exist_ok=True)
    
    console.print("🔍 [bold blue]Starting comprehensive model validation...[/bold blue]")
    
    # Discover all models and environments
    console.print("📋 Discovering available models and environments...", style="dim")
    all_models = list_model_organisms()
    all_environments = list_environments()
    
    console.print(f"📦 Found {len(all_models)} models: {', '.join(all_models)}")
    console.print(f"🌍 Found {len(all_environments)} environments: {', '.join(all_environments)}")
    
    # Find all compatible combinations
    compatible_combinations = []
    total_combinations = 0
    
    for model_name in all_models:
        try:
            model_organism = load_model_organism(model_name)
            compatible_envs = _find_compatible_environments(model_organism)
            
            for env_name in compatible_envs:
                compatible_combinations.append((model_name, env_name))
                total_combinations += 1
                
        except Exception as e:
            console.print(f"⚠️  [yellow]Error loading model {model_name}: {e}[/yellow]")
    
    console.print(f"🔗 Found {total_combinations} compatible model-environment combinations")
    # return
    
    if total_combinations == 0:
        console.print("❌ [bold red]No compatible model-environment combinations found![/bold red]")
        return
    
    # Run validations with simple progress tracking
    results = []
    successful_runs = 0
    failed_runs = 0
    
    console.print(f"🔄 Running {total_combinations} validations...")
    
    for i, (model_name, environment_name) in enumerate(compatible_combinations, 1):
        console.print(f"📋 [{i}/{total_combinations}] Validating {model_name} on {environment_name}...", style="dim")
        
        # Create results directory for this combination
        results_dir = results_base_dir / f"{model_name}_{environment_name}_compliance"
        results_dir.mkdir(parents=True, exist_ok=True)
        
        # Run validation
        result = _run_validation_for_model_environment(
            model_name=model_name,
            environment_name=environment_name,
            eval_config=eval_config,
            results_dir=results_dir,
            head=head,
        )
        
        results.append({
            "model": model_name,
            "environment": environment_name,
            "status": result["status"],
            "error": result["error"],
            "results_dir": results_dir,
        })
        
        if result["status"] == "success":
            successful_runs += 1
            console.print(f"  ✅ Success", style="green")
        else:
            failed_runs += 1
            console.print(f"  ❌ Failed: {result['error']}", style="red")
    
    # Display comprehensive results
    console.print("\n" + "="*80)
    console.print("📊 [bold blue]COMPREHENSIVE VALIDATION RESULTS[/bold blue]")
    console.print("="*80)
    
    # Summary statistics
    console.print(f"✅ Successful validations: {successful_runs}")
    console.print(f"❌ Failed validations: {failed_runs}")
    console.print(f"📁 Results stored in: {results_base_dir}")
    
    # Detailed results table
    if results:
        table = Table(title="Detailed Validation Results")
        table.add_column("Model", style="cyan")
        table.add_column("Environment", style="green")
        table.add_column("Status", style="bold")
        table.add_column("Error", style="red")
        
        for result in results:
            status_style = "green" if result["status"] == "success" else "red"
            status_icon = "✅" if result["status"] == "success" else "❌"
            
            table.add_row(
                result["model"],
                result["environment"],
                f"{status_icon} {result['status']}",
                result["error"] or "N/A"
            )
        
        console.print(table)
    
    # Save summary to file
    summary_file = results_base_dir / "validation_summary.txt"
    with open(summary_file, "w") as f:
        f.write("COMPREHENSIVE VALIDATION SUMMARY\n")
        f.write("="*50 + "\n\n")
        f.write(f"Total combinations tested: {total_combinations}\n")
        f.write(f"Successful validations: {successful_runs}\n")
        f.write(f"Failed validations: {failed_runs}\n\n")
        
        f.write("DETAILED RESULTS:\n")
        f.write("-"*20 + "\n")
        for result in results:
            f.write(f"{result['model']} on {result['environment']}: {result['status']}")
            if result["error"]:
                f.write(f" (Error: {result['error']})")
            f.write("\n")
    
    console.print(f"📄 Summary saved to: {summary_file}")
    
    if failed_runs > 0:
        console.print(f"\n⚠️  [yellow]{failed_runs} validations failed. Check the detailed results above.[/yellow]")
    
    console.print(f"\n✅ [bold green]Comprehensive validation completed![/bold green]") 