from __future__ import annotations

import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest import mock


def _install_test_stubs() -> None:
    if "paramiko" not in sys.modules:
        paramiko = types.ModuleType("paramiko")

        class _BaseKey:
            @classmethod
            def from_private_key_file(cls, path: str, password: str | None = None) -> object:
                return object()

        class SSHException(Exception):
            pass

        class AuthenticationException(Exception):
            pass

        class AutoAddPolicy:
            pass

        class SSHClient:
            def set_missing_host_key_policy(self, policy: object) -> None:
                return None

            def connect(self, **kwargs: object) -> None:
                return None

            def get_transport(self) -> object | None:
                return None

            def close(self) -> None:
                return None

        paramiko.PKey = object
        paramiko.SSHException = SSHException
        paramiko.AuthenticationException = AuthenticationException
        paramiko.AutoAddPolicy = AutoAddPolicy
        paramiko.SSHClient = SSHClient
        paramiko.Ed25519Key = _BaseKey
        paramiko.RSAKey = _BaseKey
        paramiko.ECDSAKey = _BaseKey
        sys.modules["paramiko"] = paramiko

    if "mcp.server.fastmcp" not in sys.modules:
        mcp_module = types.ModuleType("mcp")
        mcp_server_module = types.ModuleType("mcp.server")
        fastmcp_module = types.ModuleType("mcp.server.fastmcp")

        class FastMCP:
            def __init__(self, *args: object, **kwargs: object) -> None:
                return None

            def tool(self, *args: object, **kwargs: object):
                def decorator(fn):
                    return fn

                return decorator

            def run(self, *args: object, **kwargs: object) -> None:
                return None

        fastmcp_module.FastMCP = FastMCP
        sys.modules["mcp"] = mcp_module
        sys.modules["mcp.server"] = mcp_server_module
        sys.modules["mcp.server.fastmcp"] = fastmcp_module


_install_test_stubs()

from remote_sandbox_mcp import server
from remote_sandbox_mcp import watchdog_daemon
from remote_sandbox_mcp.watchdog_daemon import WatchdogDaemon
from remote_sandbox_mcp.watchdog_store import create_watch, get_watch, init_db, save_sandbox_config


