"""Specter venv integration — patches an existing virtual environment.

Instead of creating a separate ``~/.specter/envs/<gpu>/`` directory,
``specter install`` hooks into the user's currently active venv by
replacing ``bin/python`` with a specter-proxy wrapper.  The original
binary is backed up as ``python.specter-original`` and restored by
``specter uninstall``.
"""

from __future__ import annotations

import json
import logging
import os
import re
import shutil
import stat
from pathlib import Path

log = logging.getLogger(__name__)

_BACKUP_SUFFIX = ".specter-original"
_MARKER_FILE = ".specter.json"

# Tools to forward to the remote GPU via SSH shims.
# Each tool gets a bin/<tool> shim that calls forward_tool().
# Adding a new tool = appending one string here.
_FORWARDED_TOOLS = ["nvidia-smi"]

# Legacy shim subdirectory names (used by previous versions).
# Kept for cleanup in restore_venv().
_TOOL_SHIM_DIR = ".specter-tools"

# Marker lines injected into bin/activate for PATH modification.
_ACTIVATE_BEGIN = "# >>> specter tool forwarding >>>"
_ACTIVATE_END = "# <<< specter tool forwarding <<<"


def detect_venv() -> Path | None:
    """Return the user's active venv path, or ``None`` if not in one.

    Only checks the ``VIRTUAL_ENV`` env var, which is set by
    ``source .venv/bin/activate``.  The ``sys.prefix`` fallback is
    intentionally omitted — it picks up the CLI's own install venv
    rather than the user's project venv.
    """
    venv_env = os.environ.get("VIRTUAL_ENV")
    if venv_env:
        p = Path(venv_env)
        if p.is_dir() and (p / "bin" / "python").exists():
            return p

    return None


