"""Sandbox runner API for executing tool calls inside isolated containers."""

from __future__ import annotations

import asyncio
import inspect
import io
import os
import secrets
import site
import subprocess
import sys
import threading
import time
import venv
from contextlib import redirect_stderr, redirect_stdout
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Annotated, Any

from fastapi import APIRouter, Depends, Header, HTTPException
from loguru import logger
from pydantic import BaseModel, Field, ValidationError

import mindroom.tool_system.sandbox_proxy as _sandbox_proxy
from mindroom.tool_system.metadata import ensure_tool_registry_loaded, get_tool_by_name
from mindroom.tool_system.sandbox_proxy import sandbox_proxy_token_matches, to_json_compatible
from mindroom.tool_system.worker_routing import (
    ToolExecutionIdentity,
    WorkerScope,
    tool_execution_identity,
    worker_dir_name,
)

if TYPE_CHECKING:
    from collections.abc import Callable

    from agno.tools.toolkit import Toolkit

    from mindroom.config.main import Config

_MAX_LEASE_TTL_SECONDS = 3600
_DEFAULT_LEASE_TTL_SECONDS = 60
_DEFAULT_SUBPROCESS_TIMEOUT_SECONDS = 120.0
_SUBPROCESS_WORKER_ARG = "--sandbox-subprocess-worker"
_RUNNER_EXECUTION_MODE_ENV = "MINDROOM_SANDBOX_RUNNER_EXECUTION_MODE"
_RUNNER_SUBPROCESS_TIMEOUT_ENV = "MINDROOM_SANDBOX_RUNNER_SUBPROCESS_TIMEOUT_SECONDS"
_RUNNER_WORKER_ROOT_ENV = "MINDROOM_SANDBOX_WORKER_ROOT"

# Sentinel written to stderr to delimit the JSON response from tool output.
_RESPONSE_MARKER = "__SANDBOX_RESPONSE__"


def _load_config_from_env() -> tuple[Config | None, Path | None]:
    """Read runner config path from environment variables."""
    from mindroom.config.main import Config as _Config  # noqa: PLC0415
    from mindroom.constants import find_config  # noqa: PLC0415

    config_path = find_config()
    if config_path.exists():
        return _Config.from_yaml(config_path), config_path
    return None, None


def ensure_registry_loaded_with_config() -> None:
    """Load config from env and ensure the tool registry is populated.

    Used by both the FastAPI startup and the subprocess worker so that
    plugin tools are registered even in fresh processes.
    """
    config, config_path = _load_config_from_env()
    ensure_tool_registry_loaded(config, config_path=config_path)


@dataclass
class _CredentialLease:
    """In-memory lease for short-lived credential overrides."""

    lease_id: str
    tool_name: str
    function_name: str
    credential_overrides: dict[str, Any]
    expires_at: float
    uses_remaining: int


@dataclass(frozen=True)
class _WorkerStatePaths:
    """Filesystem layout for one persistent worker scope."""

    root: Path
    workspace: Path
    venv_dir: Path
    cache_dir: Path
    storage_dir: Path


# NOTE: In-process dict — leases are not shared across multiple uvicorn workers.
# The sandbox runner must be deployed with a single worker for lease correctness.
_LEASES_BY_ID: dict[str, _CredentialLease] = {}
_LEASES_LOCK = threading.Lock()
_WORKER_INIT_LOCK = threading.Lock()


class SandboxRunnerExecuteRequest(BaseModel):
    """Tool call payload forwarded from a primary runtime to the sandbox runtime.

    Also used internally for in-process and subprocess execution when
    ``credential_overrides`` are resolved from a lease.
    """

    tool_name: str
    function_name: str
    args: list[Any] = Field(default_factory=list)
    kwargs: dict[str, Any] = Field(default_factory=dict)
    lease_id: str | None = None
    worker_key: str | None = None
    worker_scope: WorkerScope | None = None
    routing_agent_name: str | None = None
    execution_identity: dict[str, Any] = Field(default_factory=dict)
    credential_overrides: dict[str, Any] = Field(default_factory=dict)


class SandboxRunnerLeaseRequest(BaseModel):
    """Request for creating a short-lived credential lease."""

    tool_name: str
    function_name: str
    credential_overrides: dict[str, Any] = Field(default_factory=dict)
    ttl_seconds: int = _DEFAULT_LEASE_TTL_SECONDS
    max_uses: int = 1


class SandboxRunnerLeaseResponse(BaseModel):
    """Response describing a created credential lease."""

    lease_id: str
    expires_at: float
    max_uses: int


class SandboxRunnerExecuteResponse(BaseModel):
    """Sandbox tool execution response."""

    ok: bool
    result: Any | None = None
    error: str | None = None


