"""Dependency synchronisation — pip freeze locally, pip install on remote.

Triggered in three ways:

1. **At ``specter install`` time** — eager sync so the pod has deps
   before the first ``python train.py``.
2. **After ``pip install`` in a patched venv** — the smart wrapper
   detects pip invocations and forwards the pip command to the remote
   venv first (remote-first), then runs locally for IDE support.
3. **Direct pip forwarding** — ``python.specter-original -m
   specter_cli.deps --pip-forward <pip-args>`` forwards a pip command
   to the remote venv via SSH.

The remote ``pip install`` fetches correct platform-specific wheels
(Linux x86_64 + CUDA) regardless of the local platform.
"""

from __future__ import annotations

import logging
import subprocess
import sys
import tempfile
from pathlib import Path

from specter_cli.ssh_config import rsync_ssh_args, ssh_base_args

log = logging.getLogger(__name__)

# Path to the remote venv created by specter install.
REMOTE_VENV = "~/.specter/venv"


def _pip_freeze(venv_path: Path) -> str:
    """Run ``pip freeze`` using the venv's backed-up python.

    Uses ``python.specter-original`` if the venv has been patched,
    otherwise uses the venv's regular python.
    """
    python_original = venv_path / "bin" / "python.specter-original"
    if python_original.exists():
        python = str(python_original)
    else:
        python = str(venv_path / "bin" / "python")

    result = subprocess.run(
        [python, "-m", "pip", "freeze", "--local"],
        capture_output=True,
        text=True,
        timeout=60,
    )

    if result.returncode != 0:
        log.warning("pip freeze failed: %s", result.stderr)
        return ""

    return result.stdout


def get_local_python_version() -> str:
    """Return the local Python version as 'major.minor' (e.g. '3.12')."""
    return f"{sys.version_info.major}.{sys.version_info.minor}"


def get_remote_python_version(
    *,
    host: str,
    user: str,
    port: int = 22,
    key_path: str | None = None,
) -> str | None:
    """Query the remote pod's Python version via SSH.

    Returns a version string like '3.12', or None if the check fails.
    """
    ssh_args = ssh_base_args(host, user, port=port, key_path=key_path)
    cmd = [
        "ssh",
        *ssh_args,
        "python3 -c \"import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')\"",
    ]
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=30,
        )
        if result.returncode == 0:
            return result.stdout.strip()
    except (subprocess.TimeoutExpired, OSError):
        pass

    return None


def check_python_version(
    *,
    host: str,
    user: str,
    port: int = 22,
    key_path: str | None = None,
) -> tuple[str, str | None]:
    """Compare local and remote Python versions.

    Returns (local_version, remote_version). remote_version is None
    if the check failed.
    """
    local = get_local_python_version()
    remote = get_remote_python_version(host=host, user=user, port=port, key_path=key_path)
    return local, remote


def create_remote_venv(
    *,
    host: str,
    user: str,
    port: int = 22,
    key_path: str | None = None,
    venv_path: str = REMOTE_VENV,
) -> str:
    """Create a virtual environment on the remote pod via SSH.

    Uses system python3 to create a venv at ``venv_path`` with
    ``--system-site-packages`` so user scripts can access synced deps
    AND system packages (torch, CUDA, etc. from the Docker image).

    Idempotent: ``python3 -m venv`` on an existing venv is a no-op.

    Raises
    ------
    subprocess.CalledProcessError
        If the SSH command exits non-zero.
    """
    create_cmd = f"python3 -m venv --system-site-packages {venv_path}"
    ssh_args = ssh_base_args(host, user, port=port, key_path=key_path)
    ssh_cmd = ["ssh", *ssh_args, create_cmd]

    result = subprocess.run(
        ssh_cmd,
        capture_output=True,
        text=True,
        timeout=60,
    )
    if result.returncode != 0:
        raise subprocess.CalledProcessError(
            result.returncode, ssh_cmd, output=result.stdout, stderr=result.stderr
        )

    # Install uv in the remote venv so proxy-time uv sync works.
    install_remote_uv(host=host, user=user, port=port, key_path=key_path, venv_path=venv_path)

    return venv_path


