"""Tests for generic remote tool forwarding (forward.py + venv.py shims).

Covers:
1. forward_tool() — SSH command building, PTY, exit codes, no-session handling
2. forward.main() — CLI entry point argv parsing
3. _write_tool_shims() — shim script generation
4. _FORWARDED_TOOLS lifecycle — created by patch_venv, removed by restore_venv
"""

from __future__ import annotations

import stat
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest

# ---------------------------------------------------------------------------
# forward_tool() tests
# ---------------------------------------------------------------------------

_LOAD_CONFIG = "specter_cli.forward.load_remote_config"
_SSH_BASE = "specter_cli.forward.ssh_base_args"
_POPEN = "specter_cli.ssh_config.subprocess.Popen"


class TestForwardTool:
    """Tests for forward_tool() in forward.py."""

    @patch(_POPEN)
    @patch(_SSH_BASE, return_value=["-o", "StrictHostKeyChecking=no", "root@1.2.3.4"])
    @patch(
        _LOAD_CONFIG,
        return_value={
            "host": "1.2.3.4",
            "user": "root",
            "ssh_port": 22,
        },
    )
    def test_builds_correct_ssh_command(
        self,
        _config: MagicMock,
        _ssh_base: MagicMock,
        mock_popen: MagicMock,
    ) -> None:
        """forward_tool should build ssh <args> <tool> <user_args>."""
        from specter_cli.forward import forward_tool

        proc = MagicMock()
        proc.returncode = 0
        proc.wait.return_value = 0
        mock_popen.return_value = proc

        forward_tool("nvidia-smi", [])

        mock_popen.assert_called_once()
        cmd = mock_popen.call_args[0][0]
        assert cmd[0] == "ssh"
        assert "nvidia-smi" in " ".join(cmd)

    @patch(_POPEN)
    @patch(_SSH_BASE, return_value=["-o", "StrictHostKeyChecking=no", "root@1.2.3.4"])
    @patch(
        _LOAD_CONFIG,
        return_value={
            "host": "1.2.3.4",
            "user": "root",
            "ssh_port": 22,
        },
    )
    def test_passes_args_to_remote_tool(
        self,
        _config: MagicMock,
        _ssh_base: MagicMock,
        mock_popen: MagicMock,
    ) -> None:
        """User args should be passed through to the remote tool."""
        from specter_cli.forward import forward_tool

        proc = MagicMock()
        proc.returncode = 0
        proc.wait.return_value = 0
        mock_popen.return_value = proc

        forward_tool("nvidia-smi", ["--query-gpu=gpu_name", "--format=csv"])

        cmd = mock_popen.call_args[0][0]
        remote_cmd = cmd[-1]
        assert "--query-gpu=gpu_name" in remote_cmd
        assert "--format=csv" in remote_cmd

    @patch(_POPEN)
    @patch(_SSH_BASE, return_value=["-o", "StrictHostKeyChecking=no", "root@1.2.3.4"])
    @patch(
        _LOAD_CONFIG,
        return_value={
            "host": "1.2.3.4",
            "user": "root",
            "ssh_port": 22,
        },
    )
    def test_returns_remote_exit_code(
        self,
        _config: MagicMock,
        _ssh_base: MagicMock,
        mock_popen: MagicMock,
    ) -> None:
        """Should return the remote tool's exit code."""
        from specter_cli.forward import forward_tool

        proc = MagicMock()
        proc.returncode = 42
        proc.wait.return_value = 42
        mock_popen.return_value = proc

        rc = forward_tool("nvidia-smi", [])

        assert rc == 42

    @patch(_POPEN)
    @patch(_SSH_BASE, return_value=["-o", "StrictHostKeyChecking=no", "root@1.2.3.4"])
    @patch(
        _LOAD_CONFIG,
        return_value={
            "host": "1.2.3.4",
            "user": "root",
            "ssh_port": 22,
        },
    )
    def test_returncode_none_returns_1(
        self,
        _config: MagicMock,
        _ssh_base: MagicMock,
        mock_popen: MagicMock,
    ) -> None:
        """If proc.returncode is None (killed), should return 1."""
        from specter_cli.forward import forward_tool

        proc = MagicMock()
        proc.returncode = None
        proc.wait.return_value = None
        mock_popen.return_value = proc

        rc = forward_tool("nvidia-smi", [])

        assert rc == 1

    @patch(_LOAD_CONFIG, return_value={})
    def test_no_session_returns_127(
        self,
        _config: MagicMock,
        capsys: pytest.CaptureFixture[str],
    ) -> None:
        """No active session → exit 127 with helpful message."""
        from specter_cli.forward import forward_tool

        rc = forward_tool("nvidia-smi", [])

        assert rc == 127
        captured = capsys.readouterr()
        assert "no GPU session" in captured.err
        assert "specter install" in captured.err

    @patch(_LOAD_CONFIG, return_value=None)
    def test_none_config_returns_127(
        self,
        _config: MagicMock,
        capsys: pytest.CaptureFixture[str],
    ) -> None:
        """None config → exit 127."""
        from specter_cli.forward import forward_tool

        rc = forward_tool("nvidia-smi", [])

        assert rc == 127

    @patch(_LOAD_CONFIG, return_value={"session_id": "sess_123"})
    def test_config_without_host_returns_127(
        self,
        _config: MagicMock,
        capsys: pytest.CaptureFixture[str],
    ) -> None:
        """Config missing 'host' key → exit 127."""
        from specter_cli.forward import forward_tool

        rc = forward_tool("nvidia-smi", [])

        assert rc == 127

    @patch(_POPEN)
    @patch(_SSH_BASE, return_value=["-p", "2222", "ubuntu@5.6.7.8"])
    @patch(
        _LOAD_CONFIG,
        return_value={
            "host": "5.6.7.8",
            "user": "ubuntu",
            "ssh_port": 2222,
            "key": "/home/user/.ssh/specter_key",
        },
    )
    def test_uses_port_and_key_from_config(
        self,
        _config: MagicMock,
        mock_ssh_base: MagicMock,
        mock_popen: MagicMock,
    ) -> None:
        """SSH port and key from config should be forwarded to ssh_base_args."""
        from specter_cli.forward import forward_tool

        proc = MagicMock()
        proc.returncode = 0
        proc.wait.return_value = 0
        mock_popen.return_value = proc

        forward_tool("nvidia-smi", [])

        mock_ssh_base.assert_called_once_with(
            "5.6.7.8",
            "ubuntu",
            port=2222,
            key_path="/home/user/.ssh/specter_key",
            tty=False,
        )

    @patch(_POPEN)
    @patch(_SSH_BASE, return_value=["root@1.2.3.4"])
    @patch(
        _LOAD_CONFIG,
        return_value={
            "host": "1.2.3.4",
            "user": "root",
            "ssh_port": 22,
        },
    )
    def test_popen_uses_passthrough_io(
        self,
        _config: MagicMock,
        _ssh_base: MagicMock,
        mock_popen: MagicMock,
    ) -> None:
        """Popen should use stdin/stdout/stderr passthrough, not capture_output."""
        import sys

        from specter_cli.forward import forward_tool

        proc = MagicMock()
        proc.returncode = 0
        proc.wait.return_value = 0
        mock_popen.return_value = proc

        forward_tool("nvidia-smi", ["-l", "1"])

        kwargs = mock_popen.call_args[1]
        assert kwargs["stdin"] is sys.stdin
        assert kwargs["stdout"] is sys.stdout
        assert kwargs["stderr"] is sys.stderr
        assert kwargs["start_new_session"] is True

    @patch(_POPEN)
    @patch(_SSH_BASE, return_value=["root@1.2.3.4"])
    @patch(
        _LOAD_CONFIG,
        return_value={
            "host": "1.2.3.4",
            "user": "root",
            "ssh_port": 22,
        },
    )
    def test_signal_handlers_restored(
        self,
        _config: MagicMock,
        _ssh_base: MagicMock,
        mock_popen: MagicMock,
    ) -> None:
        """Signal handlers should be restored after proc.wait()."""
        import signal

        from specter_cli.forward import forward_tool

        proc = MagicMock()
        proc.returncode = 0
        proc.wait.return_value = 0
        mock_popen.return_value = proc

        original_int = signal.getsignal(signal.SIGINT)
        original_term = signal.getsignal(signal.SIGTERM)

        forward_tool("nvidia-smi", [])

        assert signal.getsignal(signal.SIGINT) == original_int
        assert signal.getsignal(signal.SIGTERM) == original_term

    @patch(_POPEN)
    @patch(_SSH_BASE, return_value=["root@1.2.3.4"])
    @patch(
        _LOAD_CONFIG,
        return_value={
            "host": "1.2.3.4",
            "user": "root",
            "ssh_port": 22,
        },
    )
    def test_no_args_sends_bare_tool_name(
        self,
        _config: MagicMock,
        _ssh_base: MagicMock,
        mock_popen: MagicMock,
    ) -> None:
        """No args → just the tool name, no trailing space."""
        from specter_cli.forward import forward_tool

        proc = MagicMock()
        proc.returncode = 0
        proc.wait.return_value = 0
        mock_popen.return_value = proc

        forward_tool("nvidia-smi", [])

        cmd = mock_popen.call_args[0][0]
        remote_cmd = cmd[-1]
        assert remote_cmd == "nvidia-smi"

    @patch(_POPEN)
    @patch(_SSH_BASE, return_value=["root@1.2.3.4"])
    @patch(
        _LOAD_CONFIG,
        return_value={
            "host": "1.2.3.4",
            "user": "root",
            "ssh_port": 22,
        },
    )
    def test_works_with_any_tool_name(
        self,
        _config: MagicMock,
        _ssh_base: MagicMock,
        mock_popen: MagicMock,
    ) -> None:
        """Should work with arbitrary tool names (nvcc, nvidia-debugdump, etc.)."""
        from specter_cli.forward import forward_tool

        proc = MagicMock()
        proc.returncode = 0
        proc.wait.return_value = 0
        mock_popen.return_value = proc

        forward_tool("nvcc", ["--version"])

        cmd = mock_popen.call_args[0][0]
        remote_cmd = cmd[-1]
        assert "nvcc" in remote_cmd
        assert "--version" in remote_cmd


