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

Commands::

    specter login                → authenticate with specter.host
    specter gpus                 → list available GPU types
    specter install <gpu>        → install a GPU into the active venv
    specter install --global <gpu> → install a GPU for system Python
    specter uninstall            → uninstall the GPU and restore the venv
    specter status               → show current session info
    specter ssh                  → open an SSH shell to the remote GPU
"""

from __future__ import annotations

import os
import sys
import time
from pathlib import Path

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

from specter_cli import __version__

console = Console()

_home = str(Path.home())


def _shorten_path(p: str) -> str:
    """Replace $HOME prefix with ~ for compact display."""
    if p.startswith(_home + "/"):
        return "~" + p[len(_home):]
    return p


def _scope_label(venv_path: Path | str | None, use_global: bool) -> str:
    """Return a human-readable scope label for CLI output."""
    if use_global:
        return "system Python"
    if venv_path:
        return _shorten_path(str(venv_path))
    return "this environment"


def _auto_use_global() -> bool:
    """Auto-detect scope: venv if VIRTUAL_ENV is set, global otherwise."""
    return not bool(os.environ.get("VIRTUAL_ENV"))


@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()
def gpus() -> None:
    """List available GPU types.

    Queries the broker for currently available GPU types that can be
    installed with ``specter install <gpu>``.  No authentication required.
    """
    from specter_cli.broker_client import BrokerClient, BrokerError, get_broker_url

    client = BrokerClient("", base_url=get_broker_url())
    try:
        gpu_types = client.list_gpu_types()
    except BrokerError as e:
        console.print(f"[red]✗[/red] Failed to fetch GPU types: {e}")
        sys.exit(1)
    except Exception:
        console.print(
            "[red]✗[/red] Could not reach the broker.\n  Check your connection or try again later."
        )
        sys.exit(1)
    finally:
        client.close()

    if not gpu_types:
        console.print("[yellow]No GPU types currently available.[/yellow]")
        return

    table = Table(title="Available GPU Types")
    table.add_column("Shortname", style="green")
    table.add_column("GPU", style="cyan")

    for gpu in gpu_types:
        table.add_row(gpu.get("name", "—"), gpu.get("display_name", "—"))

    console.print(table)
    console.print()
    console.print("Use [cyan]specter install <shortname>[/cyan] to install a GPU.")


@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 GPU.")
@click.option(
    "--global",
    "use_global",
    is_flag=True,
    help="Install for system Python instead of a virtual environment.",
)
def install(
    gpu_type: str,
    gpu_count: int,
    disk: int,
    image: str | None,
    use_global: bool,
) -> None:
    """Install a remote GPU into the active virtual environment.

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

    Requires an active virtual environment (or --global for system Python).
    After install, python 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,
        register_session,
        save_remote_config,
    )

    venv_path = None
    if not use_global:
        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]\n"
                "  Or use [cyan]specter install --global {gpu}[/cyan] for system Python."
            )
            sys.exit(1)

    scope = _scope_label(venv_path, use_global)

    if has_active_session(use_global=use_global):
        console.print("[yellow]⚠[/yellow] Active session exists. Run 'specter uninstall' first.")
        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:
        console.print(f"Installing {gpu_type} into {scope} ...")
        with console.status(f"[bold]  Installing {gpu_type}...[/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 GPU...[/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]  Status: {state}...[/bold]")
                time.sleep(3)

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

    except BrokerError as e:
        if e.error == "invalid_gpu_type":
            console.print(
                f"[red]✗[/red] GPU type '{gpu_type}' is not available.\n"
                "  Run [cyan]specter gpus[/cyan] to see available GPU types."
            )
        elif e.error == "usage_limit_reached":
            console.print(
                "[red]✗[/red] GPU usage limit reached.\n"
                "  Your account has used its allocated GPU time.\n"
                "  Contact your admin or visit [cyan]https://specter.host/dashboard[/cyan] "
                "to check your plan."
            )
        elif e.error == "account_blocked":
            console.print(
                "[red]✗[/red] Account suspended.\n"
                "  Your GPU access has been paused by an administrator.\n"
                "  Contact your admin for details."
            )
        elif e.error == "concurrent_limit_reached":
            console.print(
                "[red]✗[/red] Too many active GPU sessions.\n"
                "  You've reached your concurrent session limit.\n"
                "  Run [cyan]specter uninstall[/cyan] to free a slot, or visit\n"
                "  [cyan]https://specter.host/dashboard[/cyan] to manage sessions."
            )
        else:
            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]  Setting up GPU environment...[/bold]"):
            create_remote_venv(host=ssh_host, user=ssh_user, port=ssh_port)
        console.print("[green]✓[/green] GPU environment ready")
    except Exception as e:
        console.print(f"[red]✗[/red] Failed to set up GPU environment: {e}")
        sys.exit(1)

    if not use_global:
        # Patch the active venv so python routes to the remote GPU.
        # Must happen before save_remote_config — patch writes the .specter.json
        # marker, then save merges session data into it.
        patch_venv(venv_path, gpu_type)

    # Save session config to the appropriate scope (venv or global).
    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) if venv_path else None,
        use_global=use_global,
        python_local=local_ver,
        python_remote=remote_ver or "unknown",
        remote_venv=REMOTE_VENV,
    )

    # Register the session in the global registry (venv scope only).
    if not use_global and venv_path:
        register_session(
            venv_path=str(venv_path),
            session_id=session_id,
            gpu_type=gpu_type,
        )

    if not use_global:
        try:
            with console.status("[bold]  Syncing dependencies...[/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]specter ssh[/cyan], then [cyan]uv pip install .[/cyan]"
            )

    console.print()
    console.print(f"[green]✓[/green] {gpu_type} installed in {scope}")
    console.print()
    if use_global:
        console.print(f"Run Python as usual — it now runs on your {gpu_type}.")
    else:
        console.print(f"Run Python as usual — it now runs on your {gpu_type}:")
        console.print("  [cyan]python train.py[/cyan]")

        # Warn about uv run incompatibility if this is a uv project.
        from specter_cli.sync import find_project_root

        if (find_project_root() / "uv.lock").is_file():
            console.print()
            console.print(
                "  [yellow]Note:[/yellow] Use [cyan]uv sync[/cyan] for deps, "
                "not [cyan]uv run[/cyan] (uv run recreates the venv)."
            )


