"""bilge replay - Replay a previous training run with the same config."""

import json
from pathlib import Path

import typer
from rich.console import Console

from bilgeai.memory.store import ExperimentStore

console = Console()


def replay_command(
    run_id: str = typer.Argument(None, help="Run ID to replay (latest if omitted)"),
    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"),
) -> None:
    """Replay a previous run using its stored configuration."""
    store = ExperimentStore()

    if run_id:
        record = store.get_run(run_id)
    else:
        runs = store.list_runs(limit=1)
        record = runs[0] if runs else None

    if record is None:
        console.print("[red]Run not found.[/red]")
        console.print("[yellow]Hint: Use 'bilge history' to list available runs.[/yellow]")
        raise typer.Exit(code=1)

    console.print(f"[bold]Replaying run:[/bold] {record.run_id}")
    console.print(f"  Model:    {record.model_name}")
    console.print(f"  Status:   {record.status}")
    console.print(f"  Created:  {record.timestamp}")
    console.print()

    # Reconstruct config from stored JSON
    config_dict = json.loads(record.config_json)
    hardware_dict = json.loads(record.hardware_json)

    from bilgeai.optimizer.models import TrainingConfig
    from bilgeai.profiler.models import HardwareProfile

    training_config = TrainingConfig(**config_dict)

    # Reconstruct hardware profile
    if hardware_dict:
        hw = HardwareProfile.model_validate(hardware_dict)
    else:
        console.print("[yellow]No hardware info stored. Running profiler...[/yellow]")
        from bilgeai.profiler import get_hardware_profile

        hw = get_hardware_profile()

    # Determine dataset info from stored config
    dataset_path = config_dict.get("dataset_path", "./train.jsonl")
    dataset_format = config_dict.get("dataset_format", "alpaca")

    # Generate script using stored config (no LLM, pure replay)
    from bilgeai.generator import CodeGenerator

    generator = CodeGenerator()
    hw_summary = f"{hw.cpu_name}, {hw.ram_total_gb:.0f}GB RAM"
    if hw.gpus:
        gpu = hw.gpus[0]
        hw_summary += f", {gpu.name} ({gpu.vram_gb:.0f}GB)"

    script = generator.generate_sft(
        training_config,
        dataset_path=dataset_path,
        dataset_format=dataset_format,
        hardware_summary=hw_summary,
    )

    # Save script
    from bilgeai.core.constants import PROJECT_DIR

    if output is None:
        output = Path.cwd() / PROJECT_DIR / "train_replay.py"

    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(script, encoding="utf-8")
    console.print(f"[green]Replay script saved to {output}[/green]\n")

    if show:
        from rich.panel import Panel
        from rich.syntax import Syntax

        syntax = Syntax(script, "python", theme="monokai", line_numbers=True)
        console.print(Panel(syntax, title=f"Replay: {record.run_id}", border_style="blue"))

    # Create new run record linked to original
    new_record = store.create_run(
        config_dict=config_dict,
        hardware_dict=hardware_dict,
        model_name=training_config.model_name,
        notes=f"Replay of {record.run_id}",
    )
    console.print(f"[dim]New Run ID: {new_record.run_id}[/dim]")

    if dry_run:
        console.print("\n[yellow]Dry run mode.[/yellow] Review the script above.")
        console.print(f"To start training: [bold]bilge replay {record.run_id} --start[/bold]")
    else:
        import subprocess
        import sys
        import time

        console.print("\n[bold]Starting replayed training...[/bold]\n")
        store.update_run(new_record.run_id, status="running")
        start_time = time.monotonic()

        cmd = [sys.executable, str(output)]
        if training_config.num_gpus > 1 and training_config.distributed_strategy in ("ddp", "fsdp"):
            cmd = [
                sys.executable,
                "-m",
                "torch.distributed.run",
                f"--nproc_per_node={training_config.num_gpus}",
                str(output),
            ]

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

        try:
            result = subprocess.run(cmd, check=False)
            elapsed = time.monotonic() - start_time
            success = result.returncode == 0
            store.update_run(
                new_record.run_id,
                status="completed" if success else "failed",
                training_time_seconds=elapsed,
                output_dir=str(output.parent),
            )
            if success:
                console.print("\n[green bold]Replay training completed![/green bold]")
            else:
                console.print(f"\n[red]Training failed with exit code {result.returncode}[/red]")
                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
