"""Programmatic API for Specter GPU provisioning.

Usage::

    import specter_cli as specter

    # Provision a GPU and patch the venv — same as `specter install h100`
    specter.install("h100")
    # Then run your script: python train.py  (runs on GPU)

    # Check session status
    info = specter.status()  # dict or None

    # Terminate session
    specter.uninstall()

    # Context manager for auto-cleanup
    with specter.gpu("h100") as session:
        subprocess.run(["python", "train.py"])
    # session auto-terminated on exit

All functions are thin wrappers around the same broker/config
machinery used by the CLI.  Raises on auth errors (explicit intent
= explicit errors).
"""

from __future__ import annotations

import contextlib
import os
import sys
import time
from collections.abc import Generator
from contextlib import contextmanager
from pathlib import Path
from typing import Any


class SpecterError(Exception):
    """Base exception for programmatic API errors."""


class AuthError(SpecterError):
    """Raised when authentication is missing or invalid."""


class ProvisionError(SpecterError):
    """Raised when GPU provisioning fails."""


def _get_token() -> str:
    """Get auth token or raise AuthError."""
    env_key = os.environ.get("SPECTER_API_KEY", "")
    if env_key:
        return env_key

    from specter_cli.config import get_token

    token = get_token()
    if not token:
        raise AuthError("Not authenticated. Run 'specter login' or set SPECTER_API_KEY.")
    return token


def _get_ssh_public_key() -> str:
    """Read SSH public key or raise SpecterError."""
    from specter_cli.auth import read_ssh_public_key

    try:
        return read_ssh_public_key()
    except FileNotFoundError as e:
        raise SpecterError(str(e)) from e