def install_remote_uv(
    *,
    host: str,
    user: str,
    port: int = 22,
    key_path: str | None = None,
    venv_path: str = REMOTE_VENV,
) -> None:
    """Install uv into the remote venv via pip.

    Called during ``specter install`` after creating the remote venv.
    Non-fatal: if the install fails, a warning is logged but
    installation continues — non-uv projects don't need it.
    """
    install_cmd = f"{venv_path}/bin/pip install -q uv"
    ssh_args = ssh_base_args(host, user, port=port, key_path=key_path)
    try:
        result = subprocess.run(
            ["ssh", *ssh_args, install_cmd],
            capture_output=True,
            text=True,
            timeout=120,
        )
        if result.returncode != 0:
            log.warning("Failed to install uv on remote (non-fatal): %s", result.stderr[:200])
    except subprocess.TimeoutExpired:
        log.warning("Remote uv install timed out (non-fatal)")
    except OSError as e:
        log.warning("Remote uv install error (non-fatal): %s", e)


def sync_deps(
    venv_path: Path,
    *,
    host: str,
    user: str,
    port: int = 22,
    key_path: str | None = None,
    remote_venv: str = REMOTE_VENV,
) -> bool:
    """Sync local pip packages to the remote pod.

    1. ``pip freeze`` locally (using the real python binary)
    2. Write to a temp file
    3. rsync the temp file to ``/tmp/specter-requirements.txt`` on pod
    4. SSH ``pip install -r /tmp/specter-requirements.txt`` on pod

    Returns True if deps were synced, False if nothing to sync.

    Raises
    ------
    subprocess.CalledProcessError
        If rsync or remote pip install fails.
    """
    # Step 1: pip freeze locally.
    freeze_output = _pip_freeze(venv_path)
    if not freeze_output.strip():
        log.info("No local packages to sync")
        return False

    # Filter out specter-host itself — it's the CLI, not needed on the pod.
    # Normalize underscores: pip may output specter_host or specter-host.
    lines = [
        line
        for line in freeze_output.splitlines()
        if not line.lower().replace("_", "-").startswith(("specter-host==", "specter-host @"))
    ]
    if not lines:
        log.info("No packages to sync (only specter-host)")
        return False

    filtered_freeze = "\n".join(lines) + "\n"

    # Step 2: Write to temp file.
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".txt", prefix="specter-reqs-", delete=False
    ) as f:
        f.write(filtered_freeze)
        local_reqs = f.name

    remote_reqs = "/tmp/specter-requirements.txt"

    try:
        # Step 3: rsync the requirements file to the pod.
        rsync_ssh = rsync_ssh_args(port=port, key_path=key_path)
        dest = f"{user}@{host}"

        rsync_cmd = [
            "rsync",
            "-az",
            "-e",
            rsync_ssh,
            local_reqs,
            f"{dest}:{remote_reqs}",
        ]
        subprocess.run(
            rsync_cmd,
            capture_output=True,
            text=True,
            timeout=30,
            check=True,
        )

        # Step 4: pip install into the remote venv.
        install_cmd = f"{remote_venv}/bin/pip install -q -r {remote_reqs}"
        ssh_args = ssh_base_args(host, user, port=port, key_path=key_path)
        ssh_cmd = ["ssh", *ssh_args, install_cmd]

        result = subprocess.run(
            ssh_cmd,
            capture_output=True,
            text=True,
            timeout=600,  # Large installs (torch) can take a while.
        )

        if result.returncode != 0:
            # Log stderr but don't fail hard — some packages may not
            # be available on the remote platform (e.g. macOS-only).
            log.warning("Remote pip install had issues: %s", result.stderr)
            # Print a summary of what failed for the user.
            stderr = result.stderr.strip()
            if stderr:
                print(
                    f"  Warning: some packages failed to install remotely:\n  {stderr[:200]}",
                    file=sys.stderr,
                )

        return True

    finally:
        # Clean up temp file.
        Path(local_reqs).unlink(missing_ok=True)


def pip_forward(raw_args: list[str]) -> int:
    """Forward a pip command to the remote venv via SSH.

    ``raw_args`` comes from the wrapper's ``$@`` — the first element is
    the local pip script path (e.g. ``/path/to/.venv/bin/pip``), the rest
    are pip arguments (e.g. ``install``, ``numpy``).  We skip the script
    path and forward the pip subcommand to ``~/.specter/venv/bin/pip``
    on the remote.

    Returns the remote pip exit code.
    """
    import shlex

    from specter_cli.config import load_remote_config

    # Skip local pip script path if first arg looks like a path.
    pip_args = raw_args[1:] if raw_args and "/" in raw_args[0] else raw_args

    if not pip_args:
        return 0

    config = load_remote_config()
    if not config or not config.get("host"):
        return 1  # No active session — silent exit.

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

    pip_cmd = f"{remote_venv}/bin/pip {shlex.join(pip_args)}"
    ssh_args = ssh_base_args(host, user, port=port, key_path=key_path)
    ssh_cmd = ["ssh", *ssh_args, pip_cmd]

    try:
        result = subprocess.run(ssh_cmd, capture_output=True, text=True, timeout=600)
        if result.returncode == 0:
            print("specter: \u2713 completed on remote GPU", file=sys.stderr)
        else:
            stderr = result.stderr.strip()
            if stderr:
                print(f"specter: remote pip: {stderr[:200]}", file=sys.stderr)
        return result.returncode
    except subprocess.TimeoutExpired:
        print("specter: \u26a0 remote pip timed out", file=sys.stderr)
        return 1
    except Exception as e:
        print(f"specter: \u26a0 remote pip error: {e}", file=sys.stderr)
        return 1