@cli.command()
def uninstall() -> None:
    """Uninstall the GPU and restore the virtual environment.

    Auto-detects scope: uses the active venv if one is activated,
    otherwise targets the global (system Python) session.
    """
    from specter_cli.auth import require_auth
    from specter_cli.broker_client import BrokerClient, BrokerError
    from specter_cli.config import (
        clear_remote_config,
        deregister_session,
        load_remote_config,
    )
    from specter_cli.venv import restore_venv

    use_global = _auto_use_global()
    config = load_remote_config(use_global=use_global)
    if not config or not config.get("session_id"):
        if use_global:
            console.print(
                "[yellow]No GPU installed globally.[/yellow]\n"
                "  If you installed in a venv, activate it first:\n"
                "  [cyan]source .venv/bin/activate[/cyan]"
            )
        else:
            console.print("[yellow]No GPU installed in this environment.[/yellow]")
        return

    session_id = config["session_id"]
    gpu_type = config.get("gpu_type", "GPU")
    scope = _scope_label(config.get("venv_path"), use_global)
    token = require_auth()

    from specter_cli.broker_client import get_broker_url

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

    if not use_global:
        # Restore the venv to its original state.
        venv_path = config.get("venv_path")
        if venv_path:
            restore_venv(Path(venv_path))
            console.print(f"[green]✓[/green] Restored venv: {_shorten_path(venv_path)}")

            # Remove from session registry.
            deregister_session(venv_path)

    # Clean up session config for the active scope.
    clear_remote_config(
        venv_path=config.get("venv_path") if not use_global else None,
        use_global=use_global,
    )

    console.print(f"[green]✓[/green] {gpu_type} uninstalled from {scope}.")


