"""Startup bootstrap for ``specter-host`` dependency mode.

When ``specter-host`` is installed as a project dependency with a GPU extra
(e.g. ``uv add specter-host[h100]``), this module is imported on every Python
startup via a ``.pth`` file in ``site-packages``.

The bootstrap detects which GPU extra was requested, checks for an existing
session, and provisions a remote GPU if needed — making GPU access transparent
to user code.

This module is designed to be **fast and safe**:
- No-ops in < 1ms when conditions aren't met (pipx/uvx, no extras, etc.)
- Never raises — all errors are caught (including BaseException)
- Never interferes with the CLI tool installed via pipx/uvx

Flow:
1. Safety checks (venv? not pipx/uvx? not already running?)
2. Detect GPU type from project's pyproject.toml extras
3. Check for existing session (reuse if GPU type matches)
4. If no session: check auth → prompt login if interactive → provision
5. Print dim status line so user knows they're on a GPU
"""

from __future__ import annotations

import contextlib
import importlib.metadata
import logging
import os
import re
import sys
from pathlib import Path

log = logging.getLogger("specter.pth_hook")

# Guard against re-entrant calls (e.g. subprocess spawns).
_BOOTSTRAP_LOCK_VAR = "_SPECTER_PTH_ACTIVE"

# ANSI escape for dim text (no dependency on rich at startup).
_DIM = "\033[2m"
_RESET = "\033[0m"


def _is_pipx_or_uvx() -> bool:
    """Return True if running inside a pipx/uvx managed environment."""
    # pipx sets PIPX_HOME; uvx uses UV_TOOL_DIR.
    if os.environ.get("PIPX_HOME"):
        return True
    if os.environ.get("UV_TOOL_DIR"):
        return True

    # pipx venvs live under ~/.local/pipx/venvs/
    prefix = Path(sys.prefix).resolve()
    pipx_venvs = Path.home() / ".local" / "pipx" / "venvs"
    if pipx_venvs.exists():
        try:
            prefix.relative_to(pipx_venvs)
            return True
        except ValueError:
            pass

    # uvx venvs live under ~/.local/share/uv/tools/
    uv_tools = Path.home() / ".local" / "share" / "uv" / "tools"
    if uv_tools.exists():
        try:
            prefix.relative_to(uv_tools)
            return True
        except ValueError:
            pass

    return False


def _in_virtualenv() -> bool:
    """Return True if running inside a virtualenv (not system Python)."""
    return hasattr(sys, "real_prefix") or (
        hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix
    )


def _find_project_root() -> Path | None:
    """Find the project root by looking for pyproject.toml.

    Starts from the venv parent directory (common for .venv/) and walks
    up to 5 levels looking for pyproject.toml.
    """
    # sys.prefix is the most reliable — always points to the current
    # interpreter's venv, even when VIRTUAL_ENV is stale or from an outer shell.
    venv_dir = Path(sys.prefix)
    candidates = [venv_dir.parent]

    # Also try VIRTUAL_ENV as a fallback (useful when venv is not at .venv/)
    venv_env = os.environ.get("VIRTUAL_ENV")
    if venv_env:
        venv_parent = Path(venv_env).parent
        if venv_parent not in candidates:
            candidates.append(venv_parent)

    for start in candidates:
        current = start
        for _ in range(5):
            if (current / "pyproject.toml").is_file():
                return current
            parent = current.parent
            if parent == current:
                break
            current = parent

    return None


# W13 fix: regex handles comma-separated extras like specter-host[h100,a100]
_EXTRAS_PATTERN = re.compile(r"specter[-_]host\[([\w][\w,\s-]*)\]", re.IGNORECASE)


