"""Dependency synchronisation — pyproject.toml on the remote.

If the project has a ``pyproject.toml``, rsync the workspace to the
remote and install via **uv** (``uv sync --frozen`` for uv projects,
``uv pip install .`` otherwise).  No filtering — the developer controls
their dependency spec.

If no ``pyproject.toml`` exists, no dep sync is performed.  The developer
installs packages manually via the pip/uv forwarding shims.

Also supports:

- **pip forwarding** — ``python.specter-original -m
  specter_cli.deps --pip-forward <pip-args>`` forwards a pip command
  to the remote venv via SSH.
- **uv forwarding** — ``python.specter-original -m
  specter_cli.deps --uv-forward <uv-args>`` forwards a uv command
  to the remote workspace via SSH.
"""

from __future__ import annotations

import logging
import subprocess
import sys
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"

# Default remote workspace location (matches proxy.py).
_REMOTE_WORKSPACE = "~/.specter/workspace"


def _is_uv_project(project_root: Path) -> bool:
    """Return True if the project uses uv (has ``uv.lock``)."""
    return (project_root / "uv.lock").is_file()


def _pyproject_sync(
    project_root: Path,
    *,
    host: str,
    user: str,
    port: int = 22,
    key_path: str | None = None,
    remote_venv: str = REMOTE_VENV,
    remote_workspace: str = _REMOTE_WORKSPACE,
) -> bool:
    """Install project deps on the remote from ``pyproject.toml``.

    Always uses **uv** on the remote (faster than pip).  Detects whether
    the project has a ``uv.lock`` to choose the install strategy:

    - **uv.lock present**: ``uv sync --frozen`` (deterministic from lock file)
    - **No uv.lock**: ``uv pip install .`` (PEP 517, works for poetry/hatch/setuptools)

    No filtering — the developer controls their dependency spec.

    Returns True if the install was attempted.

    Non-fatal: if the install fails, a warning is logged but the
    install continues — the developer can fix their spec and re-run.
    """
    from specter_cli.sync import rsync_workspace

    has_lockfile = _is_uv_project(project_root)
    strategy = "uv sync --frozen" if has_lockfile else "uv pip install ."
    log.info("Remote install strategy: %s", strategy)

    rsync_ssh = rsync_ssh_args(port=port, key_path=key_path)
    dest = f"{user}@{host}"

    # Rsync workspace so pyproject.toml + lock files are on the remote.
    rsync_workspace(project_root, rsync_ssh, dest, remote_workspace)

    # Build the install command — always uv on the remote.
    if has_lockfile:
        install_cmd = (
            f"cd {remote_workspace} && "
            f"UV_PROJECT_ENVIRONMENT={remote_venv} uv sync --frozen"
        )
    else:
        install_cmd = (
            f"cd {remote_workspace} && "
            f"uv pip install -q --python {remote_venv}/bin/python ."
        )

    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,
    )

    if result.returncode != 0:
        log.warning("Remote uv install had issues: %s", result.stderr)
        stderr = result.stderr.strip()
        # Detect missing uv and give a clear, actionable message.
        if "uv: not found" in stderr or "No such file" in stderr:
            print(
                "\n\u26a0 uv not available on remote \u2014 dependency sync failed.\n"
                "  Dependencies were not installed. Try: specter ssh, then pip install .\n",
                file=sys.stderr,
            )
        elif stderr:
            print(
                f"  Warning: uv install failed on remote:\n  {stderr[:200]}",
                file=sys.stderr,
            )

    return True


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 GPU'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 GPU 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 project dependencies to the remote GPU.

    If ``pyproject.toml`` exists in the project root, rsyncs the workspace
    and installs via uv on the remote.  No filtering — the developer
    controls their dependency spec.

    If no ``pyproject.toml`` exists, no dep sync is performed.

    Parameters
    ----------
    venv_path:
        Local venv path (kept for backward compatibility with callers).
    host, user, port, key_path:
        SSH connection parameters.
    remote_venv:
        Path to the remote venv.

    Returns True if deps were synced, False if nothing to sync.
    """
    from specter_cli.sync import find_project_root

    project_root = find_project_root()
    pyproject = project_root / "pyproject.toml"

    if not pyproject.is_file():
        log.warning("No pyproject.toml found — skipping dep sync")
        print(
            "\n\u26a0 No pyproject.toml found — dependencies were not synced to the GPU.\n"
            "\n"
            "  Specter uses pyproject.toml to install your project's dependencies on\n"
            "  the remote GPU. To set one up:\n"
            "\n"
            "    pip install hatch && hatch init\n"
            "    # or: uv init\n"
            "\n"
            "  Then run 'specter install' again.\n",
            file=sys.stderr,
        )
        return False

    log.info("Found pyproject.toml — installing deps on remote via uv")
    return _pyproject_sync(
        project_root,
        host=host,
        user=user,
        port=port,
        key_path=key_path,
        remote_venv=remote_venv,
    )


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} && UV_PROJECT_ENVIRONMENT={remote_venv} 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``
       — dep sync (installs from pyproject.toml if present).

    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)

    # Dep sync mode — install from pyproject.toml if present.
    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")
    except Exception as e:
        # Don't crash pip — just warn.
        print(f"specter: dep sync warning: {e}", file=sys.stderr)


if __name__ == "__main__":
    main()