async def _validate_runner_token(x_mindroom_sandbox_token: Annotated[str | None, Header()] = None) -> None:
    if _sandbox_proxy._PROXY_TOKEN is None:
        raise HTTPException(status_code=503, detail="Sandbox runner token is not configured.")
    if not sandbox_proxy_token_matches(x_mindroom_sandbox_token):
        raise HTTPException(status_code=401, detail="Unauthorized sandbox runner request")


router = APIRouter(
    prefix="/api/sandbox-runner",
    tags=["sandbox-runner"],
    dependencies=[Depends(_validate_runner_token)],
)


async def _maybe_await(value: object) -> object:
    if inspect.isawaitable(value):
        return await value
    return value


def _resolve_entrypoint(
    *,
    tool_name: str,
    function_name: str,
    credential_overrides: dict[str, object] | None = None,
    runtime_overrides: dict[str, object] | None = None,
    worker_scope: WorkerScope | None = None,
    routing_agent_name: str | None = None,
) -> tuple[Toolkit, Callable[..., object]]:
    ensure_registry_loaded_with_config()
    try:
        toolkit = get_tool_by_name(
            tool_name,
            disable_sandbox_proxy=True,
            credential_overrides=credential_overrides,
            runtime_overrides=runtime_overrides,
            worker_scope=worker_scope,
            routing_agent_name=routing_agent_name,
        )
    except ValueError as exc:
        raise HTTPException(status_code=404, detail=str(exc)) from exc
    function = toolkit.functions.get(function_name) or toolkit.async_functions.get(function_name)
    if function is None or function.entrypoint is None:
        raise HTTPException(status_code=404, detail=f"Tool '{tool_name}' does not expose '{function_name}'.")
    return toolkit, function.entrypoint


def _bounded_ttl_seconds(raw_ttl_seconds: int) -> int:
    return max(1, min(_MAX_LEASE_TTL_SECONDS, raw_ttl_seconds))


def _bounded_max_uses(raw_max_uses: int) -> int:
    return max(1, min(10, raw_max_uses))


def _cleanup_expired_leases(now: float) -> None:
    expired_ids = [lease_id for lease_id, lease in _LEASES_BY_ID.items() if lease.expires_at <= now]
    for lease_id in expired_ids:
        _LEASES_BY_ID.pop(lease_id, None)


def _create_credential_lease(request: SandboxRunnerLeaseRequest) -> _CredentialLease:
    ttl_seconds = _bounded_ttl_seconds(request.ttl_seconds)
    max_uses = _bounded_max_uses(request.max_uses)
    now = time.time()
    expires_at = now + ttl_seconds
    lease = _CredentialLease(
        lease_id=secrets.token_urlsafe(24),
        tool_name=request.tool_name,
        function_name=request.function_name,
        credential_overrides=dict(request.credential_overrides),
        expires_at=expires_at,
        uses_remaining=max_uses,
    )
    with _LEASES_LOCK:
        _cleanup_expired_leases(now)
        _LEASES_BY_ID[lease.lease_id] = lease
    return lease


def _consume_credential_lease(lease_id: str, *, tool_name: str, function_name: str) -> dict[str, object]:
    now = time.time()
    with _LEASES_LOCK:
        _cleanup_expired_leases(now)
        lease = _LEASES_BY_ID.get(lease_id)
        if lease is None:
            raise HTTPException(status_code=400, detail="Credential lease is invalid or expired.")
        if lease.tool_name != tool_name or lease.function_name != function_name:
            raise HTTPException(status_code=400, detail="Credential lease does not match tool/function.")

        lease.uses_remaining -= 1
        if lease.uses_remaining <= 0:
            _LEASES_BY_ID.pop(lease_id, None)

    return dict(lease.credential_overrides)


def _runner_execution_mode() -> str:
    return os.getenv(_RUNNER_EXECUTION_MODE_ENV, "inprocess").strip().lower()


def _runner_uses_subprocess() -> bool:
    return _runner_execution_mode() == "subprocess"


def _runner_subprocess_timeout_seconds() -> float:
    raw_timeout = os.getenv(_RUNNER_SUBPROCESS_TIMEOUT_ENV, str(_DEFAULT_SUBPROCESS_TIMEOUT_SECONDS))
    try:
        timeout = float(raw_timeout)
    except ValueError:
        timeout = _DEFAULT_SUBPROCESS_TIMEOUT_SECONDS
    return max(1.0, timeout)


def _runner_worker_root() -> Path:
    configured_root = os.getenv(_RUNNER_WORKER_ROOT_ENV)
    if configured_root:
        return Path(configured_root).expanduser().resolve()

    storage_path = os.getenv("MINDROOM_STORAGE_PATH", "/app/workspace/.mindroom")
    return Path(storage_path).expanduser().resolve() / "workers"