def install(
    gpu_type: str,
    *,
    gpu_count: int = 1,
    disk_gb: int = 20,
    image: str | None = None,
) -> dict[str, Any]:
    """Install a remote GPU into the local venv.

    Same as ``specter install <gpu_type>`` on the CLI. Installs the
    GPU, patches the venv so ``python`` proxies to the remote machine,
    and returns. Subsequent ``python`` invocations run on the GPU.

    Idempotent: if a session with the same GPU type is already active,
    prints a dim status message and returns the existing session info.
    If a different GPU type is active, raises ``SpecterError``.

    Parameters
    ----------
    gpu_type:
        GPU shortname (e.g. ``"h100"``, ``"a100"``, ``"4090"``).
    gpu_count:
        Number of GPUs.
    disk_gb:
        Disk size in GB.
    image:
        Docker image for the GPU environment.

    Returns
    -------
    Session info dict with keys: ``session_id``, ``gpu_type``,
    ``host``, ``user``, ``ssh_port``.

    Raises
    ------
    AuthError
        If not authenticated.
    ProvisionError
        If provisioning fails.
    SpecterError
        If a session with a different GPU type is already active.
    """
    from specter_cli.config import (
        has_active_session,
        load_remote_config,
        register_session,
        save_remote_config,
    )

    # Detect venv scope.
    use_global = not bool(os.environ.get("VIRTUAL_ENV"))

    # Check for existing session — idempotent no-op if same GPU type.
    if has_active_session(use_global=use_global):
        config = load_remote_config(use_global=use_global)
        existing_gpu = config.get("gpu_type")
        if existing_gpu == gpu_type:
            _print_status(gpu_type, config.get("session_id"))
            return config
        raise SpecterError(
            f"Active session uses {existing_gpu}. "
            f"Run specter.uninstall() before installing {gpu_type}."
        )

    token = _get_token()
    ssh_pub_key = _get_ssh_public_key()

    # Detect venv path for patching.
    venv_path: Path | None = None
    if not use_global:
        venv_env = os.environ.get("VIRTUAL_ENV")
        if venv_env:
            venv_path = Path(venv_env)

    # Provision via broker.
    from specter_cli.broker_client import BrokerClient, BrokerError, get_broker_url

    client = BrokerClient(token, base_url=get_broker_url())
    try:
        session = client.provision(
            gpu_type=gpu_type,
            ssh_public_key=ssh_pub_key,
            gpu_count=gpu_count,
            disk_gb=disk_gb,
            image=image,
        )
        session_id = session["session_id"]

        # Poll until ready (max 10 minutes).
        deadline = time.monotonic() + 600
        while True:
            status_resp = client.get_session(session_id)
            state = status_resp["status"]

            if state == "ready":
                break
            elif state == "error":
                error_msg = status_resp.get("error_message", "unknown error")
                raise ProvisionError(f"GPU installation failed: {error_msg}")

            if time.monotonic() > deadline:
                raise ProvisionError("Timed out waiting for GPU (10 min).")

            time.sleep(3)

    except BrokerError as e:
        _raise_provision_error(e, gpu_type)
    finally:
        client.close()

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

    # Create remote venv.
    from specter_cli.deps import REMOTE_VENV, create_remote_venv, sync_deps

    create_remote_venv(host=ssh_host, user=ssh_user, port=ssh_port)

    # Patch local venv if in one.
    if venv_path:
        from specter_cli.venv import patch_venv

        patch_venv(venv_path, gpu_type)

    # Save session config.
    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,
        remote_venv=REMOTE_VENV,
    )

    # Register in global session 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,
        )

    # Sync deps (non-fatal).
    if venv_path:
        with contextlib.suppress(Exception):
            sync_deps(venv_path, host=ssh_host, user=ssh_user, port=ssh_port)

    result = {
        "session_id": session_id,
        "gpu_type": gpu_type,
        "host": ssh_host,
        "user": ssh_user,
        "ssh_port": ssh_port,
    }

    _print_status(gpu_type, session_id)
    return result


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

    Auto-detects scope: uses the active venv if ``VIRTUAL_ENV`` is set,
    otherwise targets the global (system Python) session.

    Raises
    ------
    AuthError
        If not authenticated.
    SpecterError
        If no active session.
    """
    from specter_cli.config import (
        clear_remote_config,
        deregister_session,
        load_remote_config,
    )

    use_global = not bool(os.environ.get("VIRTUAL_ENV"))
    config = load_remote_config(use_global=use_global)

    if not config or not config.get("session_id"):
        raise SpecterError("No active GPU session.")

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

    # Terminate via broker.
    from specter_cli.broker_client import BrokerClient, BrokerError, get_broker_url

    client = BrokerClient(token, base_url=get_broker_url())
    try:
        client.destroy_session(session_id)
    except BrokerError as e:
        raise SpecterError(f"Failed to terminate session: {e}") from e
    finally:
        client.close()

    # Restore venv if applicable.
    if not use_global:
        venv_path = config.get("venv_path")
        if venv_path:
            from specter_cli.venv import restore_venv

            restore_venv(Path(venv_path))
            deregister_session(venv_path)

    # Clean up config.
    clear_remote_config(
        venv_path=config.get("venv_path") if not use_global else None,
        use_global=use_global,
    )

    _print_dim("specter: GPU uninstalled")


def status() -> dict[str, Any] | None:
    """Return current session info, or ``None`` if no active session.

    Auto-detects scope: venv if ``VIRTUAL_ENV`` is set, global otherwise.
    Does NOT query the broker — reads local config only (fast, no auth).
    """
    from specter_cli.config import load_remote_config

    use_global = not bool(os.environ.get("VIRTUAL_ENV"))
    config = load_remote_config(use_global=use_global)

    if not config or not config.get("session_id"):
        return None

    return {
        "session_id": config.get("session_id"),
        "gpu_type": config.get("gpu_type"),
        "host": config.get("host"),
        "user": config.get("user"),
        "ssh_port": config.get("ssh_port"),
    }


@contextmanager
def gpu(
    gpu_type: str,
    *,
    gpu_count: int = 1,
    disk_gb: int = 20,
    image: str | None = None,
) -> Generator[dict[str, Any], None, None]:
    """Context manager that provisions a GPU on enter and terminates on exit.

    Usage::

        with specter.gpu("h100") as session:
            print(session["session_id"])
            # GPU available here
        # session terminated automatically

    Parameters are the same as :func:`install`.
    """
    session = install(gpu_type, gpu_count=gpu_count, disk_gb=disk_gb, image=image)
    try:
        yield session
    finally:
        with contextlib.suppress(SpecterError):
            uninstall()


def _print_status(gpu_type: str, session_id: str | None = None) -> None:
    """Print dim status line to stderr."""
    parts = [f"specter: {gpu_type} installed"]
    if session_id:
        parts.append(f"session {session_id}")
    _print_dim(" — ".join(parts))


def _print_dim(text: str) -> None:
    """Print dim ANSI text to stderr."""
    if hasattr(sys.stderr, "isatty") and sys.stderr.isatty():
        print(f"\033[2m{text}\033[0m", file=sys.stderr)
    else:
        print(text, file=sys.stderr)


def _raise_provision_error(e: Any, gpu_type: str) -> None:
    """Convert BrokerError to ProvisionError with a clear message."""
    error = getattr(e, "error", "unknown")
    if error == "invalid_gpu_type":
        raise ProvisionError(
            f"GPU type '{gpu_type}' is not available. Run 'specter gpus' to see available types."
        ) from e
    elif error == "usage_limit_reached":
        raise ProvisionError("GPU usage limit reached.") from e
    elif error == "account_blocked":
        raise ProvisionError("Account suspended.") from e
    elif error == "concurrent_limit_reached":
        raise ProvisionError("Concurrent session limit reached.") from e
    else:
        raise ProvisionError(str(e)) from e