def patch_venv(venv_path: Path, gpu_type: str) -> None:
    """Patch a venv so that ``python`` routes through specter-proxy.

    1. Back up ``bin/python`` → ``bin/python.specter-original``
    2. Back up ``bin/python3`` → ``bin/python3.specter-original``
    3. Write proxy wrapper as ``bin/python`` (uses backed-up binary)
    4. Symlink ``bin/python3`` → ``python``
    5. Back up and redirect ``bin/python3.X`` → ``python`` (the wrapper)
    6. Write ``.specter.json`` marker in the venv root

    Step 5 is critical for pip forwarding — standard venvs use the
    versioned binary (e.g. ``python3.12``) in pip's shebang line.
    Without this redirect, pip bypasses the wrapper entirely.

    The wrapper invokes ``python.specter-original -m specter_cli.proxy``
    directly to avoid an infinite recursion loop — ``specter-proxy``'s
    shebang points to ``bin/python``, which *is* the wrapper after patching.
    """
    bin_dir = venv_path / "bin"
    python_bin = bin_dir / "python"
    python3_bin = bin_dir / "python3"
    python_backup = bin_dir / f"python{_BACKUP_SUFFIX}"
    python3_backup = bin_dir / f"python3{_BACKUP_SUFFIX}"

    # Safety: if the venv is already patched, restore first to get a clean
    # state.  This prevents the wrapper-backed-up-as-original bug where
    # re-patching without restoring causes infinite recursion.
    if is_venv_patched(venv_path):
        log.info("Venv already patched — restoring before re-patching")
        restore_venv(venv_path)

    # If a compat symlink (created by restore_venv) exists, remove it
    # before backing up.  Compat symlinks point to "python" — they let
    # entry-point shebangs survive uninstall.  We must NOT remove real
    # backup symlinks (e.g. python.specter-original → /usr/bin/python3.12)
    # or re-patching will back up the wrapper script instead of the binary.
    if python_backup.is_symlink() and os.readlink(python_backup) == "python":
        python_backup.unlink()
        log.info("Removed compat symlink %s", python_backup)

    # ── Phase 1: Pre-resolve all symlink targets BEFORE any renames ──
    # The binaries form a symlink web (python → python3.12 → /usr/bin/python3.12).
    # Renaming any node breaks other links.  Pre-resolve everything first.
    resolved_bin_dir = bin_dir.resolve()

    python_resolved = python_bin.resolve() if python_bin.is_symlink() else None
    python3_resolved = python3_bin.resolve() if python3_bin.is_symlink() else None

    # Collect versioned python info: {name: resolved_path_or_None}
    versioned_info: dict[str, Path | None] = {}
    for entry in bin_dir.iterdir():
        if re.match(r"^python3\.\d+$", entry.name):
            versioned_info[entry.name] = entry.resolve() if entry.is_symlink() else None

    # ── Phase 2: Execute backups ──
    # Helper: create a backup that won't be broken by subsequent redirects.
    # - Resolved target outside bin_dir → absolute symlink (stable).
    # - Resolved target inside bin_dir → copy file content (survives renames).
    # - Real file (not symlink) → rename (move the actual file).
    def _backup(src: Path, dst: Path, resolved: Path | None) -> None:
        if resolved is not None:
            if src.exists() or src.is_symlink():
                src.unlink()
            if resolved.parent == resolved_bin_dir:
                # Target is inside bin dir — copy to survive rename/redirect.
                # The source may have been renamed already (e.g. python →
                # python.specter-original); check the backup path too.
                copy_src = resolved
                if not copy_src.exists():
                    copy_src = resolved.parent / (resolved.name + _BACKUP_SUFFIX)
                shutil.copy2(str(copy_src), str(dst))
            else:
                dst.symlink_to(resolved)
        else:
            src.rename(dst)

    # Back up python.
    if python_bin.exists() and not (python_backup.exists() or python_backup.is_symlink()):
        _backup(python_bin, python_backup, python_resolved)
        log.info("Backed up %s → %s", python_bin, python_backup)
    elif python_backup.is_symlink() and not python_backup.exists():
        # Dangling symlink from a previous session — remove before backup.
        python_backup.unlink()
        if python_bin.exists():
            _backup(python_bin, python_backup, python_resolved)
            log.info("Replaced dangling backup %s", python_backup)

    # Back up python3.
    if python3_bin.exists() and not (python3_backup.exists() or python3_backup.is_symlink()):
        _backup(python3_bin, python3_backup, python3_resolved)
        log.info("Backed up %s → %s", python3_bin, python3_backup)
    elif python3_backup.is_symlink() and not python3_backup.exists():
        python3_backup.unlink()
        if python3_bin.exists():
            _backup(python3_bin, python3_backup, python3_resolved)
            log.info("Replaced dangling backup %s", python3_backup)

    # Back up versioned python binaries (python3.12, etc.).
    versioned_names: list[str] = []
    for name, resolved in versioned_info.items():
        entry = bin_dir / name
        versioned_backup = bin_dir / f"{name}{_BACKUP_SUFFIX}"
        # Remove compat symlinks only (target == the unversioned name, e.g. "python3.12").
        # Real backup symlinks (→ /usr/bin/python3.12) must be preserved.
        if versioned_backup.is_symlink() and os.readlink(versioned_backup) == name:
            versioned_backup.unlink()
            log.info("Removed compat symlink %s", versioned_backup)
        if versioned_backup.is_symlink() and not versioned_backup.exists():
            # Dangling symlink — clean up before backup.
            versioned_backup.unlink()
        if not (versioned_backup.exists() or versioned_backup.is_symlink()) and (entry.exists() or entry.is_symlink()):
            _backup(entry, versioned_backup, resolved)
            log.info("Backed up %s → %s", entry, versioned_backup)
        versioned_names.append(name)

    # Resolve the site-packages directory containing specter_cli so the
    # wrapper can set PYTHONPATH.  This is critical when specter-host is
    # installed via pipx — the user's venv won't have specter_cli in its
    # own site-packages, so the backed-up python needs a hint.
    import specter_cli as _sc

    specter_lib = str(Path(_sc.__file__).resolve().parent.parent)

    # Write the proxy wrapper.
    #
    # Uses python.specter-original directly (not specter-proxy entry point)
    # to avoid infinite recursion — specter-proxy's shebang points to
    # bin/python, which IS this wrapper after patching.
    #
    # Smart entry-point detection: when the kernel processes a shebang like
    # #!/path/to/.venv/bin/python, it passes the script path as $1. If that
    # script lives in the venv's bin/ directory, it's a pip-installed entry
    # point (specter, pip, pytest, etc.) — pass through to real python so
    # local tools keep working. User scripts (python train.py) are never
    # in bin/ — those get proxied to the remote GPU.
    #
    # Pip hook: remote-first pip forwarding.  When pip/pip3 is detected,
    # the wrapper forwards the command to the remote venv first (where
    # python runs), then runs locally for IDE support.  If remote succeeds
    # but local fails (e.g. CUDA package on macOS), a warning is shown.
    python_bin.write_text(f"""\
#!/bin/bash
# Specter proxy — transparently executes Python on remote {gpu_type} GPU.
# Auto-generated by specter install. Do not edit.
# Original python backed up as python{_BACKUP_SUFFIX}
REAL_PYTHON="$(dirname "$0")/python{_BACKUP_SUFFIX}"
BIN_DIR="$(cd "$(dirname "$0")" && pwd)"
export PYTHONPATH="{specter_lib}:$PYTHONPATH"
export VIRTUAL_ENV="${{VIRTUAL_ENV:-$(cd "$(dirname "$0")/.." && pwd)}}"

# Tool introspection queries → always local.
# Package managers (uv, pip, setuptools) query the interpreter with -I
# (isolated mode) to detect version, platform, and site-packages paths.
# These must run locally — proxying them to remote breaks package managers.
# --version/-V are also local (informational, used by version detection).
case "${{1:-}}" in
    -I|--version|-V) exec "$REAL_PYTHON" "$@" ;;
esac

# If $1 is a file inside this venv's bin/, it's a shebang invocation
# of a pip-installed entry point (specter, pip, pytest, etc.).
if [ -n "$1" ] && [ -f "$1" ]; then
    SCRIPT_DIR="$(cd "$(dirname "$1")" 2>/dev/null && pwd)"
    if [ "$SCRIPT_DIR" = "$BIN_DIR" ]; then
        SCRIPT_NAME="$(basename "$1")"
        # Pip: remote-first, then local sync-back for IDE support.
        # Remote is where python runs — it's the primary target.
        if [ "$SCRIPT_NAME" = "pip" ] || [ "$SCRIPT_NAME" = "pip3" ]; then
            # 1. Forward to remote venv first (where python runs).
            "$REAL_PYTHON" -m specter_cli.deps --pip-forward "$@"
            REMOTE_EXIT=$?
            # 2. Run locally for IDE support (best-effort).
            "$REAL_PYTHON" "$@"
            LOCAL_EXIT=$?
            if [ $LOCAL_EXIT -ne 0 ] && [ $REMOTE_EXIT -eq 0 ]; then
                echo "specter: local install failed — IDE may not recognize this package" >&2
                echo "specter: ✓ installed on remote GPU" >&2
            fi
            # Remote success = overall success.
            if [ $REMOTE_EXIT -eq 0 ]; then
                exit 0
            fi
            exit $REMOTE_EXIT
        fi
        # Other entry points: pass through to real python.
        exec "$REAL_PYTHON" "$@"
    fi
fi

# Everything else → proxy to remote GPU.
exec "$REAL_PYTHON" -m specter_cli.proxy "$@"
""")
    python_bin.chmod(python_bin.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)

    # Symlink python3 → python.
    if python3_bin.exists() or python3_bin.is_symlink():
        python3_bin.unlink()
    python3_bin.symlink_to("python")

    # Redirect versioned python binaries (python3.12, etc.) to the wrapper
    # so pip shebangs (#!/.../bin/python3.12) route through it.
    # Backups were already created above (before writing the wrapper).
    for name in versioned_names:
        entry = bin_dir / name
        if entry.exists() or entry.is_symlink():
            entry.unlink()
        entry.symlink_to("python")
        log.info("Redirected %s → python (wrapper)", entry)

    # Write uv shim for remote-first uv forwarding.
    # When the venv is activated, bin/ is on PATH so this shim shadows
    # the system uv binary.  Package-modifying commands (sync, add,
    # remove, pip) are forwarded to the remote first, then run locally
    # for IDE support — same pattern as the pip hook.
    uv_shim = bin_dir / "uv"
    uv_shim.write_text(f"""\
#!/bin/bash
# Specter uv wrapper — remote-first uv forwarding.
# Auto-generated by specter install. Do not edit.
REAL_PYTHON="$(dirname "$0")/python{_BACKUP_SUFFIX}"
BIN_DIR="$(cd "$(dirname "$0")" && pwd)"
export PYTHONPATH="{specter_lib}:$PYTHONPATH"
export VIRTUAL_ENV="${{VIRTUAL_ENV:-$(cd "$(dirname "$0")/.." && pwd)}}"

# Find real uv binary (skip our shim).
CLEAN_PATH="${{PATH//$BIN_DIR:/}}"
CLEAN_PATH="${{CLEAN_PATH%:$BIN_DIR}}"
REAL_UV=$(PATH="$CLEAN_PATH" command -v uv 2>/dev/null)
if [ -z "$REAL_UV" ]; then
    echo "specter: uv not found on system PATH" >&2
    exit 1
fi

# Package-modifying commands → remote-first, then local.
case "${{1:-}}" in
    sync|add|remove|pip)
        # 1. Forward to remote first.
        "$REAL_PYTHON" -m specter_cli.deps --uv-forward "$@"
        REMOTE_EXIT=$?
        # 2. Run locally for IDE support (best-effort).
        "$REAL_UV" "$@"
        LOCAL_EXIT=$?
        if [ $LOCAL_EXIT -ne 0 ] && [ $REMOTE_EXIT -eq 0 ]; then
            echo "specter: local uv failed — some packages unavailable for this platform" >&2
            echo "specter: ✓ completed on remote GPU" >&2
        fi
        # Remote success = overall success.
        if [ $REMOTE_EXIT -eq 0 ]; then exit 0; fi
        exit $REMOTE_EXIT
        ;;
    *)
        # Non-modifying commands (version, lock, etc.): local only.
        exec "$REAL_UV" "$@"
        ;;
esac
""")
    uv_shim.chmod(uv_shim.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)
    log.info("Wrote uv shim %s", uv_shim)

    # Write tool-forwarding shims (nvidia-smi, etc.) directly into bin/.
    _write_tool_shims(bin_dir, specter_lib)

    # Write marker file.
    marker = venv_path / _MARKER_FILE
    marker.write_text(json.dumps({"gpu_type": gpu_type, "patched": True}, indent=2) + "\n")

    log.info("Patched venv %s for GPU %s", venv_path, gpu_type)


