"""Specter CLI — command-line interface for remote GPU execution.

Commands::

    specter login          → authenticate with specter.host
    specter install <gpu>  → provision a GPU pod and patch active venv
    specter uninstall      → tear down the active GPU pod and restore venv
    specter status         → show current session info
"""

from __future__ import annotations

import sys
import time

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

from specter_cli import __version__

console = Console()


@click.group()
@click.version_option(version=__version__, prog_name="specter")
def cli() -> None:
    """Specter — remote GPUs, locally.

    Run Python on cloud GPUs as if they were local machines.
    """


@cli.command()
def login() -> None:
    """Authenticate with specter.host via browser."""
    from specter_cli.auth import AuthError, login_with_browser

    try:
        login_with_browser()
    except AuthError as e:
        console.print(f"[red]✗[/red] {e}")
        sys.exit(1)


@cli.command()
def logout() -> None:
    """Remove stored authentication."""
    from specter_cli.auth import logout as do_logout

    do_logout()


@cli.command()
@click.argument("gpu_type")
@click.option("--gpu-count", default=1, help="Number of GPUs.")
@click.option("--disk", default=20, help="Disk size in GB.")
@click.option("--image", default=None, help="Docker image for the pod.")
def install(gpu_type: str, gpu_count: int, disk: int, image: str | None) -> None:
    """Provision a remote GPU and patch the active virtual environment.

    GPU_TYPE is a shorthand like h100, a100, 4090, l40s, etc.

    Requires an active virtual environment. After install, python
    in the venv transparently routes to the remote GPU.
    """
    from specter_cli.auth import read_ssh_public_key, require_auth
    from specter_cli.broker_client import BrokerClient, BrokerError
    from specter_cli.config import has_active_session, save_remote_config
    from specter_cli.venv import detect_venv, patch_venv

    # Require an active virtual environment.
    venv_path = detect_venv()
    if not venv_path:
        console.print(
            "[red]✗[/red] No active virtual environment detected.\n"
            "  Activate one first:\n"
            "  [cyan]python -m venv .venv && source .venv/bin/activate[/cyan]"
        )
        sys.exit(1)

    console.print(f"[cyan]Virtual env detected:[/cyan] {venv_path}")

    if has_active_session():
        console.print(
            "[yellow]⚠[/yellow] Active session exists. "
            "Run 'specter uninstall' first, or use a different GPU type."
        )
        if not click.confirm("Continue anyway?"):
            return

    token = require_auth()

    # Read SSH public key.
    try:
        ssh_pub_key = read_ssh_public_key()
    except FileNotFoundError as e:
        console.print(f"[red]✗[/red] {e}")
        sys.exit(1)

    # Provision via broker.
    from specter_cli.broker_client import get_broker_url

    client = BrokerClient(token, base_url=get_broker_url())
    try:
        with console.status(f"[bold]Provisioning {gpu_type} GPU...[/bold]"):
            session = client.provision(
                gpu_type=gpu_type,
                ssh_public_key=ssh_pub_key,
                gpu_count=gpu_count,
                disk_gb=disk,
                image=image,
            )
            session_id = session["session_id"]
            console.print(f"  Session: [cyan]{session_id}[/cyan]")

        # Poll until ready.
        error_msg = None
        with console.status("[bold]Waiting for pod to be ready...[/bold]") as spinner:
            while True:
                status = client.get_session(session_id)
                state = status["status"]

                if state == "ready":
                    break
                elif state == "error":
                    error_msg = status.get("error_message", "unknown error")
                    break

                spinner.update(f"[bold]Pod status: {state}...[/bold]")
                time.sleep(3)

        if error_msg:
            console.print(f"[red]✗[/red] Provisioning failed: {error_msg}")
            sys.exit(1)

    except BrokerError as e:
        console.print(f"[red]✗[/red] {e}")
        sys.exit(1)
    finally:
        client.close()

    # Extract SSH connection info.
    ssh_info = status["ssh"]
    ssh_host = ssh_info["host"]
    ssh_user = ssh_info["user"]
    ssh_port = ssh_info.get("port", 22)

    # Check Python version compatibility.
    from specter_cli.deps import REMOTE_VENV, check_python_version, create_remote_venv, sync_deps

    local_ver, remote_ver = check_python_version(host=ssh_host, user=ssh_user, port=ssh_port)
    console.print(f"[cyan]Python:[/cyan] local {local_ver}", end="")
    if remote_ver:
        console.print(f", remote {remote_ver}")
        if local_ver != remote_ver:
            console.print(
                f"[yellow]⚠[/yellow] Python version mismatch: "
                f"local {local_ver} ≠ remote {remote_ver}\n"
                f"  Compiled packages may not work. Consider matching versions."
            )
    else:
        console.print()
        console.print("[yellow]⚠[/yellow] Could not detect remote Python version")

    # Create remote venv — hard failure if this fails.
    try:
        with console.status("[bold]Creating remote virtual environment...[/bold]"):
            create_remote_venv(host=ssh_host, user=ssh_user, port=ssh_port)
        console.print("[green]✓[/green] Remote venv ready")
    except Exception as e:
        console.print(f"[red]✗[/red] Failed to create remote venv: {e}")
        sys.exit(1)

    # Save remote config (include venv_path and python versions for debugging).
    save_remote_config(
        host=ssh_host,
        user=ssh_user,
        ssh_port=ssh_port,
        session_id=session_id,
        gpu_type=gpu_type,
        venv_path=str(venv_path),
        python_local=local_ver,
        python_remote=remote_ver or "unknown",
        remote_venv=REMOTE_VENV,
    )

    # Patch the active venv so python routes to the remote GPU.
    patch_venv(venv_path, gpu_type)

    try:
        with console.status("[bold]Syncing dependencies to remote...[/bold]"):
            synced = sync_deps(
                venv_path,
                host=ssh_host,
                user=ssh_user,
                port=ssh_port,
            )
        if synced:
            console.print("[green]✓[/green] Dependencies synced")
        else:
            console.print("[dim]No dependencies to sync[/dim]")
    except Exception as e:
        console.print(f"[yellow]⚠[/yellow] Dep sync failed: {e}")
        console.print("  You can sync later with: [cyan]pip install -r requirements.txt[/cyan]")

    console.print()
    console.print(f"[green]✓[/green] {gpu_type} GPU ready!")
    console.print()
    console.print("Run Python as usual — it now executes on the remote GPU:")
    console.print("  [cyan]python train.py[/cyan]")