# ---------------------------------------------------------------------------
# forward.main() tests
# ---------------------------------------------------------------------------


class TestForwardMain:
    """Tests for forward.main() CLI entry point."""

    @patch("specter_cli.forward.forward_tool", return_value=0)
    @patch("specter_cli.forward.sys")
    def test_main_parses_tool_and_args(
        self,
        mock_sys: MagicMock,
        mock_forward: MagicMock,
    ) -> None:
        """main() should parse argv into tool + args."""
        mock_sys.argv = ["forward", "nvidia-smi", "--query-gpu=name"]
        mock_sys.exit = MagicMock()

        from specter_cli.forward import main

        main()

        mock_forward.assert_called_once_with("nvidia-smi", ["--query-gpu=name"])
        mock_sys.exit.assert_called_once_with(0)

    @patch("specter_cli.forward.forward_tool", return_value=127)
    @patch("specter_cli.forward.sys")
    def test_main_forwards_exit_code(
        self,
        mock_sys: MagicMock,
        mock_forward: MagicMock,
    ) -> None:
        """main() should forward the exit code from forward_tool."""
        mock_sys.argv = ["forward", "nvidia-smi"]
        mock_sys.exit = MagicMock()

        from specter_cli.forward import main

        main()

        mock_sys.exit.assert_called_once_with(127)

    @patch("specter_cli.forward.sys")
    def test_main_no_args_exits_1(self, mock_sys: MagicMock) -> None:
        """main() with no tool arg should exit 1."""
        mock_sys.argv = ["forward"]
        mock_sys.exit = MagicMock(side_effect=SystemExit(1))

        from specter_cli.forward import main

        with pytest.raises(SystemExit):
            main()

        mock_sys.exit.assert_called_once_with(1)