def _write_tool_shims(bin_dir: Path, specter_lib: str) -> None:
    """Write forwarding shims for all tools in ``_FORWARDED_TOOLS``.

    Shims are placed directly in ``bin/`` so they're on PATH immediately
    without requiring the user to re-source ``activate``.  Each shim is
    a thin bash script that calls
    ``python.specter-original -m specter_cli.forward <tool> "$@"``.
    """
    for tool in _FORWARDED_TOOLS:
        shim = bin_dir / tool
        shim.write_text(f"""\
#!/bin/bash
# Specter remote tool shim — forwards {tool} to remote GPU.
# Auto-generated by specter install. Do not edit.
REAL_PYTHON="$(dirname "$0")/python{_BACKUP_SUFFIX}"
export PYTHONPATH="{specter_lib}:$PYTHONPATH"
export VIRTUAL_ENV="${{VIRTUAL_ENV:-$(cd "$(dirname "$0")/.." && pwd)}}"
exec "$REAL_PYTHON" -m specter_cli.forward {tool} "$@"
""")
        shim.chmod(shim.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)
        log.info("Wrote %s shim %s", tool, shim)


def _patch_activate(bin_dir: Path) -> None:
    """Inject specter tool PATH into ``bin/activate``.

    Adds ``$VIRTUAL_ENV/bin/.specter-tools`` to PATH so tool shims (nvidia-smi,
    etc.) are available when the venv is activated.  Idempotent — won't
    add duplicate entries if activate is re-sourced.
    """
    activate = bin_dir / "activate"
    if not activate.exists():
        log.warning("activate script not found at %s — skipping PATH patch", activate)
        return

    text = activate.read_text()

    # Already patched?
    if _ACTIVATE_BEGIN in text:
        log.info("activate already patched — skipping")
        return

    # Inject after the existing PATH export.  The standard venv activate
    # script contains: export PATH="$VIRTUAL_ENV/bin:$PATH"
    # We add our block right after it.
    snippet = f"""
{_ACTIVATE_BEGIN}
if [[ ":$PATH:" != *":$VIRTUAL_ENV/bin/{_TOOL_SHIM_DIR}:"* ]]; then
    export PATH="$VIRTUAL_ENV/bin/{_TOOL_SHIM_DIR}:$PATH"
fi
{_ACTIVATE_END}
"""

    # Find the PATH export line and insert after it.
    anchor = 'export PATH="$VIRTUAL_ENV/bin:$PATH"'
    if anchor in text:
        text = text.replace(anchor, anchor + snippet, 1)
    else:
        # Fallback: append to end of file.
        text += snippet

    activate.write_text(text)
    log.info("Patched %s with specter tool PATH", activate)