class ServerWatchdogTests(unittest.TestCase):
    def test_ensure_connected_reconnects_inactive_transport_immediately(self) -> None:
        session = server.SandboxSession(
            server.SandboxConfig(name="gpu-a", host="example", user="ubuntu", password="secret")
        )
        stale_transport = mock.Mock()
        stale_transport.is_active.return_value = False
        stale_client = mock.Mock()
        stale_client.get_transport.return_value = stale_transport
        fresh_client = mock.Mock()

        session._client = stale_client
        session._last_health_check = 10_000.0

        with mock.patch.object(server.SandboxSession, "_make_client", return_value=fresh_client):
            client = session.ensure_connected()

        self.assertIs(client, fresh_client)
        stale_client.close.assert_called_once()

    def test_exec_on_channel_reconnects_once_before_opening_a_channel(self) -> None:
        session = mock.Mock()
        stale_transport = mock.Mock()
        stale_transport.is_active.return_value = False
        stale_client = mock.Mock()
        stale_client.get_transport.return_value = stale_transport
        setattr(stale_client, "_rsmcp_session", session)

        fresh_channel = mock.Mock()
        fresh_channel.recv_ready.return_value = False
        fresh_channel.recv_stderr_ready.return_value = False
        fresh_channel.exit_status_ready.return_value = True
        fresh_channel.recv_exit_status.return_value = 0

        fresh_transport = mock.Mock()
        fresh_transport.is_active.return_value = True
        fresh_transport.open_session.return_value = fresh_channel

        fresh_client = mock.Mock()
        fresh_client.get_transport.return_value = fresh_transport
        setattr(fresh_client, "_rsmcp_session", session)
        session.reconnect.return_value = fresh_client

        result = server._exec_on_channel(stale_client, "echo ok", timeout_s=5, max_output_chars=200)

        self.assertEqual(result["exit_code"], 0)
        session.reconnect.assert_called_once()
        fresh_transport.open_session.assert_called_once()

    def test_exec_bash_background_returns_watch_id(self) -> None:
        fake_session = mock.Mock()
        fake_session.ensure_connected.return_value = object()
        fake_session.id = "gpu-a"

        with (
            mock.patch.object(server, "_get_session", return_value=fake_session),
            mock.patch.object(server, "_exec_on_channel", return_value={"exit_code": 0, "stdout": "", "stderr": ""}),
            mock.patch.object(server, "_ensure_watchdog_ready", return_value={"ok": True}),
            mock.patch.object(server, "_remote_home", return_value="/home/ubuntu"),
            mock.patch.object(
                server,
                "watch_background_task",
                return_value={"watch": {"id": 42, "tmux_session": "bg-123", "log_file": ".codex_logs/bg-123.log"}},
            ) as watch_background_task,
            mock.patch.object(server.time, "time", return_value=123),
            mock.patch.object(server.sys, "platform", "darwin"),
        ):
            result = server.exec_bash_background(command="sleep 5")

        self.assertEqual(result["status"], "started")
        self.assertEqual(result["tmux_session"], "bg-123")
        self.assertEqual(result["watch_id"], 42)
        self.assertEqual(result["watch_query_id"], 42)
        self.assertIn("check_background_task(watch_id=42)", result["note"])
        watch_background_task.assert_called_once()
        self.assertEqual(watch_background_task.call_args.kwargs["ensure_watchdog"], False)

    def test_exec_bash_background_builds_remote_script_and_resolves_default_log_path_under_cwd(self) -> None:
        fake_session = mock.Mock()
        fake_session.ensure_connected.return_value = object()
        fake_session.id = "gpu-a"
        exec_calls: list[str] = []

        def fake_exec(_client: object, command: str, **kwargs: object) -> dict:
            exec_calls.append(command)
            return {"exit_code": 0, "stdout": "", "stderr": ""}

        with (
            mock.patch.object(server, "_get_session", return_value=fake_session),
            mock.patch.object(server, "_exec_on_channel", side_effect=fake_exec),
            mock.patch.object(server, "_ensure_watchdog_ready", return_value={"ok": True}),
            mock.patch.object(server, "_remote_home", return_value="/home/ubuntu"),
            mock.patch.object(
                server,
                "watch_background_task",
                return_value={
                    "watch": {
                        "id": 99,
                        "tmux_session": "demo",
                        "log_file": "/home/ubuntu/project/.codex_logs/demo.log",
                    }
                },
            ) as watch_background_task,
            mock.patch.object(server.sys, "platform", "darwin"),
        ):
            result = server.exec_bash_background(
                command="python3 scripts/job.py",
                session_name="demo",
                cwd="project",
            )

        self.assertEqual(result["log_file"], "/home/ubuntu/project/.codex_logs/demo.log")
        self.assertEqual(
            watch_background_task.call_args.kwargs["log_file"],
            "/home/ubuntu/project/.codex_logs/demo.log",
        )
        self.assertEqual(watch_background_task.call_args.kwargs["cwd"], "/home/ubuntu/project")
        self.assertIn("tmux new-session -d -s demo", exec_calls[0])
        self.assertIn("cat > /home/ubuntu/project/.codex_logs/.rsmcp/demo.sh", exec_calls[0])
        self.assertIn("bash /home/ubuntu/project/.codex_logs/.rsmcp/demo.sh", exec_calls[0])
        self.assertIn("cd /home/ubuntu/project", exec_calls[0])

    def test_exec_bash_background_recovers_if_tmux_launch_times_out_but_session_is_running(self) -> None:
        fake_session = mock.Mock()
        fake_session.ensure_connected.return_value = object()
        fake_session.id = "gpu-a"

        with (
            mock.patch.object(server, "_get_session", return_value=fake_session),
            mock.patch.object(
                server,
                "_exec_on_channel",
                return_value={
                    "exit_code": -1,
                    "stdout": "",
                    "stderr": "timeout",
                    "timed_out": True,
                },
            ),
            mock.patch.object(server, "_ensure_watchdog_ready", return_value={"ok": True}),
            mock.patch.object(server, "_remote_home", return_value="/home/ubuntu"),
            mock.patch.object(
                server,
                "check_background_task",
                return_value={
                    "running": True,
                    "exit_code": None,
                    "log_tail": "starting demo task",
                },
            ),
            mock.patch.object(
                server,
                "watch_background_task",
                return_value={
                    "watch": {
                        "id": 55,
                        "tmux_session": "demo",
                        "log_file": "/home/ubuntu/project/logs/demo.log",
                    }
                },
            ),
            mock.patch.object(server.sys, "platform", "darwin"),
        ):
            result = server.exec_bash_background(
                command="python3 scripts/job.py",
                session_name="demo",
                cwd="project",
                log_file="logs/demo.log",
            )

        self.assertEqual(result["status"], "started")
        self.assertTrue(result["launch_recovered"])
        self.assertEqual(result["watch_id"], 55)

    def test_check_background_task_accepts_watch_id(self) -> None:
        fake_session = mock.Mock()
        fake_session.ensure_connected.return_value = object()
        fake_session.id = "gpu-a"

        with (
            mock.patch.object(server, "init_watchdog_db"),
            mock.patch.object(
                server,
                "get_watch",
                return_value={
                    "id": 7,
                    "tmux_session": "job-7",
                    "log_file": "/tmp/job-7.log",
                    "sandbox_name": "gpu-a",
                },
            ),
            mock.patch.object(server, "_get_session", return_value=fake_session),
            mock.patch.object(
                server,
                "_exec_on_channel",
                side_effect=[
                    {"exit_code": 0, "stdout": "RUNNING\n", "stderr": ""},
                    {"exit_code": 0, "stdout": "line1\nEXIT_CODE=0\n", "stderr": ""},
                ],
            ),
        ):
            result = server.check_background_task(watch_id=7, last_n_lines=10)

        self.assertTrue(result["running"])
        self.assertEqual(result["watch_id"], 7)
        self.assertEqual(result["tmux_session"], "job-7")
        self.assertEqual(result["log_file"], "/tmp/job-7.log")
        self.assertEqual(result["watch"]["id"], 7)


