"""Sandbox execution engine — orchestrates sandboxed MCP tool execution."""

from __future__ import annotations

import asyncio
import json
import logging
import signal
import time
from pathlib import Path
from typing import Any

from hatchdx.sandbox.config import (
    generate_dockerfile,
    policy_to_container_config,
)
from hatchdx.utils.jsonrpc import (
    build_initialize_request,
    build_initialized_notification,
    build_tool_call_request,
)
from hatchdx.sandbox.models import ContainerStats, SandboxPolicy, SandboxResult
from hatchdx.sandbox.policy import PolicyError, load_policy, load_preset
from hatchdx.sandbox.runtime import ContainerRuntime, SandboxError, detect_runtime

logger = logging.getLogger(__name__)


class SandboxEngine:
    """Orchestrate the full sandbox lifecycle: build, run, and cleanup.

    Usage::

        engine = SandboxEngine(project_path="/path/to/mcp-server")
        result = await engine.run_tool("my_tool", {"arg": "value"})
        engine.cleanup()

    Or as a context manager::

        async with SandboxEngine(project_path="/path/to/mcp-server") as engine:
            result = await engine.run_tool("my_tool", {"arg": "value"})
    """

    def __init__(
        self,
        project_path: str,
        *,
        server_command: str = "python -m server",
        policy: SandboxPolicy | None = None,
        policy_path: str | None = None,
        preset: str | None = None,
        runtime: ContainerRuntime | None = None,
        force_build: bool = False,
        keep_container: bool = False,
        timeout_override: int | None = None,
        image_tag: str | None = None,
    ) -> None:
        self._project_path = str(Path(project_path).resolve())
        self._server_command = server_command
        self._force_build = force_build
        self._keep_container = keep_container
        self._container_id: str | None = None
        self._image_id: str | None = None
        self._started = False
        self._network_enforced = False

        # Resolve the runtime
        if runtime is not None:
            self._runtime = runtime
        else:
            self._runtime = detect_runtime()

        # Resolve the policy — explicit policy object > policy file > preset > standard
        if policy is not None:
            self._policy = policy
        elif policy_path is not None:
            self._policy = load_policy(policy_path)
        elif preset is not None:
            self._policy = load_preset(preset)
        else:
            self._policy = load_preset("standard")

        # Apply timeout override
        if timeout_override is not None:
            self._policy.resources.timeout_seconds = timeout_override

        # Image tag
        project_name = Path(self._project_path).name.lower().replace(" ", "-")
        self._image_tag = image_tag or f"hdx-sandbox-{project_name}:latest"

    # -- Properties -----------------------------------------------------------

    @property
    def policy(self) -> SandboxPolicy:
        """The active sandbox policy."""
        return self._policy

    @property
    def container_id(self) -> str | None:
        """The ID of the current container, or None if not started."""
        return self._container_id

    @property
    def image_tag(self) -> str:
        """The container image tag being used."""
        return self._image_tag

    @property
    def network_enforced(self) -> bool:
        """Whether domain allowlist filtering is actively enforced via iptables."""
        return self._network_enforced

    # -- Context manager ------------------------------------------------------

    async def __aenter__(self) -> SandboxEngine:
        await self.start()
        return self

    async def __aexit__(self, *exc: object) -> None:
        self.cleanup()

    # -- Lifecycle ------------------------------------------------------------

    async def start(self) -> None:
        """Build the image and start the sandbox container.

        Raises ``SandboxError`` if the runtime is unavailable or the build
        fails.
        """
        if self._started:
            return

        if not self._runtime.is_available():
            raise SandboxError(
                "Container runtime is not available.\n\n"
                "Make sure your container runtime (docker/podman/nerdctl) is installed\n"
                "and the daemon is running."
            )

        # Build image
        self._image_id = await self.build_image()

        # Create and start container
        config = policy_to_container_config(
            self._policy,
            image=self._image_tag,
            project_path=self._project_path,
        )

        self._container_id = self._runtime.create_container(
            self._image_tag, config
        )
        self._runtime.start_container(self._container_id)
        self._started = True

        # Apply firewall rules as root if domain filtering is configured.
        # The firewall script is baked into the image; we exec it as root
        # because iptables requires root privileges.  The container's default
        # user is hdx (non-root), so all subsequent exec calls — including
        # tool execution — run unprivileged and cannot modify the rules.
        if self._policy.network.enabled and self._policy.network.allowed_domains:
            self._apply_firewall()

        # Register signal handlers for clean shutdown
        self._register_signal_handlers()

    async def build_image(self) -> str:
        """Build the container image, generating a Dockerfile if needed."""
        allowed = (
            self._policy.network.allowed_domains
            if self._policy.network.enabled and self._policy.network.allowed_domains
            else None
        )
        dockerfile_content = generate_dockerfile(
            self._project_path, allowed_domains=allowed
        )

        return self._runtime.build_image(
            self._project_path,
            self._image_tag,
            dockerfile=dockerfile_content,
        )

    def cleanup(self) -> None:
        """Stop and remove the sandbox container.

        Safe to call multiple times — subsequent calls are no-ops.
        """
        if self._container_id is None:
            return

        container_id = self._container_id
        self._container_id = None
        self._started = False

        self._runtime.stop_container(
            container_id, remove=not self._keep_container
        )
        self._restore_signal_handlers()

    # -- Tool execution -------------------------------------------------------

    def _get_runtime_binary(self) -> str:
        """Return the runtime binary path for building subprocess commands."""
        if hasattr(self._runtime, "binary"):
            return self._runtime.binary
        return "docker"

    async def run_tool(
        self,
        tool_name: str,
        arguments: dict[str, Any] | None = None,
    ) -> SandboxResult:
        """Execute a single MCP tool call inside the sandbox.

        Spawns an async subprocess via ``{runtime} exec -i`` and communicates
        with the MCP server using JSON-RPC over stdin/stdout.  The stdin pipe
        is kept open during the entire interaction so the server's stdio
        transport doesn't tear down prematurely.

        Returns a ``SandboxResult`` with the tool output and resource stats.
        """
        if not self._started or self._container_id is None:
            raise SandboxError(
                "Sandbox is not running. Call start() first or use as context manager."
            )

        arguments = arguments or {}
        start_time = time.monotonic()
        timeout = self._policy.resources.timeout_seconds

        # Build the exec command: {binary} exec -i {container_id} sh -c {server_command}
        binary = self._get_runtime_binary()
        exec_cmd = [
            binary, "exec", "-i", self._container_id,
            "sh", "-c", self._server_command,
        ]

        try:
            proc = await asyncio.create_subprocess_exec(
                *exec_cmd,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
        except FileNotFoundError:
            raise SandboxError(
                f"Container runtime binary not found: {binary}\n\n"
                "Make sure docker, podman, or nerdctl is installed and on your PATH."
            )

        try:
            # Step 1: Send initialize request, wait for response
            init_message = build_initialize_request(
                request_id=1, client_name="hdx-sandbox",
            )
            await self._send_line(proc, init_message, timeout)
            await self._read_response(proc, expected_id=1, timeout=timeout)

            # Step 2: Send initialized notification (no response expected)
            initialized_notification = build_initialized_notification()
            await self._send_line(proc, initialized_notification, timeout)

            # Step 3: Send tool call request, wait for response
            call_message = build_tool_call_request(tool_name, arguments, request_id=2)
            await self._send_line(proc, call_message, timeout)
            response = await self._read_response(proc, expected_id=2, timeout=timeout)

            elapsed_ms = (time.monotonic() - start_time) * 1000

            # Parse the tool call result
            if "error" in response:
                err = response["error"]
                content = [{"type": "text", "text": err.get("message", str(err))}]
                is_error = True
            else:
                result = response.get("result", {})
                content = result.get("content", [])
                is_error = result.get("isError", False)

        except asyncio.TimeoutError:
            raise SandboxError(
                f"Tool execution timed out after {timeout}s. "
                "Increase timeout with --timeout or in your sandbox policy."
            )
        finally:
            # Close stdin to signal the server to shut down, then clean up
            await _shutdown_process(proc)

        # Collect resource stats
        stats = self._get_stats_safe()

        return SandboxResult(
            tool_name=tool_name,
            arguments=arguments,
            content=content,
            is_error=is_error,
            latency_ms=elapsed_ms,
            container_id=self._container_id[:12] if self._container_id else "",
            stats=stats,
        )

    async def exec_shell(self) -> tuple[int, str]:
        """Start an interactive shell session inside the container.

        Returns ``(exit_code, output)`` when the shell exits.
        """
        if not self._started or self._container_id is None:
            raise SandboxError(
                "Sandbox is not running. Call start() first or use as context manager."
            )

        exit_code = self._runtime.exec_interactive(
            self._container_id,
            ["/bin/sh"],
        )
        return exit_code, ""

    def get_stats(self) -> ContainerStats | None:
        """Return current resource usage, or None if unavailable."""
        return self._get_stats_safe()

    # -- Async stdin/stdout helpers -------------------------------------------

    @staticmethod
    async def _send_line(
        proc: asyncio.subprocess.Process,
        message: str,
        timeout: float,
    ) -> None:
        """Write a single JSON-RPC message line to the process stdin."""
        if proc.stdin is None:
            raise SandboxError("Process stdin is not available.")
        line = (message + "\n").encode()
        proc.stdin.write(line)
        await asyncio.wait_for(proc.stdin.drain(), timeout=timeout)

    @staticmethod
    async def _read_response(
        proc: asyncio.subprocess.Process,
        *,
        expected_id: int,
        timeout: float,
    ) -> dict[str, Any]:
        """Read stdout lines until a JSON-RPC response with *expected_id* arrives."""
        if proc.stdout is None:
            raise SandboxError("Process stdout is not available.")

        while True:
            raw_line = await asyncio.wait_for(
                proc.stdout.readline(),
                timeout=timeout,
            )

            if not raw_line:
                # EOF — process exited before sending a response
                stderr_output = ""
                if proc.stderr is not None:
                    stderr_bytes = await proc.stderr.read()
                    stderr_output = stderr_bytes.decode(errors="replace")
                raise SandboxError(
                    "Server process exited before responding."
                    + (f"\nStderr:\n{stderr_output}" if stderr_output else "")
                )

            line = raw_line.decode(errors="replace").strip()
            if not line:
                continue

            try:
                msg = json.loads(line)
            except json.JSONDecodeError:
                # Skip non-JSON output (e.g. debug logs, banner text)
                continue

            if not isinstance(msg, dict) or "id" not in msg:
                # Skip notifications
                continue

            if msg["id"] == expected_id:
                return msg

    # -- Internal helpers -----------------------------------------------------

    def _apply_firewall(self) -> None:
        """Execute the baked-in firewall script as root to set up iptables rules.

        Called after the container starts when domain filtering is configured.
        The container's default user is hdx (non-root), so we must exec as
        root to modify iptables.  After this, the rules are locked in — the
        non-root server process cannot alter them.
        """
        if self._container_id is None:
            return

        exit_code, output = self._runtime.exec_in_container(
            self._container_id,
            ["sh", "/etc/hdx/firewall.sh"],
            user="root",
        )
        if exit_code == 0:
            self._network_enforced = True
            logger.info("Firewall rules applied successfully")
        else:
            logger.warning(
                "Firewall setup failed (exit %d): %s", exit_code, output.strip()
            )

    def _get_stats_safe(self) -> ContainerStats | None:
        """Attempt to get container stats, returning None on failure."""
        if self._container_id is None:
            return None
        try:
            return self._runtime.get_stats(self._container_id)
        except SandboxError:
            return None

    def _register_signal_handlers(self) -> None:
        """Register signal handlers so Ctrl+C triggers cleanup.

        Original handlers are saved and restored during cleanup.
        """
        self._original_sigint: Any = None
        self._original_sigterm: Any = None

        def _handle_signal(signum: int, frame: Any) -> None:
            logger.info("Received signal %s, cleaning up sandbox...", signum)
            self.cleanup()
            self._restore_signal_handlers()

        try:
            self._original_sigint = signal.getsignal(signal.SIGINT)
            self._original_sigterm = signal.getsignal(signal.SIGTERM)
            signal.signal(signal.SIGINT, _handle_signal)
            signal.signal(signal.SIGTERM, _handle_signal)
        except (OSError, ValueError):
            # Can't set signal handlers outside the main thread — that's fine
            pass

    def _restore_signal_handlers(self) -> None:
        """Restore original signal handlers."""
        try:
            if self._original_sigint is not None:
                signal.signal(signal.SIGINT, self._original_sigint)
            if self._original_sigterm is not None:
                signal.signal(signal.SIGTERM, self._original_sigterm)
        except (OSError, ValueError):
            pass


async def _shutdown_process(proc: asyncio.subprocess.Process) -> None:
    """Gracefully shut down an exec subprocess.

    Closes stdin to signal the server, waits briefly, then terminates/kills
    if it doesn't exit on its own.
    """
    if proc.returncode is not None:
        return

    if proc.stdin is not None:
        proc.stdin.close()
        try:
            await proc.stdin.wait_closed()
        except (BrokenPipeError, ConnectionResetError):
            pass

    try:
        await asyncio.wait_for(proc.wait(), timeout=3.0)
    except asyncio.TimeoutError:
        proc.terminate()
        try:
            await asyncio.wait_for(proc.wait(), timeout=2.0)
        except asyncio.TimeoutError:
            proc.kill()
            await proc.wait()


def _parse_tool_response(output: str) -> tuple[list[dict[str, Any]], bool]:
    """Extract the tool call result from raw container output.

    Looks for JSON-RPC response lines and extracts the ``tools/call`` result.
    Returns ``(content, is_error)``.
    """
    content: list[dict[str, Any]] = []
    is_error = False

    for line in output.strip().splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            msg = json.loads(line)
        except json.JSONDecodeError:
            continue

        # Look for a response with id=2 (our tools/call request)
        if isinstance(msg, dict) and msg.get("id") == 2:
            if "error" in msg:
                is_error = True
                err = msg["error"]
                content = [{"type": "text", "text": err.get("message", str(err))}]
            elif "result" in msg:
                result = msg["result"]
                content = result.get("content", [])
                is_error = result.get("isError", False)
            break

    return content, is_error
