"""bilge run - Generate and execute training scripts."""

import json
import os
import subprocess
import sys
import time
from pathlib import Path

import typer
from rich.console import Console
from rich.panel import Panel
from rich.syntax import Syntax

from bilgeai.analyzer.models import DatasetProfile
from bilgeai.core.constants import PROJECT_DIR
from bilgeai.core.exceptions import ConfigError
from bilgeai.memory.store import ExperimentStore
from bilgeai.optimizer.engine import HybridEngine
from bilgeai.profiler.models import HardwareProfile

console = Console()


def run_command(
    dry_run: bool = typer.Option(
        True,
        "--dry-run/--start",
        help="Generate only (default) or start training",
    ),
    output: Path = typer.Option(None, "--output", "-o", help="Output script path"),
    show: bool = typer.Option(True, "--show/--no-show", help="Show generated script"),
    seed: int = typer.Option(42, "--seed", help="Random seed for reproducibility"),
) -> None:
    """Generate optimized training script, optionally start training."""
    # Load config
    try:
        from bilgeai.core.config import BilgeConfig

        config = BilgeConfig.load()
        model_name = config.model.name
        dataset_path = config.data.path
        user_overrides = config.get_overrides()
        llm_config = {
            "provider": config.llm.provider,
            "model": config.llm.model,
            "api_key_env": config.llm.api_key_env,
            "base_url": config.llm.base_url,
        }
    except ConfigError:
        console.print("[red]No bilge.yaml found. Run 'bilge init' first.[/red]")
        raise typer.Exit(1) from None

    # API key check (mandatory for bilge run)
    _check_api_key(llm_config)

    # Show cost estimation
    _show_cost_estimate(llm_config)

    # Load hardware profile
    hw_path = Path.cwd() / PROJECT_DIR / "hardware_profile.json"
    if hw_path.exists():
        hw = HardwareProfile.model_validate(json.loads(hw_path.read_text(encoding="utf-8")))
    else:
        console.print("[yellow]No hardware profile. Running profiler...[/yellow]\n")
        from bilgeai.profiler import get_hardware_profile

        hw = get_hardware_profile()

    # Load dataset profile
    ds_path = Path.cwd() / PROJECT_DIR / "dataset_profile.json"
    if ds_path.exists():
        ds = DatasetProfile.model_validate(json.loads(ds_path.read_text(encoding="utf-8")))
        dataset_format = ds.format
    else:
        console.print(
            "[yellow]No dataset profile. Run 'bilge analyze' for better results.[/yellow]\n"
        )
        ds = DatasetProfile(path=dataset_path, format="unknown", num_rows=1000)
        dataset_format = "alpaca"

    # Optimize config
    engine = HybridEngine()
    training_config = engine.optimize(
        hw, ds, model_name=model_name, user_overrides=user_overrides, llm_config=llm_config
    )

    # Apply seed
    training_config.seed = seed

    # Generate script
    console.print("[bold]Generating training script...[/bold]\n")
    script = engine.generate_script(
        training_config,
        hardware=hw,
        dataset=ds,
        dataset_path=dataset_path,
        dataset_format=dataset_format,
        llm_config=llm_config,
    )

    # Determine output path
    if output is None:
        output = Path.cwd() / PROJECT_DIR / "train.py"

    # Save script
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(script, encoding="utf-8")

    # Extract dependencies from generated code
    from bilgeai.generator import extract_dependencies

    deps = extract_dependencies(script)

    # Save requirements.txt
    req_path = output.parent / "requirements.txt"
    req_path.write_text("\n".join(deps) + "\n", encoding="utf-8")

    # Save training config for 'bilge explain'
    config_save_path = Path.cwd() / PROJECT_DIR / "training_config.json"
    config_save_path.write_text(
        training_config.model_dump_json(indent=2, exclude_none=True),
        encoding="utf-8",
    )

    console.print(f"[green]Script saved to {output}[/green]")
    console.print(f"[green]Requirements saved to {req_path}[/green]\n")

    # Show script
    if show:
        syntax = Syntax(script, "python", theme="monokai", line_numbers=True)
        console.print(Panel(syntax, title="Generated Training Script", border_style="blue"))

    # Show config summary
    console.print("\n[bold]Config Summary:[/bold]")
    console.print(f"  Model:          {training_config.model_name}")
    console.print(f"  Quantization:   {training_config.quantization}")
    console.print(f"  Batch size:     {training_config.batch_size}")
    console.print(f"  Grad accum:     {training_config.gradient_accumulation_steps}")
    console.print(f"  Learning rate:  {training_config.learning_rate}")
    console.print(f"  Epochs:         {training_config.num_epochs}")
    console.print(f"  Seq length:     {training_config.max_seq_length}")
    if training_config.lora_enabled:
        console.print(f"  LoRA rank:      {training_config.lora_rank}")
    prec = "bf16" if training_config.bf16 else "fp16" if training_config.fp16 else "fp32"
    console.print(f"  Precision:      {prec}")
    console.print(f"  Seed:           {training_config.seed}")
    if training_config.num_gpus > 1:
        dist = training_config.distributed_strategy
        console.print(f"  GPUs:           {training_config.num_gpus} ({dist})")
    if training_config.num_nodes > 1:
        console.print(f"  Nodes:          {training_config.num_nodes}")

    # Show required dependencies
    console.print(f"\n[bold]Required Dependencies ({len(deps)}):[/bold]")
    console.print(f"  pip install {' '.join(deps)}")

    # Create experiment record
    store = ExperimentStore()
    run_record = store.create_run(
        config_dict=training_config.model_dump(exclude={"explanations"}),
        hardware_dict=hw.model_dump(),
        model_name=training_config.model_name,
    )
    console.print(f"\n[dim]Run ID: {run_record.run_id}[/dim]")

    if dry_run:
        console.print("\n[yellow]Dry run mode.[/yellow] Review the script above.")
        # Show launch command for multi-GPU/node setups
        if training_config.torchrun_command:
            console.print(f"\n[bold]Launch command:[/bold]\n  {training_config.torchrun_command}")
        console.print("To start training: [bold]bilge run --start[/bold]")
    else:
        console.print("\n[bold]Starting training...[/bold]\n")
        store.update_run(run_record.run_id, status="running")
        start_time = time.monotonic()
        success = _execute_training(output, training_config)
        elapsed = time.monotonic() - start_time
        store.update_run(
            run_record.run_id,
            status="completed" if success else "failed",
            training_time_seconds=elapsed,
            output_dir=str(output.parent),
        )

    from bilgeai.cli.hints import show_next_hint

    show_next_hint("run", console)


