"""bilge chart - Generate loss curve and metric charts from training logs."""

import json
from pathlib import Path

import typer
from rich.console import Console

from bilgeai.core.constants import PROJECT_DIR

console = Console()


def chart_command(
    log_dir: Path = typer.Argument(
        None, help="Training output directory (default: .bilge/output/)"
    ),
    output: Path = typer.Option(None, "--output", "-o", help="Save chart to file (PNG/PDF)"),
    metric: str = typer.Option(
        "loss", "--metric", "-m", help="Metric to plot: loss, lr, grad_norm"
    ),
) -> None:
    """Generate training metric charts from trainer logs."""
    # Find log file
    if log_dir is None:
        log_dir = Path.cwd() / PROJECT_DIR

    log_file = _find_log_file(log_dir)
    if log_file is None:
        console.print("[red]No training log found.[/red]")
        console.print(
            "[yellow]Hint: Run 'bilge run --start' first to generate training logs.[/yellow]"
        )
        raise typer.Exit(code=1)

    console.print(f"[bold]Reading logs from:[/bold] {log_file}\n")

    # Parse log entries
    entries = _parse_trainer_state(log_file)
    if not entries:
        console.print("[red]No log entries found in trainer state.[/red]")
        raise typer.Exit(code=1)

    # Extract metric data
    steps = []
    values = []
    for entry in entries:
        if metric in entry:
            steps.append(entry.get("step", 0))
            values.append(entry[metric])

    if not values:
        console.print(f"[red]No '{metric}' data found in logs.[/red]")
        console.print(f"[dim]Available metrics: {_available_metrics(entries)}[/dim]")
        raise typer.Exit(code=1)

    console.print(f"  Metric:     {metric}")
    console.print(f"  Data points: {len(values)}")
    console.print(f"  Min value:  {min(values):.6f}")
    console.print(f"  Max value:  {max(values):.6f}")
    console.print(f"  Final value: {values[-1]:.6f}")
    console.print()

    # Try matplotlib
    try:
        _plot_chart(steps, values, metric, output, log_file)
    except ImportError:
        console.print("[yellow]matplotlib not installed. Showing ASCII chart.[/yellow]\n")
        _ascii_chart(steps, values, metric)
        if output:
            console.print(
                "\n[dim]Install matplotlib for PNG/PDF export: pip install matplotlib[/dim]"
            )


def _find_log_file(base_dir: Path) -> Path | None:
    """Find trainer_state.json in output directories."""
    # Search common locations
    candidates = list(base_dir.rglob("trainer_state.json"))
    if candidates:
        # Return the most recently modified one
        return max(candidates, key=lambda p: p.stat().st_mtime)
    return None


def _parse_trainer_state(log_file: Path) -> list[dict]:
    """Parse HuggingFace trainer_state.json log entries."""
    data = json.loads(log_file.read_text(encoding="utf-8"))
    return data.get("log_history", [])


def _available_metrics(entries: list[dict]) -> str:
    """List available metric names from log entries."""
    keys = set()
    for entry in entries:
        keys.update(k for k in entry if k not in ("step", "epoch"))
    return ", ".join(sorted(keys)) if keys else "none"


def _plot_chart(
    steps: list[int],
    values: list[float],
    metric: str,
    output: Path | None,
    log_file: Path,
) -> None:
    """Generate chart with matplotlib."""
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots(figsize=(10, 6))
    ax.plot(steps, values, linewidth=1.5, color="#2196F3")
    ax.set_xlabel("Step")
    ax.set_ylabel(metric.replace("_", " ").title())
    ax.set_title(f"Training {metric.replace('_', ' ').title()} Curve")
    ax.grid(True, alpha=0.3)

    # Add min marker
    min_idx = values.index(min(values))
    ax.scatter([steps[min_idx]], [values[min_idx]], color="#4CAF50", zorder=5, s=50)
    ax.annotate(
        f"min: {values[min_idx]:.4f}",
        xy=(steps[min_idx], values[min_idx]),
        xytext=(10, 10),
        textcoords="offset points",
        fontsize=9,
        color="#4CAF50",
    )

    fig.tight_layout()

    if output:
        output.parent.mkdir(parents=True, exist_ok=True)
        fig.savefig(str(output), dpi=150)
        console.print(f"[green]Chart saved to {output}[/green]")
    else:
        # Save to default location
        default_path = log_file.parent / f"{metric}_curve.png"
        fig.savefig(str(default_path), dpi=150)
        console.print(f"[green]Chart saved to {default_path}[/green]")

    plt.close(fig)


def _ascii_chart(steps: list[int], values: list[float], metric: str) -> None:
    """Render a simple ASCII chart as fallback."""
    height = 15
    width = min(60, len(values))

    # Downsample if needed
    if len(values) > width:
        step_size = len(values) / width
        sampled = [values[int(i * step_size)] for i in range(width)]
    else:
        sampled = values

    min_val = min(sampled)
    max_val = max(sampled)
    val_range = max_val - min_val or 1.0

    console.print(f"  {metric} curve ({len(values)} steps)")
    console.print(f"  {max_val:.4f} |")

    for row in range(height, 0, -1):
        threshold = min_val + (row / height) * val_range
        line = ""
        for v in sampled:
            line += "#" if v >= threshold else " "
        if row == height:
            console.print(f"           |{line}|")
        elif row == 1:
            console.print(f"  {min_val:.4f} |{line}|")
        else:
            console.print(f"           |{line}|")

    console.print(f"           +{'-' * len(sampled)}+")
    console.print(f"           0{' ' * (len(sampled) - len(str(steps[-1])))}{steps[-1]}")
