"""Unit tests for :class:`flash_sandbox.Sandbox` and :class:`flash_sandbox.AsyncSandbox`.

All client interactions are mocked so the tests run without a live orchestrator.
"""

import sys
import os
from unittest.mock import MagicMock, AsyncMock

import pytest

# Ensure the package root is on sys.path so the import works from any CWD.
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from flash_sandbox.sandbox import Sandbox, AsyncSandbox
from flash_sandbox.http_client import ExecResult, MetricsResult


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_sync_client() -> MagicMock:
    """Return a mock that quacks like HTTPClient / SandboxClient."""
    client = MagicMock()
    # start_sandbox now returns a Sandbox object; we use a side_effect
    # so each call gets a fresh Sandbox bound to the client.
    client.start_sandbox.side_effect = lambda **kw: Sandbox("sb-100", client)
    client.get_status.return_value = "running"
    client.exec_command.return_value = ExecResult(stdout="hello\n", stderr="", exit_code=0)
    client.run_python.return_value = "42\n"
    client.get_platform_info.return_value = {"system": "Linux"}
    client.get_metrics.return_value = MetricsResult(
        memory_usage_bytes=1024,
        memory_limit_bytes=4096,
        memory_percent=25.0,
        cpu_percent=10.5,
        pids_current=2,
        net_rx_bytes=100,
        net_tx_bytes=200,
        block_read_bytes=300,
        block_write_bytes=400,
    )
    client.snapshot_sandbox.return_value = {
        "snapshot_path": "/snap/snap.bin",
        "mem_file_path": "/snap/mem.bin",
    }
    return client


def _make_async_client() -> MagicMock:
    """Return a mock that quacks like AsyncHTTPClient."""
    client = MagicMock()
    # start_sandbox now returns an AsyncSandbox object.
    async def _start_sandbox(**kw):
        return AsyncSandbox("sb-200", client)
    client.start_sandbox = AsyncMock(side_effect=_start_sandbox)
    client.get_status = AsyncMock(return_value="running")
    client.exec_command = AsyncMock(
        return_value=ExecResult(stdout="world\n", stderr="", exit_code=0)
    )
    client.run_python = AsyncMock(return_value="99\n")
    client.get_platform_info = AsyncMock(return_value={"system": "Linux"})
    client.get_metrics = AsyncMock(
        return_value=MetricsResult(
            memory_usage_bytes=2048,
            memory_limit_bytes=8192,
            memory_percent=25.0,
            cpu_percent=5.0,
            pids_current=3,
            net_rx_bytes=10,
            net_tx_bytes=20,
            block_read_bytes=30,
            block_write_bytes=40,
        )
    )
    client.snapshot_sandbox = AsyncMock(
        return_value={
            "snapshot_path": "/snap/async.bin",
            "mem_file_path": "/snap/async_mem.bin",
        }
    )
    client.resume_sandbox = AsyncMock()
    client.stop_sandbox = AsyncMock()
    return client


# ===========================================================================
# Sandbox (sync)
# ===========================================================================


class TestSandboxInit:
    def test_init_sets_id_and_client(self):
        client = MagicMock()
        sb = Sandbox("sb-42", client)

        assert sb.id == "sb-42"
        assert sb._client is client

    def test_repr(self):
        sb = Sandbox("sb-42", MagicMock())
        assert repr(sb) == "Sandbox(id='sb-42')"