def _detect_gpu_from_pyproject(project_root: Path) -> str | None:
    """Read pyproject.toml and find the specter-host[gpu_type] extra.

    Looks for ``specter-host[<extra>]`` in [project.dependencies],
    [project.optional-dependencies], or [dependency-groups].
    Returns the first extra name (e.g. "h100") or None if not found.
    """
    try:
        import tomllib
    except ModuleNotFoundError:
        return None

    pyproject_path = project_root / "pyproject.toml"
    try:
        with open(pyproject_path, "rb") as f:
            data = tomllib.load(f)
    except (OSError, tomllib.TOMLDecodeError):
        return None

    # Collect all dependency strings from common locations.
    dep_strings: list[str] = []

    # [project.dependencies]
    dep_strings.extend(data.get("project", {}).get("dependencies", []))

    # [project.optional-dependencies.*]
    for deps in data.get("project", {}).get("optional-dependencies", {}).values():
        dep_strings.extend(deps)

    # [dependency-groups.*] (PEP 735)
    for deps in data.get("dependency-groups", {}).values():
        if isinstance(deps, list):
            dep_strings.extend(d for d in deps if isinstance(d, str))

    # Search for specter-host[<extra>] — take the first extra if multiple.
    for dep in dep_strings:
        match = _EXTRAS_PATTERN.search(dep)
        if match:
            extras = match.group(1)
            # Take the first extra (e.g. "h100" from "h100,a100")
            first = extras.split(",")[0].strip().lower()
            if first:
                return first

    return None


def _detect_gpu_from_metadata() -> str | None:
    """Detect GPU type by checking which specter-gpu-* marker package is installed."""
    known_gpus = [
        "h100", "a100", "a100-40", "l40s", "l4", "a10g",
        "a6000", "t4", "5090", "4090", "3090", "rtx-pro-6000",
    ]
    for gpu in known_gpus:
        marker = f"specter-gpu-{gpu}"
        try:
            importlib.metadata.version(marker)
            return gpu
        except importlib.metadata.PackageNotFoundError:
            continue
    return None


def _detect_gpu_type() -> str | None:
    """Detect which GPU type was requested.

    Tries two methods:
    1. Read project's pyproject.toml for specter-host[gpu] dependency
    2. Check for installed specter-gpu-* marker packages

    Returns GPU type string (e.g. "h100") or None.
    """
    # Method 1: pyproject.toml (most reliable for dependency mode)
    project_root = _find_project_root()
    if project_root:
        gpu = _detect_gpu_from_pyproject(project_root)
        if gpu:
            return gpu

    # Method 2: marker packages (fallback)
    return _detect_gpu_from_metadata()


def _check_auth() -> str | None:
    """Check for auth token. Returns token or None.

    Checks:
    1. SPECTER_API_KEY env var
    2. Stored token in ~/.specter/auth.json
    """
    env_key = os.environ.get("SPECTER_API_KEY", "")
    if env_key:
        return env_key

    from specter_cli.config import get_token

    return get_token()


def _prompt_login() -> str | None:
    """Prompt interactive login if no auth token exists.

    Returns the token on success, None on failure/skip.
    """
    if not sys.stdin.isatty() or not sys.stderr.isatty():
        print(
            f"{_DIM}specter: no API key found. Set SPECTER_API_KEY or run 'specter login'.{_RESET}",
            file=sys.stderr,
        )
        return None

    print(
        f"\n{_DIM}specter: GPU dependency detected but no API key found.{_RESET}",
        file=sys.stderr,
    )
    print(
        f"{_DIM}  Run 'specter login' to authenticate, or set SPECTER_API_KEY.{_RESET}\n",
        file=sys.stderr,
    )
    return None


def _print_status(gpu_type: str, session_id: str | None = None) -> None:
    """Print a dim status line to stderr so user knows they're on a GPU."""
    sid = f" ({session_id[:12]})" if session_id else ""
    print(
        f"{_DIM}specter: {gpu_type} session active{sid}{_RESET}",
        file=sys.stderr,
    )