# ---------------------------------------------------------------------------
# _write_tool_shims() tests
# ---------------------------------------------------------------------------


class TestWriteToolShims:
    """Tests for _write_tool_shims() in venv.py."""

    @pytest.fixture
    def bin_dir(self, tmp_path: Path) -> Path:
        d = tmp_path / "bin"
        d.mkdir()
        return d

    def test_creates_nvidia_smi_shim(self, bin_dir: Path) -> None:
        """Should write a bin/nvidia-smi shim."""
        from specter_cli.venv import _write_tool_shims

        _write_tool_shims(bin_dir, "/fake/lib")

        shim = bin_dir /"nvidia-smi"
        assert shim.exists()

    def test_shim_is_executable(self, bin_dir: Path) -> None:
        """Shim should have execute permission."""
        from specter_cli.venv import _write_tool_shims

        _write_tool_shims(bin_dir, "/fake/lib")

        shim = bin_dir /"nvidia-smi"
        assert shim.stat().st_mode & stat.S_IEXEC

    def test_shim_content_has_bash_shebang(self, bin_dir: Path) -> None:
        """Shim should start with #!/bin/bash."""
        from specter_cli.venv import _write_tool_shims

        _write_tool_shims(bin_dir, "/fake/lib")

        content = (bin_dir /"nvidia-smi").read_text()
        assert content.startswith("#!/bin/bash\n")

    def test_shim_calls_forward_module(self, bin_dir: Path) -> None:
        """Shim should invoke python -m specter_cli.forward <tool>."""
        from specter_cli.venv import _write_tool_shims

        _write_tool_shims(bin_dir, "/fake/lib")

        content = (bin_dir /"nvidia-smi").read_text()
        assert "specter_cli.forward" in content
        assert "nvidia-smi" in content

    def test_shim_sets_pythonpath(self, bin_dir: Path) -> None:
        """Shim should set PYTHONPATH for specter_cli."""
        from specter_cli.venv import _write_tool_shims

        _write_tool_shims(bin_dir, "/opt/specter/lib")

        content = (bin_dir /"nvidia-smi").read_text()
        assert "PYTHONPATH" in content
        assert "/opt/specter/lib" in content

    def test_shim_sets_virtual_env(self, bin_dir: Path) -> None:
        """Shim should set VIRTUAL_ENV."""
        from specter_cli.venv import _write_tool_shims

        _write_tool_shims(bin_dir, "/fake/lib")

        content = (bin_dir /"nvidia-smi").read_text()
        assert "VIRTUAL_ENV" in content

    def test_shim_uses_specter_original_python(self, bin_dir: Path) -> None:
        """Shim should use python.specter-original to avoid recursion."""
        from specter_cli.venv import _write_tool_shims

        _write_tool_shims(bin_dir, "/fake/lib")

        content = (bin_dir /"nvidia-smi").read_text()
        assert "python.specter-original" in content

    def test_creates_shim_for_each_tool(self, bin_dir: Path) -> None:
        """Should create one shim per tool in _FORWARDED_TOOLS."""
        from specter_cli.venv import _FORWARDED_TOOLS, _write_tool_shims

        _write_tool_shims(bin_dir, "/fake/lib")

        for tool in _FORWARDED_TOOLS:
            shim = bin_dir /tool
            assert shim.exists(), f"Missing shim for {tool}"
            assert shim.stat().st_mode & stat.S_IEXEC, f"{tool} not executable"


