import json as json_lib
import os

import psutil
import typer
from rich.console import Console
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
from rich.table import Table

from .engine import LLAMA_CPP_AVAILABLE, StressTester
from .probe import NVML_AVAILABLE, HardwareProbe

app = typer.Typer(
    help="⚡ stressllm — Find the breaking point of your local LLM hardware.",
    no_args_is_help=True,
    rich_markup_mode="rich",
)
console = Console()


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _check_ollama() -> bool:
    """Return True if Ollama is reachable, else print error and return False."""
    try:
        import ollama
        ollama.list()
        return True
    except Exception:
        console.print(Panel(
            "[bold red]✗ Ollama is not running or not installed.[/bold red]\n\n"
            "  Install it from: [link=https://ollama.com/download]https://ollama.com/download[/link]\n"
            "  Then start it with: [cyan]ollama serve[/cyan]",
            title="Connection Error",
            border_style="red",
        ))
        return False


def _check_model(model: str) -> bool:
    """Return True if the model is available in Ollama, else print error."""
    try:
        import ollama
        response = ollama.list()
        available = [m.model for m in response.models] if hasattr(response, 'models') else []
        # Exact match only — user must provide the full name including tag
        if model in available:
            return True
        # Also accept bare name if there's exactly one match
        # (e.g. "gemma" when only "gemma:2b" exists)
        bare_matches = [n for n in available if n.split(":")[0] == model]
        if len(bare_matches) == 1:
            return True
        if len(bare_matches) > 1:
            _error_panel(
                "Ambiguous Model Name",
                f"✗ '{model}' matches multiple models.",
                "Ollama has several variants with that name.",
                f"Be specific. Available: {', '.join(bare_matches)}",
            )
            return False
        _error_panel(
            "Model Not Found",
            f"✗ Model '{model}' is not available in Ollama.",
            "The model hasn't been pulled yet.",
            f"Pull it with: [cyan]ollama pull {model}[/cyan]\n"
            f"  Available models: {', '.join(available) if available else 'none'}",
        )
        return False
    except Exception:
        # If we can't list models, let the benchmark try anyway
        return True


def _error_panel(title: str, what: str, why: str, fix: str):
    """Print a structured Rich error panel."""
    console.print(Panel(
        f"[bold red]{what}[/bold red]\n\n"
        f"  [dim]Why:[/dim] {why}\n"
        f"  [dim]Fix:[/dim] {fix}",
        title=title,
        border_style="red",
    ))


def _print_results(results: list[dict], *, gpu_enabled: bool, as_json: bool = False):
    """Render the results table and verdict panel, or JSON."""
    if as_json:
        console.print_json(json_lib.dumps(results, indent=2))
        return

    table = Table(title="Benchmark Results", title_style="bold")
    table.add_column("Context", justify="center", style="cyan")
    table.add_column("TPS", style="green")
    table.add_column("TTFT", style="blue")
    table.add_column("Elapsed", style="dim")

    if gpu_enabled:
        table.add_column("VRAM", style="yellow")
        table.add_column("GPU Temp", style="red")

    table.add_column("RAM", style="dim")
    table.add_column("CPU%", style="dim")
    table.add_column("Status", justify="center")

    for r in results:
        row = [r["context"], str(r["tps"]), r["ttft"], r.get("elapsed", "—")]
        if gpu_enabled:
            row += [r["vram"], r["gpu_temp"]]
        row += [r["ram"], r["cpu"], r["status"]]
        table.add_row(*row)

    console.print(table)
    _print_verdict(results, gpu_enabled=gpu_enabled)


