"""bilge diff - Compare two training runs."""

import json

import typer
from rich.console import Console
from rich.table import Table

from bilgeai.memory.store import ExperimentStore

console = Console()


def diff_command(
    run_a: str = typer.Argument(..., help="First run ID"),
    run_b: str = typer.Argument(..., help="Second run ID"),
    json_output: bool = typer.Option(False, "--json", help="Output as JSON"),
) -> None:
    """Compare two training runs side by side."""
    store = ExperimentStore()
    result = store.diff_runs(run_a, run_b)

    if "error" in result:
        console.print(f"[red]{result['error']}[/red]")
        console.print("[yellow]Hint: Use 'bilge history' to see valid run IDs.[/yellow]")
        raise typer.Exit(code=1)

    if json_output:
        console.print_json(json.dumps(result, indent=2))
        return

    # Header
    console.print("\n[bold]Comparing:[/bold]")
    console.print(f"  A: {result['run_a']['id']} ({result['run_a']['status']})")
    console.print(f"  B: {result['run_b']['id']} ({result['run_b']['status']})")

    if result["same_config"]:
        console.print("\n[green]Same config hash[/green] — configs are identical.")
    else:
        console.print("\n[yellow]Different config hashes[/yellow]")

    # Config diff
    if result["config_diff"]:
        table = Table(title="Config Differences")
        table.add_column("Parameter", style="cyan")
        table.add_column("Run A", style="red")
        table.add_column("Run B", style="green")

        for key, vals in sorted(result["config_diff"].items()):
            table.add_row(key, str(vals["run_a"]), str(vals["run_b"]))

        console.print(table)

    # Results diff
    if result["results_diff"]:
        table = Table(title="Results Differences")
        table.add_column("Metric", style="cyan")
        table.add_column("Run A", style="red")
        table.add_column("Run B", style="green")

        for key, vals in sorted(result["results_diff"].items()):
            a_val = f"{vals['run_a']}" if vals["run_a"] is not None else "-"
            b_val = f"{vals['run_b']}" if vals["run_b"] is not None else "-"
            table.add_row(key, a_val, b_val)

        console.print(table)

    if not result["config_diff"] and not result["results_diff"]:
        console.print("\n[green]Runs are identical in config and results.[/green]")
