"""Specter Remote Proxy — transparent remote Python execution.

This is the ``bin/python`` wrapper that makes remote execution feel local.
When invoked as ``python train.py``, it:

1. Detects the project root (git repo, pyproject.toml, etc.)
2. Rsyncs the workspace to the remote machine
3. Executes the command via SSH with stdin/stdout/stderr passthrough
4. Returns the remote exit code

For MVP (Phase 1), all transport is plain SSH — no custom agent protocol.
SSH handles PTY, streams, signals, auth, and encryption.
"""

from __future__ import annotations

import logging
import os
import shlex
import subprocess
import sys
from pathlib import Path

from specter_cli.config import load_remote_config
from specter_cli.deps import REMOTE_VENV
from specter_cli.ssh_config import exec_ssh_passthrough, rsync_ssh_args, ssh_base_args
from specter_cli.sync import (
    _load_cached_hash,
    _save_cached_hash,
    find_project_root,
    rsync_workspace,
    workspace_hash,
)

log = logging.getLogger(__name__)

# Default remote workspace location.
_REMOTE_WORKSPACE = "~/.specter/workspace"


def _get_termination_reason(config: dict) -> str | None:
    """Query the broker for the session termination reason.

    Returns the ``termination_reason`` string (e.g. ``"usage_limit"``,
    ``"admin_block"``, ``"idle_timeout"``), or None if the session is
    still active or the broker is unreachable.

    Non-fatal: any error silently returns None so the caller falls
    back to a generic SSH-255 message.
    """
    session_id = config.get("session_id")
    if not session_id:
        return None

    try:
        from specter_cli.broker_client import BrokerClient, get_broker_url
        from specter_cli.config import get_token

        token = get_token()
        if not token:
            return None

        client = BrokerClient(token, base_url=get_broker_url(), timeout=5.0)
        try:
            session = client.get_session(session_id)
            if session.get("status") == "terminated":
                return session.get("termination_reason")
        finally:
            client.close()
    except Exception:
        # Broker unreachable, auth missing, etc. — fall back to generic message.
        log.debug("Could not query broker for termination reason", exc_info=True)
    return None


def _compute_remote_cwd(project_root: Path, local_cwd: Path, remote_workspace: str) -> str:
    """Compute the remote cwd corresponding to the local cwd.

    If the user is in ``/home/user/project/src/`` and the project root is
    ``/home/user/project/``, the remote cwd is ``~/.specter/workspace/src/``.
    """
    try:
        relative = local_cwd.resolve().relative_to(project_root.resolve())
    except ValueError:
        # cwd is outside project root — just use workspace root.
        return remote_workspace
    return f"{remote_workspace}/{relative}" if str(relative) != "." else remote_workspace


def _build_env_forwards(config: dict) -> list[str]:
    """Build remote env setup for key environment variables.

    We forward select env vars that affect Python/PyTorch behavior.
    """
    forward_vars = [
        "CUDA_VISIBLE_DEVICES",
        "TORCH_DTYPE",
        "PYTHONPATH",
        "PYTHONDONTWRITEBYTECODE",
        "PYTHONUNBUFFERED",
    ]
    env_parts = []
    for var in forward_vars:
        val = os.environ.get(var)
        if val is not None:
            env_parts.append(f"export {var}={shlex.quote(val)}")
    return env_parts


def _ensure_remote_dir(
    host: str,
    user: str,
    port: int,
    key_path: str | None,
    remote_dir: str,
) -> None:
    """Ensure the remote workspace directory exists via SSH mkdir -p.

    Handles tilde (~) expansion by letting the remote shell interpret it.
    """
    # Don't quote the path — let the remote shell expand ~ and $HOME.
    mkdir_cmd = f"mkdir -p {remote_dir}"
    args = ["ssh", *ssh_base_args(host, user, port=port, key_path=key_path), mkdir_cmd]
    subprocess.run(args, capture_output=True, text=True, timeout=30)