class TestSandboxCreate:
    def test_basic(self):
        client = _make_sync_client()
        sb = Sandbox.create(
            client,
            type="docker",
            image="alpine:latest",
            command=["sleep", "3600"],
            memory_mb=256,
            cpu_cores=0.5,
        )

        assert isinstance(sb, Sandbox)
        assert sb.id == "sb-100"
        # The returned Sandbox is the one produced by start_sandbox,
        # which is already bound to the client.
        assert sb._client is client
        client.start_sandbox.assert_called_once_with(
            type="docker",
            image="alpine:latest",
            command=["sleep", "3600"],
            memory_mb=256,
            cpu_cores=0.5,
        )

    def test_defaults(self):
        client = _make_sync_client()
        sb = Sandbox.create(client, type="docker", image="ubuntu:22.04")

        assert sb.id == "sb-100"
        client.start_sandbox.assert_called_once_with(
            type="docker",
            image="ubuntu:22.04",
            command=None,
            memory_mb=512,
            cpu_cores=1.0,
        )

    def test_extra_kwargs_forwarded(self):
        client = _make_sync_client()
        Sandbox.create(
            client,
            type="firecracker",
            image="rootfs.ext4",
            kernel_image="/boot/vmlinux",
            snapshot_path="/snaps/snap.bin",
        )

        _, kwargs = client.start_sandbox.call_args
        assert kwargs["kernel_image"] == "/boot/vmlinux"
        assert kwargs["snapshot_path"] == "/snaps/snap.bin"

    def test_works_with_string_returning_client(self):
        """Clients that return a plain string ID (e.g. gRPC) still work."""
        client = MagicMock()
        client.start_sandbox.return_value = "sb-grpc"
        sb = Sandbox.create(client, type="docker", image="alpine:latest")

        assert isinstance(sb, Sandbox)
        assert sb.id == "sb-grpc"
        assert sb._client is client


class TestSandboxOperations:
    def _make_sandbox(self) -> tuple:
        client = _make_sync_client()
        sb = Sandbox("sb-100", client)
        return sb, client

    def test_get_status(self):
        sb, client = self._make_sandbox()
        status = sb.get_status()

        assert status == "running"
        client.get_status.assert_called_once_with("sb-100")

    def test_exec_command(self):
        sb, client = self._make_sandbox()
        result = sb.exec_command(["echo", "hello"])

        assert result.stdout == "hello\n"
        assert result.exit_code == 0
        client.exec_command.assert_called_once_with("sb-100", ["echo", "hello"])

    def test_run_python(self):
        sb, client = self._make_sandbox()
        result = sb.run_python("print(6 * 7)")

        assert result == "42\n"
        client.run_python.assert_called_once_with("sb-100", "print(6 * 7)")

    def test_get_platform_info(self):
        sb, client = self._make_sandbox()
        info = sb.get_platform_info()

        assert info == {"system": "Linux"}
        assert isinstance(info, dict)
        client.get_platform_info.assert_called_once_with("sb-100")

    def test_get_metrics(self):
        sb, client = self._make_sandbox()
        metrics = sb.get_metrics()

        assert metrics.memory_usage_bytes == 1024
        assert metrics.cpu_percent == pytest.approx(10.5)
        client.get_metrics.assert_called_once_with("sb-100")

    def test_snapshot(self):
        sb, client = self._make_sandbox()
        snap = sb.snapshot()

        assert snap["snapshot_path"] == "/snap/snap.bin"
        assert snap["mem_file_path"] == "/snap/mem.bin"
        client.snapshot_sandbox.assert_called_once_with("sb-100")

    def test_resume(self):
        sb, client = self._make_sandbox()
        sb.resume()

        client.resume_sandbox.assert_called_once_with("sb-100")

    def test_stop(self):
        sb, client = self._make_sandbox()
        sb.stop()

        client.stop_sandbox.assert_called_once_with("sb-100", cleanup=True)

    def test_stop_without_cleanup(self):
        sb, client = self._make_sandbox()
        sb.stop(cleanup=False)

        client.stop_sandbox.assert_called_once_with("sb-100", cleanup=False)

    def test_stop_with_explicit_cleanup(self):
        sb, client = self._make_sandbox()
        sb.stop(cleanup=True)

        client.stop_sandbox.assert_called_once_with("sb-100", cleanup=True)