def _worker_state_paths(worker_key: str) -> _WorkerStatePaths:
    worker_root = _runner_worker_root() / worker_dir_name(worker_key)
    return _WorkerStatePaths(
        root=worker_root,
        workspace=worker_root / "workspace",
        venv_dir=worker_root / "venv",
        cache_dir=worker_root / "cache",
        storage_dir=worker_root,
    )


def _ensure_worker_state(paths: _WorkerStatePaths) -> None:
    with _WORKER_INIT_LOCK:
        paths.workspace.mkdir(parents=True, exist_ok=True)
        paths.cache_dir.mkdir(parents=True, exist_ok=True)
        if (paths.venv_dir / "bin" / "python").exists():
            return

        builder = venv.EnvBuilder(with_pip=True, system_site_packages=True)
        builder.create(paths.venv_dir)


def _project_src_path() -> Path:
    return Path(__file__).resolve().parents[2]


def _current_runtime_site_packages() -> list[str]:
    site_package_paths = list(site.getsitepackages())
    user_site = site.getusersitepackages()
    if isinstance(user_site, str):
        site_package_paths.append(user_site)

    discovered_paths: list[str] = []
    for path_text in site_package_paths:
        path = Path(path_text).expanduser()
        if path.is_dir():
            discovered_paths.append(str(path.resolve()))

    return list(dict.fromkeys(discovered_paths))


def _worker_subprocess_env(paths: _WorkerStatePaths) -> dict[str, str]:
    env = os.environ.copy()
    env["HOME"] = str(paths.root)
    env["MINDROOM_STORAGE_PATH"] = str(paths.storage_dir)
    env["XDG_CACHE_HOME"] = str(paths.cache_dir)
    env["PIP_CACHE_DIR"] = str(paths.cache_dir / "pip")
    env["UV_CACHE_DIR"] = str(paths.cache_dir / "uv")
    env["PYTHONPYCACHEPREFIX"] = str(paths.cache_dir / "pycache")
    env["VIRTUAL_ENV"] = str(paths.venv_dir)

    current_path = env.get("PATH", "")
    env["PATH"] = f"{paths.venv_dir / 'bin'}:{current_path}" if current_path else str(paths.venv_dir / "bin")

    project_src = str(_project_src_path())
    python_path_parts = [project_src, *_current_runtime_site_packages()]
    existing_python_path = env.get("PYTHONPATH", "")
    if existing_python_path:
        python_path_parts.append(existing_python_path)
    env["PYTHONPATH"] = ":".join(python_path_parts)
    return env


async def _execute_request_inprocess(request: SandboxRunnerExecuteRequest) -> SandboxRunnerExecuteResponse:
    runtime_overrides: dict[str, object] | None = None
    if request.worker_key is not None:
        paths = _worker_state_paths(request.worker_key)
        _ensure_worker_state(paths)
        runtime_overrides = {"base_dir": paths.workspace}

    execution_identity: ToolExecutionIdentity | None = None
    if request.execution_identity:
        execution_identity = ToolExecutionIdentity(**request.execution_identity)

    with tool_execution_identity(execution_identity):
        toolkit, entrypoint = _resolve_entrypoint(
            tool_name=request.tool_name,
            function_name=request.function_name,
            credential_overrides=request.credential_overrides or None,
            runtime_overrides=runtime_overrides,
            worker_scope=request.worker_scope,
            routing_agent_name=request.routing_agent_name,
        )

        try:
            if toolkit.requires_connect:
                await _maybe_await(toolkit.connect())
                try:
                    result = await _maybe_await(entrypoint(*request.args, **request.kwargs))
                finally:
                    await _maybe_await(toolkit.close())
            else:
                result = await _maybe_await(entrypoint(*request.args, **request.kwargs))
        except Exception as exc:
            logger.opt(exception=True).warning(
                f"Sandbox tool execution failed: {request.tool_name}.{request.function_name}",
            )
            return SandboxRunnerExecuteResponse(
                ok=False,
                error=f"Sandbox tool execution failed: {type(exc).__name__}: {exc}",
            )

    return SandboxRunnerExecuteResponse(ok=True, result=to_json_compatible(result))


def _subprocess_worker_command(python_executable: str | None = None) -> list[str]:
    return [python_executable or sys.executable, "-m", "mindroom.api.sandbox_runner", _SUBPROCESS_WORKER_ARG]