def _uv_sync_remote(
    host: str,
    user: str,
    port: int,
    key_path: str | None,
    remote_workspace: str,
    remote_venv: str,
) -> None:
    """Run ``uv sync --frozen`` on the remote to mirror local uv deps.

    Called at proxy time after rsync sends the workspace (including
    ``pyproject.toml`` and ``uv.lock``) to the remote.  ``--frozen``
    installs exactly what's in the lock file — no resolution, fast
    (~72 ms no-op).

    Non-fatal: if uv sync fails (uv not installed, network issue, etc.),
    a warning is logged but the user's command still executes.
    """
    sync_cmd = f"cd {remote_workspace} && UV_PROJECT_ENVIRONMENT={remote_venv} uv sync --frozen"
    args = ["ssh", *ssh_base_args(host, user, port=port, key_path=key_path), sync_cmd]
    try:
        result = subprocess.run(args, capture_output=True, text=True, timeout=300)
        if result.returncode != 0:
            log.warning("Remote uv sync failed (non-fatal): %s", result.stderr[:200])
    except subprocess.TimeoutExpired:
        log.warning("Remote uv sync timed out (non-fatal)")
    except OSError as e:
        log.warning("Remote uv sync error (non-fatal): %s", e)


def run_remote(argv: list[str]) -> int:
    """Execute a command remotely via SSH.

    This is the main entry point for the specter-proxy wrapper.

    Parameters
    ----------
    argv:
        The command to execute, e.g. ["python", "train.py", "--lr", "0.01"].
        The first element ("python") is replaced with the remote Python path.

    Returns
    -------
    The remote process exit code.
    """
    config = load_remote_config()
    if not config:
        print(
            "specter: not configured. Run 'specter install <gpu>' first.",
            file=sys.stderr,
        )
        return 1

    host = config["host"]
    user = config["user"]
    port = config.get("ssh_port", 22)
    key_path = config.get("key")
    remote_workspace = config.get("workspace", _REMOTE_WORKSPACE)
    remote_venv = config.get("remote_venv", REMOTE_VENV)
    dest = f"{user}@{host}"

    # Step 1: Detect project root and sync.
    project_root = find_project_root()
    local_cwd = Path.cwd()

    # Ensure remote workspace directory exists.
    _ensure_remote_dir(host, user, port, key_path, remote_workspace)

    log.info("syncing %s → %s:%s", project_root, dest, remote_workspace)

    # File sync optimization: skip rsync when workspace hasn't changed.
    # Uses git state (HEAD + diff + untracked) to detect changes.
    current_hash = workspace_hash(project_root)
    cached_hash = _load_cached_hash(project_root)
    skip_sync = current_hash is not None and current_hash == cached_hash

    if skip_sync:
        log.info("workspace unchanged (hash %s) — skipping rsync", current_hash)
    else:
        # Build rsync -e argument using ssh_config module.
        rsync_ssh = rsync_ssh_args(port=port, key_path=key_path)

        try:
            rsync_workspace(
                project_root=project_root,
                rsync_ssh=rsync_ssh,
                remote_dest=dest,
                remote_workspace=remote_workspace,
            )
        except subprocess.CalledProcessError as e:
            stderr = (e.stderr or "").strip()
            if "Connection refused" in stderr or "Connection timed out" in stderr:
                print(
                    "specter: cannot reach GPU — connection refused/timed out.\n"
                    "  Run 'specter status' to check your session.",
                    file=sys.stderr,
                )
            elif "No route to host" in stderr:
                print(
                    "specter: cannot reach GPU — no route to host.\n"
                    "  The session may have ended. Run 'specter status'.",
                    file=sys.stderr,
                )
            else:
                print(f"specter: workspace sync failed: {stderr}", file=sys.stderr)
            return 1
        except subprocess.TimeoutExpired:
            print(
                "specter: workspace sync timed out.\n"
                "  Large projects may need a .specterignore file to exclude big directories.",
                file=sys.stderr,
            )
            return 1

        # Cache the hash after successful sync.
        if current_hash is not None:
            _save_cached_hash(project_root, current_hash)

    # Step 1.5: Auto-sync uv deps if this is a uv-managed project.
    # rsync already sent pyproject.toml + uv.lock to the remote workspace.
    # uv sync --frozen installs exactly what's in the lock file (no resolution).
    if (project_root / "uv.lock").is_file():
        _uv_sync_remote(host, user, port, key_path, remote_workspace, remote_venv)

    # Step 2: Build the remote command.
    remote_cwd = _compute_remote_cwd(project_root, local_cwd, remote_workspace)
    env_exports = _build_env_forwards(config)
    env_prefix = " && ".join(env_exports) + " && " if env_exports else ""

    # Replace "python" in argv[0] with the remote python (venv python).
    remote_argv = list(argv)
    if remote_argv and remote_argv[0].startswith("python"):
        remote_argv[0] = "python3"

    # Make script paths relative to local CWD (which maps to remote CWD).
    for i, arg in enumerate(remote_argv[1:], 1):
        arg_path = Path(arg)
        if arg_path.is_file():
            try:
                # Relative to CWD works because remote CWD mirrors local CWD.
                rel = arg_path.resolve().relative_to(local_cwd.resolve())
                remote_argv[i] = str(rel)
            except ValueError:
                # File is outside CWD — try project root with absolute
                # workspace path so it resolves correctly.
                try:
                    rel = arg_path.resolve().relative_to(project_root.resolve())
                    remote_argv[i] = f"{remote_workspace}/{rel}"
                except ValueError:
                    pass  # Not in project — leave as-is.

    quoted_cmd = " ".join(shlex.quote(a) for a in remote_argv)
    # Activate the remote venv so user-installed packages are available,
    # then cd to the workspace.  Paths are unquoted for ~ expansion.
    venv_activate = f"source {remote_venv}/bin/activate"
    remote_command = f"cd {remote_cwd} && {venv_activate} && {env_prefix}{quoted_cmd}"

    # Step 3: SSH exec with passthrough I/O.
    # Use PTY (-t) only for interactive sessions (when stdin is a terminal).
    # Without PTY, SSH keeps stdout/stderr as separate streams.
    tty = sys.stdin.isatty()
    ssh_args = ssh_base_args(host, user, port=port, key_path=key_path, tty=tty)
    full_cmd = ["ssh", *ssh_args, remote_command]

    log.debug("remote exec: %s", " ".join(full_cmd))

    # Execute via SSH with full I/O passthrough and signal forwarding.
    rc = exec_ssh_passthrough(full_cmd)

    # SSH exit code 255 = connection lost (network error, session ended, etc.).
    # Query the broker for the termination reason to show a specific message.
    if rc == 255:
        reason = _get_termination_reason(config)
        if reason == "usage_limit":
            print(
                "\nspecter: session terminated — GPU usage limit reached.\n"
                "  Run 'specter status' for details, or contact your admin.",
                file=sys.stderr,
            )
        elif reason == "admin_block":
            print(
                "\nspecter: session terminated — account suspended by admin.\n"
                "  Contact your admin for details.",
                file=sys.stderr,
            )
        elif reason == "admin_pause":
            print(
                "\nspecter: session paused by admin.\n"
                "  Your session may resume. Run 'specter status' to check.",
                file=sys.stderr,
            )
        elif reason == "idle_timeout":
            print(
                "\nspecter: session terminated — idle timeout.\n"
                "  Run 'specter install' to start a new session.",
                file=sys.stderr,
            )
        else:
            print(
                "\nspecter: connection to GPU lost (SSH exit 255).\n"
                "  Run 'specter status' to check your session.",
                file=sys.stderr,
            )

    return rc


def main() -> None:
    """Entry point for the specter-proxy bin/python wrapper.

    Called when the user runs ``python train.py`` through the venv-like
    specter environment.
    """
    argv = sys.argv[1:]

    if not argv:
        # Interactive Python REPL — forward to remote.
        argv = ["python3"]

    # Prepend "python3" if the first arg looks like a script or flag.
    if argv[0].startswith("-") or argv[0].endswith(".py"):
        argv = ["python3", *argv]

    exit_code = run_remote(argv)
    sys.exit(exit_code)


if __name__ == "__main__":
    main()