def _unpatch_activate(bin_dir: Path) -> None:
    """Remove specter tool PATH injection from ``bin/activate``."""
    activate = bin_dir / "activate"
    if not activate.exists():
        return

    text = activate.read_text()
    if _ACTIVATE_BEGIN not in text:
        return

    # Remove everything from BEGIN to END (inclusive), plus surrounding newlines.
    pattern = rf"\n?{re.escape(_ACTIVATE_BEGIN)}.*?{re.escape(_ACTIVATE_END)}\n?"
    text = re.sub(pattern, "", text, flags=re.DOTALL)
    activate.write_text(text)
    log.info("Removed specter PATH patch from %s", activate)


def restore_venv(venv_path: Path) -> bool:
    """Restore a patched venv to its original state.

    Returns True if the venv was restored, False if nothing to restore.
    """
    bin_dir = venv_path / "bin"
    python_bin = bin_dir / "python"
    python3_bin = bin_dir / "python3"
    python_backup = bin_dir / f"python{_BACKUP_SUFFIX}"
    python3_backup = bin_dir / f"python3{_BACKUP_SUFFIX}"

    restored = False

    # Restore python.
    # Use is_symlink() to catch dangling symlinks (backups resolved to
    # absolute paths that may not exist on this machine, e.g. in tests).
    if python_backup.exists() or python_backup.is_symlink():
        if python_bin.exists() or python_bin.is_symlink():
            python_bin.unlink()
        python_backup.rename(python_bin)
        log.info("Restored %s", python_bin)
        restored = True

    # Restore python3.
    if python3_backup.exists() or python3_backup.is_symlink():
        if python3_bin.exists() or python3_bin.is_symlink():
            python3_bin.unlink()
        python3_backup.rename(python3_bin)
        log.info("Restored %s", python3_bin)
        restored = True

    # Restore versioned python binaries (python3.12.specter-original → python3.12).
    for backup in bin_dir.glob(f"python3.*{_BACKUP_SUFFIX}"):
        original_name = backup.name.removesuffix(_BACKUP_SUFFIX)
        # Skip non-versioned backups (python3.specter-original is handled above).
        if not re.match(r"^python3\.\d+$", original_name):
            continue
        original = bin_dir / original_name
        if original.exists() or original.is_symlink():
            original.unlink()
        backup.rename(original)
        log.info("Restored %s", original)
        restored = True

    # Leave a compatibility symlink so entry points installed while
    # specter was active (whose shebangs point to python.specter-original)
    # continue to resolve correctly.
    if restored:
        compat_link = bin_dir / f"python{_BACKUP_SUFFIX}"
        if not compat_link.exists():
            compat_link.symlink_to("python")
            log.info("Created compat symlink %s → python", compat_link)

    # Remove uv shim (created by patch_venv, not a backup).
    uv_shim = bin_dir / "uv"
    if uv_shim.exists():
        uv_shim.unlink()
        log.info("Removed uv shim %s", uv_shim)

    # Remove tool-forwarding shims from bin/ (current location).
    for tool in _FORWARDED_TOOLS:
        tool_shim = bin_dir / tool
        if tool_shim.exists():
            tool_shim.unlink()
            log.info("Removed %s shim %s", tool, tool_shim)
    # Clean up legacy shim directories from previous versions.
    for legacy_dir_name in (_TOOL_SHIM_DIR, "specter"):
        legacy_dir = bin_dir / legacy_dir_name
        if legacy_dir.is_dir():
            shutil.rmtree(legacy_dir)
            log.info("Removed legacy shim directory %s", legacy_dir)

    # Remove specter PATH patch from activate script.
    _unpatch_activate(bin_dir)

    # Remove marker.
    marker = venv_path / _MARKER_FILE
    if marker.exists():
        marker.unlink()

    if restored:
        log.info("Restored venv %s", venv_path)

    return restored


def is_venv_patched(venv_path: Path) -> bool:
    """Check if a venv has been patched by specter."""
    return (venv_path / _MARKER_FILE).exists()