class WatchdogDaemonStateTests(unittest.TestCase):
    def test_daemon_start_background_task_uses_absolute_remote_paths(self) -> None:
        fake_session = mock.Mock()
        fake_session.ensure_connected.return_value = object()
        exec_calls: list[str] = []

        def fake_exec(_client: object, command: str, **kwargs: object) -> dict:
            exec_calls.append(command)
            if 'printf "%s" "$HOME"' in command:
                return {"exit_code": 0, "stdout": "/home/ubuntu", "stderr": ""}
            return {"exit_code": 0, "stdout": "", "stderr": ""}

        with mock.patch.object(watchdog_daemon, "_exec_on_channel", side_effect=fake_exec):
            result = watchdog_daemon._start_background_task(
                fake_session,
                "python3 scripts/job.py --resume",
                "project",
                "job-1",
                "project/logs/job-1.log",
            )

        self.assertTrue(result["started"])
        self.assertIn("cat > /home/ubuntu/project/logs/.rsmcp/job-1.sh", exec_calls[-1])
        self.assertIn("bash /home/ubuntu/project/logs/.rsmcp/job-1.sh", exec_calls[-1])
        self.assertIn("cd /home/ubuntu/project", exec_calls[-1])
        self.assertIn("tee -a logs/job-1.log", exec_calls[-1])

    def test_daemon_alerts_after_threshold_and_resumes_with_same_tmux_session(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            tmp = Path(tmpdir)
            config_path = tmp / "config.json"
            db_path = tmp / "watchdog.sqlite3"

            save_sandbox_config(
                [{"name": "gpu-a", "host": "example", "user": "ubuntu", "password": "secret"}],
                path=str(config_path),
            )
            init_db(str(db_path))
            watch = create_watch(
                "gpu-a",
                "job-1",
                "/tmp/job-1.log",
                run_id="run-1",
                launch_command="python task.py",
                resume_command="python resume.py",
                checkpoint_path="/tmp/checkpoint.json",
                interval_s=300,
                max_log_lines=20,
                auto_resume=True,
                max_resume_attempts=2,
                alert_after_failures=2,
                notify_local=False,
                resume_delay_s=300,
                path=str(db_path),
            )

            daemon = WatchdogDaemon(config_path=str(config_path), db_path=str(db_path), loop_sleep_s=1)
            fake_session = mock.Mock()
            fake_session.ensure_connected.side_effect = [
                ConnectionError("network down"),
                ConnectionError("network down"),
                object(),
                object(),
            ]
            daemon.sessions.get = mock.Mock(return_value=fake_session)

            events: list[tuple[str, str]] = []

            with (
                mock.patch(
                    "remote_sandbox_mcp.watchdog_daemon._emit_state_event",
                    side_effect=lambda watch, event_type, summary, **kwargs: events.append((event_type, summary)),
                ),
                mock.patch(
                    "remote_sandbox_mcp.watchdog_daemon._read_checkpoint",
                    return_value={"text": "", "digest": "", "parsed": None},
                ),
                mock.patch(
                    "remote_sandbox_mcp.watchdog_daemon._check_task",
                    side_effect=[
                        {"running": False, "exit_code": None, "log_tail": "", "log_digest": "a"},
                        {"running": False, "exit_code": None, "log_tail": "", "log_digest": "b"},
                    ],
                ),
                mock.patch(
                    "remote_sandbox_mcp.watchdog_daemon._start_background_task",
                    return_value={"started": True},
                ) as start_background_task,
            ):
                daemon._process_watch(get_watch(watch["id"], path=str(db_path)), daemon._load_configs(), 1000)
                current = get_watch(watch["id"], path=str(db_path))
                self.assertEqual(current["consecutive_ssh_failures"], 1)
                self.assertFalse(current["ssh_alert_sent"])
                self.assertEqual(events, [])

                daemon._process_watch(get_watch(watch["id"], path=str(db_path)), daemon._load_configs(), 1300)
                current = get_watch(watch["id"], path=str(db_path))
                self.assertEqual(current["consecutive_ssh_failures"], 2)
                self.assertTrue(current["ssh_alert_sent"])
                self.assertEqual(events[0][0], "ssh_unreachable")

                daemon._process_watch(get_watch(watch["id"], path=str(db_path)), daemon._load_configs(), 1600)
                current = get_watch(watch["id"], path=str(db_path))
                self.assertEqual(current["last_state"], "interrupted")
                self.assertEqual(current["resume_attempts"], 0)
                self.assertEqual(current["resume_due_ts"], 1900)
                self.assertEqual(events[1][0], "interrupted")

                daemon._process_watch(get_watch(watch["id"], path=str(db_path)), daemon._load_configs(), 1901)
                current = get_watch(watch["id"], path=str(db_path))
                self.assertEqual(current["last_state"], "running")
                self.assertEqual(current["resume_attempts"], 1)
                self.assertEqual(events[2][0], "resume_started")
                self.assertEqual(start_background_task.call_args.args[3], "job-1")


if __name__ == "__main__":
    unittest.main()