def bootstrap() -> None:
    """Main entry point called by the .pth file on Python startup.

    This function is designed to never raise — all errors are caught
    including BaseException (W12 fix). It must be fast (< 1ms) when
    conditions aren't met.
    """
    # Re-entrancy guard.
    if os.environ.get(_BOOTSTRAP_LOCK_VAR):
        return
    os.environ[_BOOTSTRAP_LOCK_VAR] = "1"

    try:
        _bootstrap_inner()
    except BaseException as exc:  # W12: catch KeyboardInterrupt, SystemExit too
        log.debug("specter .pth hook error: %s", exc)
    finally:
        os.environ.pop(_BOOTSTRAP_LOCK_VAR, None)


def _bootstrap_inner() -> None:
    """Inner bootstrap logic, separated for testability."""
    # 1. Must be in a virtualenv.
    if not _in_virtualenv():
        return

    # 2. Must not be pipx/uvx.
    if _is_pipx_or_uvx():
        return

    # 3. Detect GPU type.
    gpu_type = _detect_gpu_type()
    if not gpu_type:
        return

    log.debug("specter: detected GPU dependency: %s", gpu_type)

    # 4. Check for existing session that matches.
    from specter_cli.config import load_remote_config

    config = load_remote_config()
    if config.get("host") and config.get("session_id"):
        existing_gpu = config.get("gpu_type")
        if existing_gpu == gpu_type:
            log.debug("specter: reusing existing %s session", gpu_type)
            _print_status(gpu_type, config.get("session_id"))
            return
        # GPU type mismatch — warn but don't auto-reprovision.
        print(
            f"{_DIM}specter: GPU type mismatch — session has '{existing_gpu}' "
            f"but dependency requests '{gpu_type}'.{_RESET}\n"
            f"{_DIM}  Run 'specter uninstall && specter install {gpu_type}' to switch.{_RESET}",
            file=sys.stderr,
        )
        return

    # 5. Check auth.
    token = _check_auth()
    if not token:
        _prompt_login()
        return

    # 6. Auto-provision the GPU.
    print(
        f"{_DIM}specter: provisioning {gpu_type} GPU (from project dependency)...{_RESET}",
        file=sys.stderr,
    )

    try:
        from specter_cli.auth import read_ssh_public_key
        from specter_cli.broker_client import BrokerClient, BrokerError, get_broker_url
        from specter_cli.config import register_session, save_remote_config
        from specter_cli.deps import REMOTE_VENV, create_remote_venv, sync_deps
        from specter_cli.venv import detect_venv, patch_venv

        venv_path = detect_venv()
        if not venv_path:
            print(
                f"{_DIM}specter: could not detect virtual environment for patching.{_RESET}",
                file=sys.stderr,
            )
            return

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

            # Poll until ready.
            import time

            for _ in range(120):  # 6 min timeout (120 * 3s)
                status = client.get_session(session_id)
                state = status["status"]
                if state == "ready":
                    break
                if state == "error":
                    msg = status.get("error_message", "unknown")
                    print(
                        f"{_DIM}specter: provisioning failed: {msg}{_RESET}",
                        file=sys.stderr,
                    )
                    return
                time.sleep(3)
            else:
                print(
                    f"{_DIM}specter: provisioning timed out.{_RESET}",
                    file=sys.stderr,
                )
                return

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

            # Create remote venv.
            create_remote_venv(host=ssh_host, user=ssh_user, port=ssh_port)

            # 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),
                remote_venv=REMOTE_VENV,
            )

            register_session(
                venv_path=str(venv_path),
                session_id=session_id,
                gpu_type=gpu_type,
            )

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

            _print_status(gpu_type, session_id)

        except BrokerError as e:
            if e.error == "invalid_gpu_type":
                print(
                    f"{_DIM}specter: GPU type '{gpu_type}' is not available. "
                    f"Run 'specter gpus' to see options.{_RESET}",
                    file=sys.stderr,
                )
            else:
                print(
                    f"{_DIM}specter: provisioning error: {e}{_RESET}",
                    file=sys.stderr,
                )
        finally:
            client.close()

    except Exception as exc:
        print(
            f"{_DIM}specter: auto-provision failed: {exc}{_RESET}",
            file=sys.stderr,
        )
        log.debug("specter auto-provision traceback:", exc_info=True)