class TestSandboxContextManager:
    def test_context_manager_stops_on_exit(self):
        client = _make_sync_client()
        sb = Sandbox("sb-100", client)

        with sb:
            pass

        client.stop_sandbox.assert_called_once_with("sb-100", cleanup=True)

    def test_context_manager_stops_on_exception(self):
        client = _make_sync_client()
        sb = Sandbox("sb-100", client)

        with pytest.raises(RuntimeError):
            with sb:
                raise RuntimeError("boom")

        client.stop_sandbox.assert_called_once_with("sb-100", cleanup=True)

    def test_context_manager_returns_sandbox(self):
        client = _make_sync_client()
        sb = Sandbox("sb-100", client)

        with sb as ctx:
            assert ctx is sb


class TestSandboxFullLifecycle:
    """Simulate a complete create → use → stop flow."""

    def test_lifecycle(self):
        client = _make_sync_client()

        with Sandbox.create(client, type="docker", image="alpine:latest") as sb:
            assert sb.id == "sb-100"

            status = sb.get_status()
            assert status == "running"

            result = sb.exec_command(["echo", "hi"])
            assert result.stdout == "hello\n"

            py_result = sb.run_python("print(42)")
            assert py_result == "42\n"

            info = sb.get_platform_info()
            assert info.get("system") == "Linux"

            metrics = sb.get_metrics()
            assert metrics.memory_usage_bytes > 0

            snap = sb.snapshot()
            assert "snapshot_path" in snap

            sb.resume()

        # Exiting the context manager should stop the sandbox.
        client.stop_sandbox.assert_called_once_with("sb-100", cleanup=True)


# ===========================================================================
# AsyncSandbox
# ===========================================================================


class TestAsyncSandboxInit:
    def test_init_sets_id_and_client(self):
        client = MagicMock()
        sb = AsyncSandbox("sb-42", client)

        assert sb.id == "sb-42"
        assert sb._client is client

    def test_repr(self):
        sb = AsyncSandbox("sb-42", MagicMock())
        assert repr(sb) == "AsyncSandbox(id='sb-42')"


@pytest.mark.asyncio
class TestAsyncSandboxCreate:
    async def test_basic(self):
        client = _make_async_client()
        sb = await AsyncSandbox.create(
            client,
            type="docker",
            image="alpine:latest",
            command=["sleep", "3600"],
            memory_mb=256,
            cpu_cores=0.5,
        )

        assert isinstance(sb, AsyncSandbox)
        assert sb.id == "sb-200"
        assert sb._client is client
        client.start_sandbox.assert_awaited_once_with(
            type="docker",
            image="alpine:latest",
            command=["sleep", "3600"],
            memory_mb=256,
            cpu_cores=0.5,
        )

    async def test_defaults(self):
        client = _make_async_client()
        sb = await AsyncSandbox.create(client, type="docker", image="ubuntu:22.04")

        assert sb.id == "sb-200"
        client.start_sandbox.assert_awaited_once_with(
            type="docker",
            image="ubuntu:22.04",
            command=None,
            memory_mb=512,
            cpu_cores=1.0,
        )

    async def test_extra_kwargs_forwarded(self):
        client = _make_async_client()
        await AsyncSandbox.create(
            client,
            type="firecracker",
            image="rootfs.ext4",
            kernel_image="/boot/vmlinux",
        )

        _, kwargs = client.start_sandbox.call_args
        assert kwargs["kernel_image"] == "/boot/vmlinux"

    async def test_works_with_string_returning_client(self):
        """Clients that return a plain string ID (e.g. gRPC) still work."""
        client = MagicMock()
        client.start_sandbox = AsyncMock(return_value="sb-grpc")
        sb = await AsyncSandbox.create(client, type="docker", image="alpine:latest")

        assert isinstance(sb, AsyncSandbox)
        assert sb.id == "sb-grpc"
        assert sb._client is client