def _print_verdict(results: list[dict], *, gpu_enabled: bool):
    """Analyze results and print a verdict panel."""
    smooth_ctx = None
    slowing_ctx = None
    cliff_ctx = None

    for r in results:
        tps = r["tps"]
        if isinstance(tps, (int, float)):
            if tps > 15:
                smooth_ctx = r["context"]
            elif tps > 5 and slowing_ctx is None:
                slowing_ctx = r["context"]
            elif tps <= 5 and cliff_ctx is None:
                cliff_ctx = r["context"]

    lines = []
    if not gpu_enabled:
        if not NVML_AVAILABLE:
            lines.append(
                "[yellow]⚠  NVIDIA GPU not detected.[/yellow] "
                "Install GPU support: [bold]pip install \"stressllm\\[gpu]\"[/bold]"
            )
        else:
            lines.append(
                "[yellow]⚠  NVIDIA GPU not detected.[/yellow] "
                "pynvml is installed but could not initialize — check your NVIDIA drivers."
            )

    if smooth_ctx:
        lines.append(
            f"Verdict: model runs well up to [bold green]{smooth_ctx}[/bold green] context."
        )
    if slowing_ctx:
        lines.append(f"Usable but slowing at [bold yellow]{slowing_ctx}[/bold yellow] context.")
    if cliff_ctx:
        lines.append(f"Performance cliff detected at [bold red]{cliff_ctx}[/bold red].")
    if not smooth_ctx and not slowing_ctx and not cliff_ctx:
        lines.append("Could not determine performance profile — check errors above.")

    lines.append("")
    lines.append(
        "[dim]Note: Results may vary ±20% between runs due to thermal"
        " throttling, background system load, and GPU cache state."
        " Run multiple times for confidence.[/dim]"
    )

    border = "green" if smooth_ctx else ("yellow" if slowing_ctx else "red")
    console.print(Panel(
        "\n".join(lines),
        title="Verdict",
        border_style=border,
    ))


def _print_verbose(summary: dict):
    """Print hardware and dependency details before a benchmark."""
    console.print(Panel(
        f"[cyan]GPU:[/cyan]      {summary['gpu_name']}\n"
        f"[cyan]VRAM:[/cyan]     {summary['vram_total']}\n"
        f"[cyan]RAM:[/cyan]      {summary['ram_total']}\n"
        f"[cyan]CPU:[/cyan]      {psutil.cpu_count(logical=True)} logical cores\n"
        f"[cyan]pynvml:[/cyan]   {'✓' if NVML_AVAILABLE else '✗'}\n"
        f"[cyan]llama-cpp:[/cyan] {'✓' if LLAMA_CPP_AVAILABLE else '✗'}",
        title="Debug Info",
        border_style="dim",
    ))


# ---------------------------------------------------------------------------
# Commands
# ---------------------------------------------------------------------------