def uv_forward(raw_args: list[str]) -> int:
    """Forward a uv command to the remote venv via SSH.

    Called by the ``bin/uv`` shim when the venv is activated.
    Rsyncs the workspace first (so ``pyproject.toml`` / ``uv.lock``
    are on the remote), then runs ``uv <args>`` in the remote
    workspace with the remote venv activated.

    Returns the remote uv exit code.
    """
    import shlex

    from specter_cli.config import load_remote_config
    from specter_cli.ssh_config import rsync_ssh_args
    from specter_cli.sync import find_project_root, rsync_workspace

    config = load_remote_config()
    if not config or not config.get("host"):
        return 1  # No active session — silent exit.

    host = config["host"]
    user = config["user"]
    port = config.get("ssh_port", 22)
    key_path = config.get("key")
    remote_workspace = config.get("workspace", "~/.specter/workspace")
    remote_venv = config.get("remote_venv", REMOTE_VENV)

    # Rsync workspace so remote has pyproject.toml + uv.lock.
    project_root = find_project_root()
    rsync_ssh = rsync_ssh_args(port=port, key_path=key_path)
    dest = f"{user}@{host}"
    try:
        rsync_workspace(project_root, rsync_ssh, dest, remote_workspace)
    except subprocess.CalledProcessError:
        print("specter: workspace sync failed for uv forward", file=sys.stderr)
        return 1

    # Run uv on remote.
    uv_cmd = (
        f"cd {remote_workspace} && source {remote_venv}/bin/activate && uv {shlex.join(raw_args)}"
    )
    ssh_args = ssh_base_args(host, user, port=port, key_path=key_path)
    ssh_cmd = ["ssh", *ssh_args, uv_cmd]

    try:
        result = subprocess.run(ssh_cmd, capture_output=True, text=True, timeout=600)
        if result.returncode == 0:
            print("specter: \u2713 uv completed on remote GPU", file=sys.stderr)
        else:
            stderr = result.stderr.strip()
            if stderr:
                print(f"specter: remote uv: {stderr[:200]}", file=sys.stderr)
        return result.returncode
    except subprocess.TimeoutExpired:
        print("specter: \u26a0 remote uv timed out", file=sys.stderr)
        return 1
    except Exception as e:
        print(f"specter: \u26a0 remote uv error: {e}", file=sys.stderr)
        return 1


def main() -> None:
    """Entry point for automatic dep sync and pip forwarding.

    Called by the smart wrapper in two modes:

    1. ``python.specter-original -m specter_cli.deps --pip-forward <args>``
       — forwards a pip command to the remote venv.
    2. ``python.specter-original -m specter_cli.deps``
       — legacy pip-freeze sync (kept for ``specter install`` bootstrapping).

    Fails silently if no active session (e.g. user hasn't run install).
    """
    args = sys.argv[1:]

    # --pip-forward mode: forward pip command to remote venv.
    if args and args[0] == "--pip-forward":
        rc = pip_forward(args[1:])
        sys.exit(rc)

    # --uv-forward mode: forward uv command to remote workspace.
    if args and args[0] == "--uv-forward":
        rc = uv_forward(args[1:])
        sys.exit(rc)

    # Legacy freeze-sync mode (kept for specter install bootstrapping).
    import os

    from specter_cli.config import load_remote_config

    config = load_remote_config()
    if not config or not config.get("host"):
        # No active session — nothing to sync.
        return

    venv_env = os.environ.get("VIRTUAL_ENV")
    if not venv_env:
        return

    venv_path = Path(venv_env)
    if not venv_path.is_dir():
        return

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

    remote_venv = config.get("remote_venv", REMOTE_VENV)

    try:
        synced = sync_deps(
            venv_path,
            host=host,
            user=user,
            port=port,
            key_path=key_path,
            remote_venv=remote_venv,
        )
        if synced:
            print("specter: dependencies synced to remote GPU")
    except Exception as e:
        # Don't crash pip — just warn.
        print(f"specter: dep sync warning: {e}", file=sys.stderr)


if __name__ == "__main__":
    main()