@cli.command("list")
def list_cmd() -> None:
    """List all active GPU sessions across all virtual environments.

    Shows every venv with an active specter install, plus the global
    session if one exists.  Local-only — no broker queries, no auth
    required.  Reads from the session registry and config files.
    """
    from specter_cli.config import list_sessions, load_remote_config

    sessions = list_sessions()
    global_config = load_remote_config(use_global=True)

    # Check if there's anything to show.
    has_global = bool(global_config and global_config.get("session_id"))
    if not sessions and not has_global:
        console.print("[dim]No active sessions.[/dim]")
        console.print("  Run [cyan]specter install <gpu>[/cyan] to get started.")
        return

    table = Table(title="Specter GPU Sessions")
    table.add_column("GPU", style="green")
    table.add_column("Session ID", style="dim")
    table.add_column("Scope", style="cyan")

    # Global session (if any).
    if has_global:
        table.add_row(
            global_config.get("gpu_type", "—"),
            global_config.get("session_id", "—"),
            "(global)",
        )

    # Per-venv sessions.
    for sess in sessions:
        venv_path = sess.get("venv_path", "—")
        if venv_path != "—":
            venv_path = _shorten_path(venv_path)
        table.add_row(
            sess.get("gpu_type", "—"),
            sess.get("session_id", "—"),
            venv_path,
        )

    console.print(table)


@cli.command()
def status() -> None:
    """Show current session status.

    Auto-detects scope: uses the active venv if one is activated,
    otherwise shows the global (system Python) session.
    """
    from specter_cli.auth import require_auth
    from specter_cli.broker_client import BrokerClient, BrokerError
    from specter_cli.config import load_remote_config

    use_global = _auto_use_global()
    config = load_remote_config(use_global=use_global)
    if not config or not config.get("session_id"):
        if use_global:
            console.print(
                "[yellow]No GPU installed globally.[/yellow]\n"
                "  If you installed in a venv, activate it first:\n"
                "  [cyan]source .venv/bin/activate[/cyan]"
            )
        else:
            console.print("[yellow]No GPU installed in this environment.[/yellow]")
        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()

    scope = "Global" if use_global else "Venv"
    table = Table(title=f"Specter Session ({scope})")
    table.add_column("Field", style="cyan")
    table.add_column("Value")

    table.add_row("Session ID", session.get("session_id", "—"))
    status_val = session.get("status", "—")
    table.add_row("Status", status_val)
    if status_val == "terminated":
        reason = session.get("termination_reason")
        if reason:
            reason_labels = {
                "usage_limit": "GPU usage limit reached",
                "concurrent_limit": "Concurrent session limit reached",
                "admin_block": "Account suspended by admin",
                "admin_pause": "Session paused by admin",
                "idle_timeout": "Idle timeout",
                "user_request": "User requested (specter uninstall)",
            }
            table.add_row("Terminated Reason", reason_labels.get(reason, reason))
    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)


@cli.command()
def ssh() -> None:
    """Open an SSH shell to the remote GPU.

    Connects to the active GPU session over SSH with a TTY.
    Auto-detects scope: uses the active venv if one is activated,
    otherwise targets the global (system Python) session.
    """
    from specter_cli.config import load_remote_config
    from specter_cli.ssh_config import ssh_base_args

    use_global = _auto_use_global()
    config = load_remote_config(use_global=use_global)
    if not config or not config.get("host"):
        if use_global:
            console.print(
                "[yellow]No GPU installed globally.[/yellow]\n"
                "  If you installed in a venv, activate it first:\n"
                "  [cyan]source .venv/bin/activate[/cyan]"
            )
        else:
            console.print(
                "[yellow]No GPU installed in this environment.[/yellow]\n"
                "  Run [cyan]specter install <gpu>[/cyan] first."
            )
        sys.exit(1)

    host = config["host"]
    user = config["user"]
    port = config.get("ssh_port", 22)
    key_path = config.get("key")

    args = ssh_base_args(host, user, port=port, key_path=key_path, tty=True)
    os.execvp("ssh", ["ssh", *args])
