from __future__ import annotations

from pathlib import Path
from typing import Literal

import typer
from typing import Annotated
import rich
from rich.panel import Panel

def config_path_completions() -> list[str]:
    from bench_af.util import list_configs
    return list_configs()

console = rich.get_console()
run_app = typer.Typer(
    name="run",
    help="🚀 Execute validation and detection operations",
    rich_markup_mode="rich",
    no_args_is_help=True,
)


def _validate_config_path(config_path: Path) -> Path:
    """
    Validate that the config file exists and is readable.
    """
    path = Path(config_path)
    if not path.exists():
        console.print(
            f"[bold red]❌ Configuration file not found:[/bold red] {config_path}",
            style="red",
        )
        console.print(
            "💡 Please check the file path and try again.", style="dim")
        raise typer.Exit(1)

    if not path.is_file():
        console.print(
            f"[bold red]❌ Path is not a file:[/bold red] {config_path}", style="red"
        )
        raise typer.Exit(1)

    return path


def _get_results_dir(
    run_type: Literal["validate_model", "run_detector"],
    model_name: str,
    config_path: Path,
) -> Path:
    """
    Get the results directory for the run. The inspect-ai logs are stored in the logs subdirectory.
    """
    import shutil
    from datetime import datetime

    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    path = Path(__file__).parent.parent / "results" / \
        run_type / model_name / timestamp
    path.mkdir(parents=True, exist_ok=True)
    shutil.copy(config_path, path / "config.yaml")
    return path


@run_app.command(name="validate-model", no_args_is_help=True)
def cli_validate_model(
    config_name: Annotated[str, typer.Argument(autocompletion=config_path_completions)] = "",
    full_experiment: bool = typer.Option(
        False,
        "--full",
        "-f",
        help="""Run both the specified model and the baseline model, in both chain-of-thought and no chain-of-thought modes, in the specified environment.""",
        rich_help_panel="🔧 Optional Configuration",
    ),
    all_models: bool = typer.Option(
        False,
        "--all",
        "-a",
        help="""Validate ALL models in ALL their compatible environments.""",
        rich_help_panel="📋 Configuration",
    ),
    eval_config: str = typer.Option(
        "{}",
        "--eval-config",
        help="""Evaluation configuration parameters as a JSON string (used with --all).""",
        rich_help_panel="🔧 Optional Configuration",
    ),
    head: int = typer.Option(
        50,
        "--head",
        help="""Number of samples from the dataset to use for evaluation (used with --all).""",
        rich_help_panel="🔧 Optional Configuration",
    ),
) -> None:
    """
    🧪 Validate model organisms on environments.

    Pass the name of a config to validate a specific model using a config file.
    Use --all/-a to validate ALL models in ALL their compatible environments.
    """
    from bench_af._abstract.configs import ValidateModelConfig
    from bench_af.cli.validate_model import validate_model_organism
    from bench_af.cli.validate_all import validate_all_models
    from bench_af.util import get_bench_af_root, load_config, load_environment, load_model_organism
    
    if config_name and all_models:
        console.print("❌ [bold red]Cannot specify both a config name and --all/-a flags simultaneously.[/bold red]")
        raise typer.Exit(1)

    if full_experiment and all_models:
        console.print("❌ [bold red]Cannot use both --full and --all flags simultaneously.[/bold red]")
        raise typer.Exit(1)

    if config_name:
        config_path = get_bench_af_root() / "configs" / config_name / "validate_model.yaml"
        # Config-based validation
        console.print("🔍 [bold blue]Starting model validation...[/bold blue]")
        console.print(f"📄 Loading configuration from: {config_path}")

        # Validate config file exists
        config_file = _validate_config_path(config_path)

        # Load and validate configuration
        config = load_config(str(config_file), "validate-model")
        assert isinstance(config, ValidateModelConfig)

        if full_experiment and config.model_organism.name.startswith("baseline"):
            console.print("❌ [bold red]Cannot use --full with baseline model.[/bold red]")
            raise typer.Exit(1)

        if full_experiment and "no_cot" in config.model_organism.name:
            console.print(f"❌ [bold red]Cannot use --full with model organism that does not support chain-of-thought.[/bold red]\n\
                            Do you mean to use {config.model_organism.name.replace('_no_cot', '')}?")
            raise typer.Exit(1)

        # Load components
        console.print("📦 Loading model organism...", style="dim")
        loaded_model_organism = load_model_organism(config.model_organism.name)

        if full_experiment:
            console.print("📦 Loading model organism without chain-of-thought...", style="dim")
            loaded_model_organism_no_cot = load_model_organism(config.model_organism.name + "_no_cot")
            console.print("📦 Loading baseline model organism...", style="dim")
            loaded_baseline = load_model_organism("baseline")
            console.print("📦 Loading baseline model organism without chain-of-thought...", style="dim")
            loaded__baseline_no_cot = load_model_organism("baseline_no_cot")

        console.print("🌍 Loading environment...", style="dim")

        # Run validation
        console.print("🧪 Running validation...", style="bold blue")
        results_dir = _get_results_dir(
            "validate_model", config.model_organism.name, config_path
        )
        console.print(f"📁 Results will be stored in: {results_dir}")
        print(config.model_organism)

        for env in config.environments:
            loaded_environment = load_environment(env.name)
            print(env.eval_config.get("head", 50))
            # breakpoint()
            validation_result = validate_model_organism(
                model_organism=loaded_model_organism,
                environment=loaded_environment,
                eval_config=env.eval_config,
                results_dir=results_dir / f"{config.model_organism.name}_{env.name}_compliance",
                raw_eval_logs=True,
                head=env.eval_config.get("head", 50),
            )

            if full_experiment:
                validation_result_no_cot = validate_model_organism(
                    model_organism=loaded_model_organism_no_cot,
                    environment=loaded_environment,
                    eval_config=env.eval_config,
                    results_dir=results_dir / f"{config.model_organism.name+'_no_cot'}_{env.name}_compliance",
                    raw_eval_logs=True,
                    head=env.eval_config.get("head", 50),
                )

                validation_result_baseline = validate_model_organism(
                    model_organism=loaded_baseline,
                    environment=loaded_environment,
                    eval_config=env.eval_config,
                    results_dir=results_dir / f"baseline_{env.name}_compliance",
                    raw_eval_logs=True,
                    head=env.eval_config.get("head", 50),
                )

                validation_result_baseline_no_cot = validate_model_organism(
                    model_organism=loaded__baseline_no_cot,
                    environment=loaded_environment,
                    eval_config=env.eval_config,
                    results_dir=results_dir / f"baseline_no_cot_{env.name}_compliance",
                    raw_eval_logs=True,
                    head=env.eval_config.get("head", 50),
                )

            # Display results
            console.print(
                "\n✅ [bold green]Validation completed successfully![/bold green]"
            )
            console.print(
                Panel(
                    str(validation_result),
                    title="📊 Validation Results for Model Organism",
                    border_style="green",
                )
            )
            if full_experiment:
                console.print(
                    Panel(
                        str(validation_result_no_cot),
                        title="📊 Validation Results for Model Organism Without Chain-of-Thought",
                        border_style="green",
                    )
                )
                console.print(
                    Panel(
                        str(validation_result_baseline),
                        title="📊 Validation Results for Baseline Model Organism",
                        border_style="green",
                    )
                )
                console.print(
                    Panel(
                        str(validation_result_baseline_no_cot),
                        title="📊 Validation Results for Baseline Model Organism Without Chain-of-Thought",
                        border_style="green",
                    )
                )
    
    elif all_models:
        # All models validation
        console.print(
            "🔍 [bold blue]Starting comprehensive model validation...[/bold blue]")

        # Parse the eval_config JSON string
        import json
        try:
            parsed_eval_config = json.loads(eval_config)
        except json.JSONDecodeError as e:
            console.print(
                f"❌ [bold red]Invalid JSON in eval-config: {e}[/bold red]")
            raise typer.Exit(1)

        # Run the comprehensive validation
        validate_all_models(eval_config=parsed_eval_config, head=head)