@pytest.mark.asyncio
class TestAsyncSandboxOperations:
    def _make_sandbox(self) -> tuple:
        client = _make_async_client()
        sb = AsyncSandbox("sb-200", client)
        return sb, client

    async def test_get_status(self):
        sb, client = self._make_sandbox()
        status = await sb.get_status()

        assert status == "running"
        client.get_status.assert_awaited_once_with("sb-200")

    async def test_exec_command(self):
        sb, client = self._make_sandbox()
        result = await sb.exec_command(["echo", "world"])

        assert result.stdout == "world\n"
        assert result.exit_code == 0
        client.exec_command.assert_awaited_once_with("sb-200", ["echo", "world"])

    async def test_run_python(self):
        sb, client = self._make_sandbox()
        result = await sb.run_python("print(99)")

        assert result == "99\n"
        client.run_python.assert_awaited_once_with("sb-200", "print(99)")

    async def test_get_platform_info(self):
        sb, client = self._make_sandbox()
        info = await sb.get_platform_info()

        assert info == {"system": "Linux"}
        assert isinstance(info, dict)
        client.get_platform_info.assert_awaited_once_with("sb-200")

    async def test_get_metrics(self):
        sb, client = self._make_sandbox()
        metrics = await sb.get_metrics()

        assert metrics.memory_usage_bytes == 2048
        assert metrics.cpu_percent == pytest.approx(5.0)
        client.get_metrics.assert_awaited_once_with("sb-200")

    async def test_snapshot(self):
        sb, client = self._make_sandbox()
        snap = await sb.snapshot()

        assert snap["snapshot_path"] == "/snap/async.bin"
        assert snap["mem_file_path"] == "/snap/async_mem.bin"
        client.snapshot_sandbox.assert_awaited_once_with("sb-200")

    async def test_resume(self):
        sb, client = self._make_sandbox()
        await sb.resume()

        client.resume_sandbox.assert_awaited_once_with("sb-200")

    async def test_stop(self):
        sb, client = self._make_sandbox()
        await sb.stop()

        client.stop_sandbox.assert_awaited_once_with("sb-200", cleanup=True)

    async def test_stop_without_cleanup(self):
        sb, client = self._make_sandbox()
        await sb.stop(cleanup=False)

        client.stop_sandbox.assert_awaited_once_with("sb-200", cleanup=False)

    async def test_stop_with_explicit_cleanup(self):
        sb, client = self._make_sandbox()
        await sb.stop(cleanup=True)

        client.stop_sandbox.assert_awaited_once_with("sb-200", cleanup=True)


@pytest.mark.asyncio
class TestAsyncSandboxContextManager:
    async def test_context_manager_stops_on_exit(self):
        client = _make_async_client()
        sb = AsyncSandbox("sb-200", client)

        async with sb:
            pass

        client.stop_sandbox.assert_awaited_once_with("sb-200", cleanup=True)

    async def test_context_manager_stops_on_exception(self):
        client = _make_async_client()
        sb = AsyncSandbox("sb-200", client)

        with pytest.raises(RuntimeError):
            async with sb:
                raise RuntimeError("boom")

        client.stop_sandbox.assert_awaited_once_with("sb-200", cleanup=True)

    async def test_context_manager_returns_sandbox(self):
        client = _make_async_client()
        sb = AsyncSandbox("sb-200", client)

        async with sb as ctx:
            assert ctx is sb


@pytest.mark.asyncio
class TestAsyncSandboxFullLifecycle:
    """Simulate a complete create → use → stop flow."""

    async def test_lifecycle(self):
        client = _make_async_client()

        async with await AsyncSandbox.create(
            client, type="docker", image="alpine:latest"
        ) as sb:
            assert sb.id == "sb-200"

            status = await sb.get_status()
            assert status == "running"

            result = await sb.exec_command(["echo", "hi"])
            assert result.stdout == "world\n"

            py_result = await sb.run_python("print(99)")
            assert py_result == "99\n"

            info = await sb.get_platform_info()
            assert info.get("system") == "Linux"

            metrics = await sb.get_metrics()
            assert metrics.memory_usage_bytes > 0

            snap = await sb.snapshot()
            assert "snapshot_path" in snap

            await sb.resume()

        # Exiting the async context manager should stop the sandbox.
        client.stop_sandbox.assert_awaited_once_with("sb-200", cleanup=True)