# ---------------------------------------------------------------------------
# Shim lifecycle: patch_venv creates, restore_venv removes
# ---------------------------------------------------------------------------


class TestToolShimLifecycle:
    """Tool shims should be created during patch_venv and removed during restore_venv."""

    @pytest.fixture
    def fake_venv(self, tmp_path: Path) -> Path:
        """Create a minimal fake venv."""
        venv = tmp_path / "fakevenv"
        bin_dir = venv / "bin"
        bin_dir.mkdir(parents=True)
        # Create a real python binary (needed for backup).
        python = bin_dir / "python"
        python.write_text("#!/bin/bash\necho fake python\n")
        python.chmod(0o755)
        python3 = bin_dir / "python3"
        python3.symlink_to("python")
        # Create a minimal activate script (needed for PATH patching).
        activate = bin_dir / "activate"
        activate.write_text(
            'VIRTUAL_ENV="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"\n'
            "export VIRTUAL_ENV\n"
            'export PATH="$VIRTUAL_ENV/bin:$PATH"\n'
        )
        return venv

    def test_patch_creates_nvidia_smi_shim(self, fake_venv: Path) -> None:
        """patch_venv should create a nvidia-smi shim."""
        from specter_cli.venv import patch_venv

        patch_venv(fake_venv, "h100")

        shim = fake_venv / "bin" /"nvidia-smi"
        assert shim.exists()
        assert shim.stat().st_mode & stat.S_IEXEC

    def test_patch_creates_all_forwarded_tool_shims(self, fake_venv: Path) -> None:
        """patch_venv should create shims for all _FORWARDED_TOOLS."""
        from specter_cli.venv import _FORWARDED_TOOLS, patch_venv

        patch_venv(fake_venv, "h100")

        for tool in _FORWARDED_TOOLS:
            shim = fake_venv / "bin" /tool
            assert shim.exists(), f"Missing shim for {tool}"

    def test_restore_removes_nvidia_smi_shim(self, fake_venv: Path) -> None:
        """restore_venv should remove the nvidia-smi shim."""
        from specter_cli.venv import patch_venv, restore_venv

        patch_venv(fake_venv, "h100")
        assert (fake_venv / "bin" /"nvidia-smi").exists()

        restore_venv(fake_venv)
        assert not (fake_venv / "bin" /"nvidia-smi").exists()

    def test_restore_removes_all_tool_shims(self, fake_venv: Path) -> None:
        """restore_venv should remove all tool shims."""
        from specter_cli.venv import _FORWARDED_TOOLS, patch_venv, restore_venv

        patch_venv(fake_venv, "h100")
        for tool in _FORWARDED_TOOLS:
            assert (fake_venv / "bin" /tool).exists()

        restore_venv(fake_venv)
        for tool in _FORWARDED_TOOLS:
            assert not (fake_venv / "bin" /tool).exists()

    def test_restore_tolerates_missing_shim(self, fake_venv: Path) -> None:
        """restore_venv should not crash if tool shim was already removed."""
        from specter_cli.venv import patch_venv, restore_venv

        patch_venv(fake_venv, "h100")

        # Manually remove the shim.
        (fake_venv / "bin" /"nvidia-smi").unlink()

        # Should not raise.
        restore_venv(fake_venv)

    def test_shim_content_matches_tool_name(self, fake_venv: Path) -> None:
        """Each shim's content should reference its specific tool name."""
        from specter_cli.venv import _FORWARDED_TOOLS, patch_venv

        patch_venv(fake_venv, "h100")

        for tool in _FORWARDED_TOOLS:
            content = (fake_venv / "bin" /tool).read_text()
            assert f"specter_cli.forward {tool}" in content

    def test_re_patch_recreates_shims(self, fake_venv: Path) -> None:
        """Re-patching (with restore in between) should recreate shims."""
        from specter_cli.venv import patch_venv

        patch_venv(fake_venv, "h100")
        # Re-patch triggers auto-restore + re-create.
        patch_venv(fake_venv, "a100")

        assert (fake_venv / "bin" /"nvidia-smi").exists()
        content = (fake_venv / "bin" /"nvidia-smi").read_text()
        assert "specter_cli.forward" in content