@app.command(
    epilog="[dim]Example: stressllm run gemma2 --depth 4 --verbose[/dim]",
)
def run(
    model: str = typer.Argument(
        ...,
        help="Ollama model name (e.g. gemma2, llama3).",
    ),
    depth: int = typer.Option(
        3, min=1, max=5,
        help="Context steps (1–5). Higher = longer test, larger contexts.",
        rich_help_panel="Benchmark Options",
    ),
    timeout: int = typer.Option(
        300,
        help="Max seconds per context step. 0 = no limit.",
        rich_help_panel="Benchmark Options",
    ),
    verbose: bool = typer.Option(
        False, "--verbose", "-v",
        help="Show detected hardware and dependency info before the test.",
        rich_help_panel="Output Options",
    ),
    json: bool = typer.Option(
        False, "--json",
        help="Output results as JSON (for scripting / CI).",
        rich_help_panel="Output Options",
    ),
):
    """Stress test a model via Ollama API."""
    if not _check_ollama():
        raise typer.Exit(code=1)

    if not _check_model(model):
        raise typer.Exit(code=1)

    tester = StressTester(model)
    summary = tester.probe.get_system_summary()

    if verbose:
        _print_verbose(summary)

    step_timeout = timeout if timeout > 0 else None

    # Hardware banner
    if not json:
        hw_line = f"RAM: {summary['ram_total']}"
        if summary["gpu_enabled"]:
            hw_line = f"{summary['gpu_name']} · {summary['vram_total']} VRAM · {hw_line}"
        else:
            hw_line = f"CPU-only mode · {hw_line}"

        console.print(Panel(
            f"[bold]⚡ stressllm — Stress Testing:[/bold] {model}\n"
            f"[dim]{hw_line}[/dim]",
            border_style="bright_blue",
        ))

    steps = [2048 * (4 ** i) for i in range(depth)]
    results = []
    try:
        if json:
            for size in steps:
                row = next(tester.iter_ollama_test_single(size, timeout=step_timeout))
                results.append(row)
                if "Timeout" in row["status"] or "OOM" in row["status"]:
                    break
        else:
            with Progress(
                SpinnerColumn(),
                TextColumn("[bold yellow]{task.description}[/bold yellow]"),
                TimeElapsedColumn(),
                console=console,
                transient=True,
            ) as progress:
                for i, size in enumerate(steps):
                    ctx_label = f"{size // 1024}k"
                    task = progress.add_task(
                        f"Testing {ctx_label} context… ({i + 1}/{len(steps)})",
                        total=None,
                    )
                    row = next(tester.iter_ollama_test_single(size, timeout=step_timeout))
                    results.append(row)
                    progress.remove_task(task)
                    console.print(
                        f"  [dim]✓[/dim] {ctx_label} — "
                        f"TPS: [green]{row['tps']}[/green]  "
                        f"TTFT: [blue]{row['ttft']}[/blue]  "
                        f"{row['status']}"
                    )
                    if "Timeout" in row["status"] or "OOM" in row["status"]:
                        console.print(
                            "[yellow]  Stopping — larger contexts will also fail.[/yellow]"
                        )
                        break
    except KeyboardInterrupt:
        if not json:
            console.print("\n[yellow]Interrupted — showing partial results.[/yellow]")
    except StopIteration:
        pass
    except Exception as e:
        _handle_ollama_error(e, model)
        if not results:
            raise typer.Exit(code=1)

    if results:
        _print_results(results, gpu_enabled=summary["gpu_enabled"], as_json=json)