def _check_api_key(llm_config: dict) -> None:
    """Verify API key is available. Exit with guidance if not."""
    provider = llm_config.get("provider", "anthropic")

    # Local provider: no API key needed
    if provider == "local":
        return

    api_key_env = llm_config.get("api_key_env", "")
    if not api_key_env:
        console.print("[red]API key environment variable not configured.[/red]")
        console.print("[yellow]Run 'bilge init' to configure your LLM provider.[/yellow]")
        raise typer.Exit(1)

    api_key = os.environ.get(api_key_env)
    if not api_key:
        console.print(f"[red]API key not found: {api_key_env} is not set.[/red]")
        console.print(f"[yellow]Set your API key: export {api_key_env}=your-key-here[/yellow]")
        console.print("[yellow]Or run 'bilge init' to reconfigure.[/yellow]")
        raise typer.Exit(1)


def _show_cost_estimate(llm_config: dict) -> None:
    """Show estimated LLM cost before running."""
    from bilgeai.llm.cost_tracker import estimate_run_cost

    provider = llm_config.get("provider", "anthropic")
    model = llm_config.get("model", "")

    estimate = estimate_run_cost(provider, model)
    if estimate > 0:
        console.print(f"[dim]Estimated LLM cost: ~${estimate:.4f} ({model})[/dim]")
        console.print(
            "[dim]Note: Prices may not be up to date. "
            "Check your provider dashboard for exact costs.[/dim]\n"
        )
    else:
        console.print("[dim]Estimated cost: $0.00 (local model)[/dim]\n")


def _execute_training(script_path: Path, config) -> bool:
    """Execute the generated training script. Returns True on success."""
    cmd = [sys.executable, str(script_path)]

    # Multi-GPU / Multi-Node launch
    if config.num_gpus > 1 and config.distributed_strategy in ("ddp", "fsdp"):
        cmd = [
            sys.executable,
            "-m",
            "torch.distributed.run",
            f"--nproc_per_node={config.num_gpus}",
        ]
        if config.num_nodes > 1:
            cmd.extend(
                [
                    f"--nnodes={config.num_nodes}",
                    "--rdzv_backend=c10d",
                    f"--rdzv_endpoint={_get_master_addr()}:{_get_master_port()}",
                ]
            )
        cmd.append(str(script_path))
    elif config.num_nodes > 1:
        # Multi-node without multi-GPU per node
        cmd = [
            sys.executable,
            "-m",
            "torch.distributed.run",
            f"--nproc_per_node={max(config.num_gpus, 1)}",
            f"--nnodes={config.num_nodes}",
            "--rdzv_backend=c10d",
            f"--rdzv_endpoint={_get_master_addr()}:{_get_master_port()}",
            str(script_path),
        ]

    console.print(f"[dim]Running: {' '.join(cmd)}[/dim]\n")

    try:
        result = subprocess.run(cmd, check=False)
        if result.returncode == 0:
            console.print("\n[green bold]Training completed successfully![/green bold]")
            return True
        else:
            console.print(f"\n[red]Training failed with exit code {result.returncode}[/red]")
            console.print("[yellow]Hint: Use 'bilge doctor' to validate your setup.[/yellow]")
            raise typer.Exit(code=result.returncode)
    except FileNotFoundError:
        console.print("[red]Python not found. Check your environment.[/red]")
        raise typer.Exit(code=1) from None


def _build_launch_command(script_path: Path, config) -> str:
    """Build the torchrun launch command string for display."""
    parts = ["torchrun"]
    parts.append(f"--nproc_per_node={config.num_gpus}")
    if config.num_nodes > 1:
        parts.append(f"--nnodes={config.num_nodes}")
        parts.append("--rdzv_backend=c10d")
        parts.append("--rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT")
    parts.append(str(script_path))
    return " ".join(parts)


def _get_master_addr() -> str:
    """Get master node address for multi-node training."""
    return os.environ.get("MASTER_ADDR", "localhost")


def _get_master_port() -> str:
    """Get master node port for multi-node training."""
    return os.environ.get("MASTER_PORT", "29500")