# ---------------------------------------------------------------------------
# _FORWARDED_TOOLS extensibility
# ---------------------------------------------------------------------------


class TestForwardedToolsConfig:
    """Verify that adding tools to _FORWARDED_TOOLS is all that's needed."""

    def test_forwarded_tools_contains_nvidia_smi(self) -> None:
        """nvidia-smi should be in the default list."""
        from specter_cli.venv import _FORWARDED_TOOLS

        assert "nvidia-smi" in _FORWARDED_TOOLS

    def test_forwarded_tools_is_a_list(self) -> None:
        """Should be a list (not tuple or set) for easy appending."""
        from specter_cli.venv import _FORWARDED_TOOLS

        assert isinstance(_FORWARDED_TOOLS, list)

    @pytest.fixture
    def bin_dir(self, tmp_path: Path) -> Path:
        d = tmp_path / "bin"
        d.mkdir()
        return d

    def test_adding_tool_creates_shim(self, bin_dir: Path) -> None:
        """Adding a tool to _FORWARDED_TOOLS should create its shim."""
        from specter_cli.venv import _FORWARDED_TOOLS, _write_tool_shims

        original = _FORWARDED_TOOLS.copy()
        try:
            _FORWARDED_TOOLS.append("nvcc")
            _write_tool_shims(bin_dir, "/fake/lib")

            assert (bin_dir /"nvcc").exists()
            content = (bin_dir /"nvcc").read_text()
            assert "specter_cli.forward nvcc" in content
        finally:
            _FORWARDED_TOOLS.clear()
            _FORWARDED_TOOLS.extend(original)
