"""Specter local configuration management.

Manages ``~/.specter/`` directory with:
- ``auth.json`` — authentication token from ``specter login``
- ``remote.json`` — global (system-python) GPU session config
- ``sessions.json`` — registry of all active venv sessions
- ``ssh/`` — ControlMaster socket directory

Session config uses **two independent scopes**:

1. **Venv scope** (default): stored in ``$VIRTUAL_ENV/.specter.json``.
   Each venv has its own GPU session.  Written by ``specter install``.

2. **Global scope**: stored in ``~/.specter/remote.json``.
   Used when ``specter install --global`` targets system Python.

There is **no fallback** between scopes.  ``status`` and ``uninstall``
auto-detect: if ``VIRTUAL_ENV`` is set they use venv scope, otherwise
global scope.
"""

from __future__ import annotations

import contextlib
import json
import logging
import os
from pathlib import Path
from typing import Any

log = logging.getLogger(__name__)


def specter_dir() -> Path:
    """Root configuration directory.

    Defaults to ``~/.specter``, overridable via ``SPECTER_HOME`` env var.
    """
    return Path(os.environ.get("SPECTER_HOME", Path.home() / ".specter"))


def auth_file() -> Path:
    """Path to auth.json."""
    return specter_dir() / "auth.json"


def remote_file() -> Path:
    """Path to remote.json (global scope session config)."""
    return specter_dir() / "remote.json"


def sessions_file() -> Path:
    """Path to sessions.json (registry of active venv sessions)."""
    return specter_dir() / "sessions.json"


def envs_dir() -> Path:
    """Path to envs/ directory."""
    return specter_dir() / "envs"


def _venv_config_path() -> Path | None:
    """Return the active venv's ``.specter.json`` path, or ``None``.

    Uses the ``VIRTUAL_ENV`` env var (set by ``source activate`` or by
    the specter wrapper's ``export VIRTUAL_ENV=…`` line).
    """
    venv_env = os.environ.get("VIRTUAL_ENV")
    if venv_env:
        return Path(venv_env) / ".specter.json"
    return None


def ensure_specter_dir() -> Path:
    """Create ~/.specter/ directory tree if it doesn't exist."""
    d = specter_dir()
    d.mkdir(parents=True, exist_ok=True)
    (d / "ssh").mkdir(exist_ok=True)
    envs_dir().mkdir(exist_ok=True)
    return d


def load_auth() -> dict[str, Any]:
    """Load auth token from ~/.specter/auth.json.

    Returns
    -------
    Dict with at least ``token`` key, or empty dict if not logged in.
    """
    af = auth_file()
    if not af.is_file():
        return {}
    try:
        return json.loads(af.read_text())
    except (json.JSONDecodeError, OSError) as e:
        log.warning("Failed to read auth config: %s", e)
        return {}


def save_auth(token: str, **extra: Any) -> None:
    """Save auth token to ~/.specter/auth.json."""
    ensure_specter_dir()
    af = auth_file()
    data = {"token": token, **extra}
    af.write_text(json.dumps(data, indent=2) + "\n")
    af.chmod(0o600)  # Restrict access to owner only.
    log.debug("Auth token saved to %s", af)


def clear_auth() -> None:
    """Remove auth token."""
    af = auth_file()
    if af.is_file():
        af.unlink()
        log.debug("Auth token removed")


def get_token() -> str | None:
    """Get the current auth token, or None if not logged in."""
    auth = load_auth()
    return auth.get("token")


# ---------------------------------------------------------------------------
# Remote session config — two independent scopes
# ---------------------------------------------------------------------------


def load_remote_config(*, use_global: bool = False) -> dict[str, Any]:
    """Load remote GPU session config.

    When *use_global* is ``True``, reads from ``~/.specter/remote.json``
    (global / system-python scope).  Otherwise reads from the active
    venv's ``$VIRTUAL_ENV/.specter.json``.

    There is **no fallback** between scopes — they are independent.

    Returns
    -------
    Dict with keys: host, user, ssh_port, session_id, workspace, etc.
    Empty dict if no active session in the requested scope.
    """
    if use_global:
        rf = remote_file()
        if not rf.is_file():
            return {}
        try:
            return json.loads(rf.read_text())
        except (json.JSONDecodeError, OSError) as e:
            log.warning("Failed to read remote config: %s", e)
            return {}

    # Per-venv config.
    venv_cfg = _venv_config_path()
    if venv_cfg and venv_cfg.is_file():
        try:
            data = json.loads(venv_cfg.read_text())
            # Only use venv config if it has session data (not just marker).
            if data.get("host") and data.get("session_id"):
                return data
        except (json.JSONDecodeError, OSError) as e:
            log.warning("Failed to read venv config %s: %s", venv_cfg, e)

    return {}