@cli.command()
def uninstall() -> None:
    """Tear down the active GPU pod and restore the virtual environment."""
    from specter_cli.auth import require_auth
    from specter_cli.broker_client import BrokerClient, BrokerError
    from specter_cli.config import clear_remote_config, load_remote_config
    from specter_cli.venv import restore_venv

    config = load_remote_config()
    if not config or not config.get("session_id"):
        console.print("[yellow]No active session to uninstall.[/yellow]")
        return

    session_id = config["session_id"]
    token = require_auth()

    from specter_cli.broker_client import get_broker_url

    client = BrokerClient(token, base_url=get_broker_url())
    try:
        with console.status("[bold]Terminating GPU pod...[/bold]"):
            client.destroy_session(session_id)
    except BrokerError as e:
        console.print(f"[red]✗[/red] {e}")
        sys.exit(1)
    finally:
        client.close()

    # Restore the venv to its original state.
    venv_path = config.get("venv_path")
    if venv_path:
        from pathlib import Path

        restore_venv(Path(venv_path))
        console.print(f"[green]✓[/green] Restored venv: {venv_path}")

    # Clean up local config.
    clear_remote_config()

    console.print("[green]✓[/green] GPU pod terminated.")


@cli.command()
def status() -> None:
    """Show current session status."""
    from specter_cli.auth import require_auth
    from specter_cli.broker_client import BrokerClient, BrokerError
    from specter_cli.config import load_remote_config

    config = load_remote_config()
    if not config or not config.get("session_id"):
        console.print("[yellow]No active session.[/yellow]")
        console.print("Run 'specter install <gpu>' to provision a GPU.")
        return

    token = require_auth()
    session_id = config["session_id"]

    from specter_cli.broker_client import get_broker_url

    client = BrokerClient(token, base_url=get_broker_url())
    try:
        session = client.get_session(session_id)
    except BrokerError as e:
        console.print(f"[red]✗[/red] {e}")
        sys.exit(1)
    finally:
        client.close()

    table = Table(title="Specter Session")
    table.add_column("Field", style="cyan")
    table.add_column("Value")

    table.add_row("Session ID", session.get("session_id", "—"))
    table.add_row("Status", session.get("status", "—"))
    table.add_row("GPU Type", session.get("gpu_type", "—"))
    table.add_row("GPU Name", session.get("gpu_name", "—"))
    table.add_row("Provider", session.get("provider", "—"))

    ssh = session.get("ssh")
    if ssh:
        table.add_row("SSH Host", ssh.get("host", "—"))
        table.add_row("SSH Port", str(ssh.get("port", "—")))
        table.add_row("SSH User", ssh.get("user", "—"))

    table.add_row("Created", session.get("created_at", "—"))

    console.print(table)