@app.command(
    epilog="[dim]Example: stressllm check ./model.gguf --n-gpu -1 --depth 3[/dim]",
)
def check(
    path: str = typer.Argument(
        ...,
        help="Path to a .gguf model file.",
    ),
    n_gpu: int = typer.Option(
        -1,
        help="GPU layers to offload (-1 = all).",
        rich_help_panel="Benchmark Options",
    ),
    depth: int = typer.Option(
        3, min=1, max=5,
        help="Context steps (1–5).",
        rich_help_panel="Benchmark Options",
    ),
    verbose: bool = typer.Option(
        False, "--verbose", "-v",
        help="Show detected hardware and dependency info before the test.",
        rich_help_panel="Output Options",
    ),
    json: bool = typer.Option(
        False, "--json",
        help="Output results as JSON (for scripting / CI).",
        rich_help_panel="Output Options",
    ),
):
    """Analyze a .gguf file directly (no Ollama needed)."""
    if not LLAMA_CPP_AVAILABLE:
        _error_panel(
            "Missing Dependency",
            "✗ llama-cpp-python is not installed.",
            "The 'check' command requires llama-cpp-python to load .gguf files directly.",
            "[cyan]pip install stressllm[gguf][/cyan]",
        )
        raise typer.Exit(code=1)

    if not os.path.isfile(path):
        _error_panel(
            "File Error",
            f"✗ File not found: {path}",
            "The path does not point to an existing file.",
            "Check the path and try again. Make sure the .gguf file exists.",
        )
        raise typer.Exit(code=1)

    tester = StressTester()
    summary = tester.probe.get_system_summary()

    if verbose:
        _print_verbose(summary)

    if not json:
        hw_line = f"RAM: {summary['ram_total']}"
        if summary["gpu_enabled"]:
            hw_line = f"{summary['gpu_name']} · {summary['vram_total']} VRAM · {hw_line}"

        console.print(Panel(
            f"[bold]⚡ stressllm — Direct Weight Analysis[/bold]\n"
            f"[dim]{path}[/dim]\n"
            f"[dim]{hw_line}[/dim]",
            border_style="magenta",
        ))

    steps = [2048 * (4 ** i) for i in range(depth)]
    results = []
    try:
        if json:
            for row in tester.iter_direct_test(path, depth, n_gpu):
                results.append(row)
        else:
            with Progress(
                SpinnerColumn(),
                TextColumn("[bold cyan]{task.description}[/bold cyan]"),
                TimeElapsedColumn(),
                console=console,
                transient=True,
            ) as progress:
                for i, row in enumerate(tester.iter_direct_test(path, depth, n_gpu)):
                    ctx_label = row["context"]
                    if i == 0 and ctx_label == "LOAD":
                        # Load failure — just collect and break
                        results.append(row)
                        continue
                    task = progress.add_task(
                        f"Testing {ctx_label} context… ({i + 1}/{len(steps)})",
                        total=None,
                    )
                    results.append(row)
                    progress.remove_task(task)
                    console.print(
                        f"  [dim]✓[/dim] {ctx_label} — "
                        f"TTFT: [blue]{row['ttft']}[/blue]  "
                        f"{row['status']}"
                    )
    except KeyboardInterrupt:
        if not json:
            console.print("\n[yellow]Interrupted — showing partial results.[/yellow]")
    except MemoryError:
        _error_panel(
            "Out of Memory",
            "✗ System ran out of memory loading the model.",
            "The .gguf file is too large for your available RAM/VRAM.",
            "Try a smaller quantization or reduce [cyan]--n-gpu[/cyan] layers.",
        )
        if not results:
            raise typer.Exit(code=1)
    except Exception as e:
        _error_panel(
            "Benchmark Error",
            f"✗ {e}",
            "An unexpected error occurred during the direct benchmark.",
            "Check that the .gguf file is valid and not corrupted.",
        )
        if not results:
            raise typer.Exit(code=1)

    if results:
        _print_results(results, gpu_enabled=summary["gpu_enabled"], as_json=json)


@app.command(
    epilog="[dim]Example: stressllm info[/dim]",
)
def info():
    """Show detected hardware, dependency status, and benchmark reference."""
    probe = HardwareProbe()
    summary = probe.get_system_summary()

    # Hardware
    hw_table = Table(title="Hardware", title_style="bold", show_lines=True)
    hw_table.add_column("Component", style="cyan")
    hw_table.add_column("Value")

    if summary["gpu_enabled"]:
        hw_table.add_row("GPU", summary["gpu_name"])
        hw_table.add_row("VRAM", summary["vram_total"])
    else:
        if not NVML_AVAILABLE:
            hw_table.add_row(
                "GPU",
                "[yellow]Not detected[/yellow] — run [bold]pip install \"stressllm\\[gpu]\"[/bold] to enable",
            )
        else:
            hw_table.add_row(
                "GPU",
                "[yellow]Not detected[/yellow] — pynvml loaded but GPU init failed, check drivers",
            )

    hw_table.add_row("RAM", summary["ram_total"])
    hw_table.add_row("CPU Cores", str(psutil.cpu_count(logical=True)))
    console.print(hw_table)

    # Dependencies
    dep_table = Table(title="Dependencies", title_style="bold", show_lines=True)
    dep_table.add_column("Package", style="cyan")
    dep_table.add_column("Status")

    ollama_status = "[red]✗ Not reachable[/red]"
    try:
        import ollama
        ollama.list()
        ollama_status = "[green]✓ Running[/green]"
    except ImportError:
        ollama_status = "[red]✗ Not installed[/red]"
    except Exception:
        ollama_status = "[red]✗ Not running[/red] — start with: ollama serve"

    dep_table.add_row("Ollama", ollama_status)
    pynvml_status = (
        "[green]✓ Available[/green]"
        if NVML_AVAILABLE
        else "[yellow]✗ Not installed[/yellow] — pip install \"stressllm\\[gpu]\""
    )
    llama_status = (
        "[green]✓ Available[/green]" if LLAMA_CPP_AVAILABLE else "[dim]✗ Not installed[/dim]"
    )
    dep_table.add_row("pynvml (GPU metrics)", pynvml_status)
    dep_table.add_row("llama-cpp-python (check cmd)", llama_status)
    console.print(dep_table)

    # Depth reference
    depth_table = Table(title="Depth Levels", title_style="bold", show_lines=True)
    depth_table.add_column("Depth", justify="center", style="cyan")
    depth_table.add_column("Context Steps Tested")

    depth_table.add_row("1", "2k")
    depth_table.add_row("2", "2k → 8k")
    depth_table.add_row("3", "2k → 8k → 32k")
    depth_table.add_row("4", "2k → 8k → 32k → 128k")
    depth_table.add_row("5", "2k → 8k → 32k → 128k → 512k")
    console.print(depth_table)

    # Status legend
    status_table = Table(title="Status Legend", title_style="bold", show_lines=True)
    status_table.add_column("Status", justify="center")
    status_table.add_column("Meaning")
    status_table.add_column("TPS Range", style="dim")

    status_table.add_row("✅ Smooth", "Model runs comfortably at this context size", "> 15")
    status_table.add_row("⚠️ Slowing", "Usable but noticeably slower — approaching limits", "5 – 15")
    status_table.add_row("💀 Cliff", "Performance cliff — too slow for practical use", "< 5")
    status_table.add_row("💀 Crash", "Inference failed at this context size", "—")
    status_table.add_row("💀 OOM", "Out of memory — hardware cannot allocate", "—")
    console.print(status_table)

    probe.shutdown()