def _execute_request_subprocess_sync(request: SandboxRunnerExecuteRequest) -> SandboxRunnerExecuteResponse:
    subprocess_env: dict[str, str] | None = None
    python_executable: str | None = None
    cwd: str | None = None
    if request.worker_key is not None:
        paths = _worker_state_paths(request.worker_key)
        _ensure_worker_state(paths)
        python_executable = str(paths.venv_dir / "bin" / "python")
        subprocess_env = _worker_subprocess_env(paths)
        cwd = str(paths.workspace)

    try:
        completed = subprocess.run(
            _subprocess_worker_command(python_executable),
            input=request.model_dump_json(),
            capture_output=True,
            text=True,
            timeout=_runner_subprocess_timeout_seconds(),
            check=False,
            env=subprocess_env,
            cwd=cwd,
        )
    except subprocess.TimeoutExpired:
        return SandboxRunnerExecuteResponse(ok=False, error="Sandbox subprocess timed out.")
    except OSError as exc:
        return SandboxRunnerExecuteResponse(ok=False, error=f"Failed to start sandbox subprocess: {exc}")

    # The worker writes the JSON response to stderr after a marker line so that
    # tool stdout (e.g. print() inside python tools) does not corrupt the protocol.
    stderr = completed.stderr or ""
    marker_pos = stderr.rfind(_RESPONSE_MARKER)
    if marker_pos != -1:
        response_json = stderr[marker_pos + len(_RESPONSE_MARKER) :].strip()
        if response_json:
            try:
                return SandboxRunnerExecuteResponse.model_validate_json(response_json)
            except ValidationError:
                pass

    if completed.returncode != 0:
        error = (
            stderr.strip() or completed.stdout.strip() or f"Sandbox subprocess exited with code {completed.returncode}."
        )
        return SandboxRunnerExecuteResponse(ok=False, error=error)

    return SandboxRunnerExecuteResponse(ok=False, error="Sandbox subprocess returned an invalid response.")


async def _execute_request_subprocess(request: SandboxRunnerExecuteRequest) -> SandboxRunnerExecuteResponse:
    return await asyncio.to_thread(_execute_request_subprocess_sync, request)


def _run_subprocess_worker() -> int:
    payload = sys.stdin.read()
    if not payload.strip():
        print(
            _RESPONSE_MARKER
            + SandboxRunnerExecuteResponse(
                ok=False,
                error="Sandbox subprocess received empty payload.",
            ).model_dump_json(),
            file=sys.stderr,
        )
        return 1

    try:
        request = SandboxRunnerExecuteRequest.model_validate_json(payload)
    except ValidationError as exc:
        print(
            _RESPONSE_MARKER
            + SandboxRunnerExecuteResponse(
                ok=False,
                error=f"Sandbox subprocess payload validation failed: {exc}",
            ).model_dump_json(),
            file=sys.stderr,
        )
        return 1

    # Redirect stdout/stderr during tool execution so tool output doesn't
    # interfere with the protocol marker we write to stderr afterwards.
    captured_out = io.StringIO()
    captured_err = io.StringIO()
    with redirect_stdout(captured_out), redirect_stderr(captured_err):
        response = asyncio.run(_execute_request_inprocess(request))

    # Flush captured tool output to real stdout/stderr (informational only).
    tool_stdout = captured_out.getvalue()
    if tool_stdout:
        sys.stdout.write(tool_stdout)
    tool_stderr = captured_err.getvalue()
    if tool_stderr:
        sys.stdout.write(tool_stderr)

    # Write the response JSON to stderr after the marker.
    print(_RESPONSE_MARKER + response.model_dump_json(), file=sys.stderr)
    return 0


@router.post("/leases", response_model=SandboxRunnerLeaseResponse)
async def create_credential_lease(
    request: SandboxRunnerLeaseRequest,
) -> SandboxRunnerLeaseResponse:
    """Create a short-lived, one-or-few-use credential lease."""
    lease = _create_credential_lease(request)
    return SandboxRunnerLeaseResponse(
        lease_id=lease.lease_id,
        expires_at=lease.expires_at,
        max_uses=lease.uses_remaining,
    )


@router.post("/execute", response_model=SandboxRunnerExecuteResponse)
async def execute_tool_call(
    request: SandboxRunnerExecuteRequest,
) -> SandboxRunnerExecuteResponse:
    """Execute a tool function locally and return the serialized result."""
    credential_overrides: dict[str, object] = {}
    if request.lease_id is not None:
        credential_overrides = _consume_credential_lease(
            request.lease_id,
            tool_name=request.tool_name,
            function_name=request.function_name,
        )

    request.credential_overrides = credential_overrides
    if request.worker_key is not None or _runner_uses_subprocess():
        return await _execute_request_subprocess(request)
    return await _execute_request_inprocess(request)


if __name__ == "__main__":
    if _SUBPROCESS_WORKER_ARG in sys.argv:
        raise SystemExit(_run_subprocess_worker())