def save_remote_config(
    *,
    host: str,
    user: str,
    ssh_port: int = 22,
    session_id: str,
    gpu_type: str,
    workspace: str = "~/.specter/workspace",
    venv_path: str | None = None,
    use_global: bool = False,
    **extra: Any,
) -> None:
    """Save remote GPU session config to a single scope.

    When *use_global* is ``True``, writes to ``~/.specter/remote.json``.
    When *venv_path* is given (and *use_global* is ``False``), writes to
    the venv's ``.specter.json``.

    Never writes to both — the two scopes are independent.
    """
    data: dict[str, Any] = {
        "host": host,
        "user": user,
        "ssh_port": ssh_port,
        "session_id": session_id,
        "gpu_type": gpu_type,
        "workspace": workspace,
        **extra,
    }

    if use_global:
        # Global scope — write to ~/.specter/remote.json only.
        ensure_specter_dir()
        rf = remote_file()
        rf.write_text(json.dumps(data, indent=2) + "\n")
        log.debug("Global session config saved to %s", rf)
        return

    if venv_path:
        data["venv_path"] = venv_path
        # Per-venv scope — merge into existing .specter.json (preserves marker).
        venv_cfg = Path(venv_path) / ".specter.json"
        existing: dict[str, Any] = {}
        if venv_cfg.is_file():
            with contextlib.suppress(json.JSONDecodeError, OSError):
                existing = json.loads(venv_cfg.read_text())
        existing.update(data)
        venv_cfg.write_text(json.dumps(existing, indent=2) + "\n")
        log.debug("Session config saved to %s", venv_cfg)
        return

    log.warning("save_remote_config called without venv_path or use_global")


def clear_remote_config(
    venv_path: str | Path | None = None,
    *,
    use_global: bool = False,
) -> None:
    """Remove remote session config from a single scope.

    When *use_global* is ``True``, removes ``~/.specter/remote.json``.
    Otherwise clears session fields from the venv's ``.specter.json``
    (preserving the ``patched`` marker).
    """
    if use_global:
        rf = remote_file()
        if rf.is_file():
            rf.unlink()
            log.debug("Global remote config removed")
        return

    session_keys = {
        "host",
        "user",
        "ssh_port",
        "session_id",
        "workspace",
        "python_local",
        "python_remote",
        "remote_venv",
        "venv_path",
    }

    # Clear session fields from venv config (keep marker fields).
    vp = str(venv_path) if venv_path else os.environ.get("VIRTUAL_ENV")
    if vp:
        venv_cfg = Path(vp) / ".specter.json"
        if venv_cfg.is_file():
            try:
                existing = json.loads(venv_cfg.read_text())
                for key in session_keys:
                    existing.pop(key, None)
                venv_cfg.write_text(json.dumps(existing, indent=2) + "\n")
                log.debug("Session fields cleared from %s", venv_cfg)
            except (json.JSONDecodeError, OSError) as e:
                log.warning("Failed to clear venv config %s: %s", venv_cfg, e)


def has_active_session(*, use_global: bool = False) -> bool:
    """Check if there's an active GPU session in the given scope."""
    config = load_remote_config(use_global=use_global)
    return bool(config.get("host") and config.get("session_id"))


# ---------------------------------------------------------------------------
# Session registry — tracks all venvs with active specter installs
# ---------------------------------------------------------------------------


def register_session(
    venv_path: str,
    session_id: str,
    gpu_type: str,
) -> None:
    """Register a venv session in the global registry.

    Called during ``specter install`` (venv scope) so that
    ``specter uninstall --all`` can find every patched venv later.
    """
    ensure_specter_dir()
    sf = sessions_file()

    registry: list[dict[str, str]] = []
    if sf.is_file():
        with contextlib.suppress(json.JSONDecodeError, OSError):
            data = json.loads(sf.read_text())
            registry = data.get("sessions", [])

    # Remove existing entry for this venv (idempotent re-install).
    registry = [s for s in registry if s.get("venv_path") != venv_path]
    registry.append(
        {
            "venv_path": venv_path,
            "session_id": session_id,
            "gpu_type": gpu_type,
        }
    )

    sf.write_text(json.dumps({"sessions": registry}, indent=2) + "\n")
    log.debug("Session registered: %s (%s)", venv_path, gpu_type)


def deregister_session(venv_path: str) -> None:
    """Remove a venv session from the global registry.

    Called during ``specter uninstall`` (venv scope).
    """
    sf = sessions_file()
    if not sf.is_file():
        return

    try:
        data = json.loads(sf.read_text())
        registry = data.get("sessions", [])
        registry = [s for s in registry if s.get("venv_path") != venv_path]
        sf.write_text(json.dumps({"sessions": registry}, indent=2) + "\n")
        log.debug("Session deregistered: %s", venv_path)
    except (json.JSONDecodeError, OSError) as e:
        log.warning("Failed to update session registry: %s", e)


def list_sessions() -> list[dict[str, str]]:
    """List all registered venv sessions.

    Returns a list of dicts with ``venv_path``, ``session_id``,
    and ``gpu_type`` keys.
    """
    sf = sessions_file()
    if not sf.is_file():
        return []
    try:
        data = json.loads(sf.read_text())
        return data.get("sessions", [])
    except (json.JSONDecodeError, OSError):
        return []