@app.command(
    epilog="[dim]Example: stressllm models[/dim]",
)
def models():
    """List available Ollama models and show sample run commands."""
    if not _check_ollama():
        raise typer.Exit(code=1)

    import ollama
    response = ollama.list()
    model_list = response.models if hasattr(response, "models") else []

    if not model_list:
        console.print(Panel(
            "[dim]No models found in Ollama.[/dim]\n\n"
            "  Pull one with: [cyan]ollama pull gemma:2b[/cyan]",
            title="Models",
            border_style="yellow",
        ))
        raise typer.Exit(code=0)

    table = Table(title="Available Ollama Models", title_style="bold", show_lines=True)
    table.add_column("Model", style="cyan")
    table.add_column("Size", style="yellow")
    table.add_column("Model Test Command", style="dim")

    for m in model_list:
        name = m.model if hasattr(m, "model") else str(m)
        size_bytes = m.size if hasattr(m, "size") else 0
        if size_bytes > 0:
            size_gb = round(size_bytes / (1024 ** 3), 1)
            size_str = f"{size_gb}GB"
        else:
            size_str = "—"
        table.add_row(name, size_str, f"stressllm run {name} --depth 3")

    console.print(table)



# ---------------------------------------------------------------------------
# Error handling for Ollama-specific failures
# ---------------------------------------------------------------------------

def _handle_ollama_error(e: Exception, model: str):
    """Route Ollama exceptions to the right Rich error panel."""
    err_str = str(e).lower()

    if "connection" in err_str or "refused" in err_str:
        _error_panel(
            "Connection Error",
            "✗ Lost connection to Ollama during the benchmark.",
            "Ollama may have crashed or been stopped.",
            "Restart with: [cyan]ollama serve[/cyan]",
        )
    elif "not found" in err_str or "model" in err_str:
        _error_panel(
            "Model Not Found",
            f"✗ Model '{model}' is not available in Ollama.",
            "The model hasn't been pulled yet.",
            f"Pull it with: [cyan]ollama pull {model}[/cyan]",
        )
    else:
        _error_panel(
            "Benchmark Error",
            f"✗ {e}",
            "An unexpected error occurred during the Ollama benchmark.",
            "Check Ollama logs or try a different model.",
        )