@run_app.command(name="run-detector", no_args_is_help=True)
def cli_run_detector(config_name: Annotated[str, typer.Argument(autocompletion=config_path_completions)]) -> None:
    """
    🔍 Run a detector algorithm on a model organism and environment.
    """
    from bench_af._abstract.configs import RunDetectorConfig
    from bench_af.cli.run_detector import run_detector
    from bench_af.util import get_bench_af_root, load_config, load_detector, load_environment, load_model_organism
    
    console.print("🔍 [bold blue]Starting detector analysis...[/bold blue]")
    config_path = get_bench_af_root() / "configs" / config_name / "run_detector.yaml"
    console.print(f"📄 Loading configuration from: {config_path}")

    # Validate config file exists
    config_file = _validate_config_path(config_path)

    # Load and validate configuration
    config = load_config(str(config_file), "run-detector")
    assert isinstance(config, RunDetectorConfig)

    # Load components
    console.print(f"🔍 Loading detector {config.detector.name}...", style="dim")
    loaded_detector = load_detector(config.detector.name)

    console.print(
        f"📦 Loading model organism {config.model_organism.name}...", style="dim"
    )
    loaded_model_organism = load_model_organism(config.model_organism.name)

    console.print(
        f"🌍 Loading environment {config.environments[0].name}...", style="dim"
    )
    # TODO: Support multiple MODELS
    loaded_environment = load_environment(config.environments[0].name)

    # Run detector
    console.print("🔍 Running detector analysis...", style="bold blue")
    results_dir = _get_results_dir(
        "run_detector", config.model_organism.name, config_path
    )
    console.print(f"📁 Results will be stored in: {results_dir}")
    console.print(
        f"📁 Inspect-AI logs will be stored in: {results_dir / 'logs'}")

    detector_result = run_detector(
        detector=loaded_detector,
        model_organism=loaded_model_organism,
        environment=loaded_environment,
        # TODO: Support multiple MODELS
        eval_config=config.environments[0].eval_config,
        results_dir=results_dir,
    )

    # Display results
    console.print("\n✅ [bold green]Detector analysis completed![/bold green]")
    console.print(
        Panel(str(detector_result),
              title="🔍 Detection Results", border_style="green")
    )
