from __future__ import annotations

import asyncio
import contextvars
import functools
import math
import os
import random
import shutil
import time
import uuid
from collections.abc import Awaitable, Callable, Mapping
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from email.utils import parsedate_to_datetime
from typing import Any, Protocol, TypeVar, override

import uvicorn
from fastapi import FastAPI, Request
from opentelemetry import context as otel_context
from opentelemetry import trace
from opentelemetry.propagate import extract
from starlette.types import ASGIApp, Receive, Scope, Send

from terminaluse.core.api_error import ApiError
from terminaluse.lib.adk._modules.filesystem import FilesystemModule
from terminaluse.lib.adk._modules.task import TaskModule
from terminaluse.lib.client import AsyncTerminalUse

# from terminaluse.lib.sdk.fastacp.types import BaseACPConfig
from terminaluse.lib.environment_variables import EnvironmentVariables, refreshed_environment_variables
from terminaluse.lib.sandbox.config import SandboxConfig, get_sandbox_config
from terminaluse.lib.sandbox.handler_ref import (
    HandlerRef,
    HandlerValidationError,
    validate_handler_for_sandbox,
)
from terminaluse.lib.sandbox.mounts import (
    Mount,
    assemble_sandbox_mounts,
    load_config,
    write_sandbox_mounts_file,
)
from terminaluse.lib.sandbox.runner import SANDBOX_SYSTEM_PATH, run_handler_sandboxed
from terminaluse.lib.sdk.fastacp.base.constants import (
    FASTACP_HEADER_SKIP_EXACT,
    FASTACP_HEADER_SKIP_PREFIXES,
)
from terminaluse.lib.types.acp import (
    PARAMS_MODEL_BY_METHOD,
    RPC_SYNC_METHODS,
    CancelTaskParams,
    CreateTaskParams,
    RPCMethod,
    SendEventParams,
)
from terminaluse.lib.types.json_rpc import JSONRPCError, JSONRPCRequest, JSONRPCResponse
from terminaluse.lib.types.task_context import TaskContext
from terminaluse.lib.utils.logging import ctx_var_request_id, make_logger
from terminaluse.lib.utils.model_utils import BaseModel
from terminaluse.lib.utils.registration import register_agent
from terminaluse.types.task import Task

logger = make_logger(__name__)
tracer = trace.get_tracer(__name__)
TRACE_CONTEXT_HEADERS = {"traceparent", "tracestate", "baggage"}
HANDLER_END_INITIAL_RETRY_DELAY_SECONDS = 1.0
HANDLER_END_MAX_RETRY_DELAY_SECONDS = 60.0
HANDLER_END_RETRY_JITTER_FACTOR = 0.2

# Contextvar for passing assembled mounts to sandbox runner (per-request, thread-safe)
_ctx_assembled_mounts: contextvars.ContextVar[list[Mount] | None] = contextvars.ContextVar(
    "assembled_mounts", default=None
)


class SandboxedHandlerError(RuntimeError):
    """Raised when a sandboxed handler fails to execute successfully."""

    def __init__(
        self,
        *,
        method: RPCMethod,
        handler_ref: HandlerRef,
        exit_code: int,
        timed_out: bool,
        status_json: dict[str, Any] | None,
    ) -> None:
        self.method = method
        self.handler_ref = handler_ref
        self.exit_code = exit_code
        self.timed_out = timed_out
        self.status_json = status_json

        detail = ""
        if status_json and status_json.get("error"):
            detail = f": {status_json['error']}"

        super().__init__(
            "Sandboxed handler failed for "
            f"{method.value} ({handler_ref.module}.{handler_ref.function}), "
            f"exit_code={exit_code}, timed_out={timed_out}{detail}"
        )


class UnsandboxableHandlerError(RuntimeError):
    """Raised when handler registration cannot guarantee sandbox execution."""


class ParamsWithTask(Protocol):
    """Protocol for RPC params that contain a task with filesystem_id."""

    task: Task


# TypeVar bound to ParamsWithTask for generic handler wrapping
P = TypeVar("P", bound=ParamsWithTask)


class RequestIDMiddleware:
    """Pure ASGI middleware to extract or generate request IDs.

    Uses pure ASGI (not BaseHTTPMiddleware) to preserve contextvars propagation,
    which is required for OpenTelemetry trace context to work correctly.

    See: https://www.starlette.io/middleware/#limitations
    """

    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        # Extract request ID from headers (ASGI headers are list of byte tuples)
        request_id = None
        for name, value in scope.get("headers", []):
            if name == b"x-request-id":
                request_id = value.decode("latin-1")
                break

        # Generate if not provided
        if not request_id:
            request_id = uuid.uuid4().hex

        # Set in contextvar with proper token handling
        token = ctx_var_request_id.set(request_id)
        try:
            await self.app(scope, receive, send)
        finally:
            ctx_var_request_id.reset(token)


class BaseACPServer(FastAPI):
    """
    AsyncAgentACP provides RPC-style hooks for agent events and commands asynchronously.
    All methods follow JSON-RPC 2.0 format.

    Available methods:
    - event/send → Send a message to a task
    - task/cancel → Cancel a task
    - task/approve → Approve a task
    """

    def __init__(self, sandbox_config: SandboxConfig | None = None):
        # Configure telemetry BEFORE creating FastAPI app
        # This ensures FastAPI instrumentation is set up before app initialization
        self._configure_telemetry()

        super().__init__(lifespan=self.get_lifespan_function())

        self.get("/healthz")(self._healthz)
        self.post("/api")(self._handle_jsonrpc)

        # Method handlers
        # this just adds a request ID to the request and response headers
        self.add_middleware(RequestIDMiddleware)

        self._handlers: dict[RPCMethod, Callable] = {}

        # Handler references for sandboxed execution
        self._handler_refs: dict[RPCMethod, HandlerRef] = {}

        # Sandbox configuration
        self._sandbox_config = sandbox_config or get_sandbox_config()

        # Branch ID to return in healthz (set during registration)
        self.branch_id: str | None = None

        # TerminalUse client for sync operations (created lazily)
        self._terminaluse_client: AsyncTerminalUse | None = None

        # Filesystem module for sync operations
        self._filesystem_module: FilesystemModule | None = None

        # Task module for system folder sync operations
        self._task_module: TaskModule | None = None

        # Background tasks set to prevent GC and track exceptions
        self._background_tasks: set[asyncio.Task] = set()

    def _configure_telemetry(self) -> None:
        """Configure OpenTelemetry for traces and metrics.

        Called automatically in __init__ before FastAPI initialization.
        Controlled via environment variables:
        - TERMINALUSE_TELEMETRY=false: Disable telemetry entirely
        - TERMINALUSE_AUTO_INSTRUMENT=false: Disable auto-instrumentation (manual spans only)

        Requires TERMINALUSE_BASE_URL and TERMINALUSE_AGENT_API_KEY to be set.
        """
        if os.environ.get("TERMINALUSE_TELEMETRY", "true").lower() == "false":
            logger.debug("Telemetry disabled via TERMINALUSE_TELEMETRY=false")
            return

        from terminaluse.lib.telemetry import configure_agent_telemetry

        auto_instrument = os.environ.get("TERMINALUSE_AUTO_INSTRUMENT", "true").lower() != "false"
        configure_agent_telemetry(auto_instrument=auto_instrument)

    def _bool_env(self, key: str, default: str = "false") -> bool:
        return os.getenv(key, default).strip().lower() in {"1", "true", "yes", "on"}

    def _log_runtime_binary_diagnostics(self) -> None:
        runtime_path = os.environ.get("PATH", "")
        runtime_node = shutil.which("node")
        runtime_npm = shutil.which("npm")
        sandbox_node = shutil.which("node", path=SANDBOX_SYSTEM_PATH)
        sandbox_npm = shutil.which("npm", path=SANDBOX_SYSTEM_PATH)

        logger.info(f"[runtime] PATH={runtime_path}")
        logger.info(f"[runtime] node={runtime_node or 'missing'}, npm={runtime_npm or 'missing'}")
        logger.info(
            f"[sandbox-path] PATH={SANDBOX_SYSTEM_PATH}, node={sandbox_node or 'missing'}, npm={sandbox_npm or 'missing'}"
        )

    def _enforce_required_runtime_binaries(self) -> None:
        # Opt-in guard for agents that depend on node during event/send.
        if not self._bool_env("TERMINALUSE_REQUIRE_NODE"):
            return

        sandbox_node = shutil.which("node", path=SANDBOX_SYSTEM_PATH)
        if sandbox_node:
            return

        raise RuntimeError(
            "TERMINALUSE_REQUIRE_NODE=true but node is not available in sandbox PATH "
            f"({SANDBOX_SYSTEM_PATH}). Install node in /usr/bin or /usr/local/bin."
        )

    @classmethod
    def create(cls):
        """Create and initialize BaseACPServer instance"""
        instance = cls()
        instance._setup_handlers()
        return instance

    def _setup_handlers(self):
        """Set up default handlers - override in subclasses"""
        # Base class has no default handlers
        pass

    def get_lifespan_function(self):
        @asynccontextmanager
        async def lifespan_context(app: FastAPI):  # noqa: ARG001
            self._log_runtime_binary_diagnostics()
            self._enforce_required_runtime_binaries()

            env_vars = EnvironmentVariables.refresh()
            if env_vars.TERMINALUSE_BASE_URL:
                await register_agent(env_vars)
                # Store branch_id for health check responses
                self.branch_id = env_vars.TERMINALUSE_BRANCH_ID
            else:
                logger.warning("TERMINALUSE_BASE_URL not set, skipping container registration")

            yield

        return lifespan_context

    async def _healthz(self):
        """Health check endpoint - returns branch_id for platform health monitoring"""
        result = {"status": "healthy"}
        if self.branch_id:
            result["branch_id"] = self.branch_id
        return result

    def _get_terminaluse_client(self) -> AsyncTerminalUse:
        """Get or create the TerminalUse client lazily"""
        if self._terminaluse_client is None:
            env_vars = EnvironmentVariables.refresh()
            self._terminaluse_client = AsyncTerminalUse(
                base_url=env_vars.TERMINALUSE_BASE_URL if env_vars else None,
                agent_api_key=env_vars.TERMINALUSE_AGENT_API_KEY if env_vars else None,
            )
        return self._terminaluse_client

    def _get_filesystem_module(self) -> FilesystemModule:
        """Get or create the FilesystemModule lazily"""
        if self._filesystem_module is None:
            self._filesystem_module = FilesystemModule(client=self._get_terminaluse_client())
        return self._filesystem_module

    def _get_task_module(self) -> TaskModule:
        """Get or create the TaskModule lazily"""
        if self._task_module is None:
            self._task_module = TaskModule(client=self._get_terminaluse_client())
        return self._task_module

    def _wrap_with_sync(
        self,
        fn: Callable[[P], Awaitable[Any]],
        handler_name: str,
        sync_up_after: bool = False,
    ) -> Callable[[P], Awaitable[Any]]:
        """
        Wrap a handler with filesystem and system folder sync.

        Each module manages its own cache internally.
        FastACP just coordinates the sync calls and assembles mounts.
        Mounts are passed directly to the runner in memory.

        Args:
            fn: The handler function to wrap (must accept params with a task)
            handler_name: Name for logging (e.g., "on_create")
            sync_up_after: If True, also sync_up after handler completes
        """

        async def handler_with_sync(params: P) -> Any:
            handler_start = time.monotonic()

            task = params.task
            filesystem_id = task.filesystem_id
            task_id = task.id

            with tracer.start_as_current_span(
                "fastacp.handler.with_sync",
                attributes={
                    "fastacp.handler.name": handler_name,
                    "fastacp.sync_up_after": sync_up_after,
                    "task.id": task_id or "",
                    "filesystem.id": filesystem_id or "",
                    "request.id": ctx_var_request_id.get(None) or "",
                },
            ) as handler_span:
                logger.info(f"[{handler_name}] Starting handler for task {task_id}, filesystem {filesystem_id}")

                if not task_id:
                    logger.warning(f"[{handler_name}] Task has no id, skipping sync")
                    handler_span.set_attribute("fastacp.sync.skipped_reason", "missing_task_id")
                    return await fn(params)

                if not filesystem_id:
                    logger.warning(
                        f"[{handler_name}] Task {task_id} has no filesystem_id, skipping sync - invoking handler directly"
                    )
                    handler_span.set_attribute("fastacp.sync.skipped_reason", "missing_filesystem_id")
                    return await fn(params)

                filesystem_module = self._get_filesystem_module()
                task_module = self._get_task_module()
                filesystem_local_path = filesystem_module.get_filesystem_path(filesystem_id)

                async def _run_with_span(
                    span_name: str,
                    coro: Awaitable[Any],
                    attributes: dict[str, Any] | None = None,
                ) -> Any:
                    with tracer.start_as_current_span(span_name, attributes=attributes):
                        return await coro

                # Pre-sync: Download filesystem and system folders in parallel
                sync_start = time.monotonic()
                logger.info(
                    f"[{handler_name}] Syncing down filesystem {filesystem_id} "
                    f"(local_path={filesystem_local_path}) and system folder for task {task_id}"
                )
                try:
                    with tracer.start_as_current_span(
                        "fastacp.handler.sync_down",
                        attributes={
                            "task.id": task_id,
                            "filesystem.id": filesystem_id,
                            "fastacp.handler.name": handler_name,
                        },
                    ) as sync_down_span:
                        with tracer.start_as_current_span(
                            "fastacp.handler.sync_down.parallel_wait",
                            attributes={
                                "task.id": task_id,
                                "filesystem.id": filesystem_id,
                            },
                        ):
                            await asyncio.gather(
                                _run_with_span(
                                    "fastacp.handler.sync_down.filesystem",
                                    filesystem_module.sync_down(filesystem_id, local_path=filesystem_local_path),
                                    {
                                        "task.id": task_id,
                                        "filesystem.id": filesystem_id,
                                    },
                                ),
                                _run_with_span(
                                    "fastacp.handler.sync_down.system_folder",
                                    task_module.sync_down_system_folder(task_id),
                                    {
                                        "task.id": task_id,
                                    },
                                ),
                            )
                        sync_down_span.set_attribute("sync.participants", 2)
                    sync_duration = time.monotonic() - sync_start
                    logger.info(f"[{handler_name}] Sync down completed in {sync_duration:.2f}s")
                    handler_span.set_attribute("fastacp.sync_down.duration_ms", int(sync_duration * 1000))
                except Exception as e:
                    sync_duration = time.monotonic() - sync_start
                    logger.error(f"[{handler_name}] Pre-sync failed after {sync_duration:.2f}s: {e}")
                    handler_span.set_attribute("fastacp.sync_down.failed", True)
                    # Continue with handler even if sync fails - may be new/empty

                # Load config and assemble mounts IN MEMORY
                # Use sandbox config's agent_dir if available, otherwise default to /app
                mount_start = time.monotonic()
                agent_workdir = str(self._sandbox_config.agent_dir) if self._sandbox_config else "/app"
                agent_config = load_config(agent_workdir)
                task_metadata = task.task_metadata or {}
                filesystem_access_mode = str(task_metadata.get("filesystem_access_mode", "rw")).lower()
                workspace_read_only = filesystem_access_mode == "ro"
                try:
                    with tracer.start_as_current_span(
                        "fastacp.handler.mount_assembly",
                        attributes={
                            "task.id": task_id,
                            "filesystem.id": filesystem_id,
                            "fastacp.handler.name": handler_name,
                            "workspace.read_only": workspace_read_only,
                        },
                    ):
                        mounts = assemble_sandbox_mounts(
                            task_id,
                            filesystem_id,
                            agent_config,
                            agent_workdir,
                            filesystem_host_path=str(filesystem_local_path),
                            workspace_read_only=workspace_read_only,
                        )
                        write_sandbox_mounts_file(task_id, mounts)
                        _ctx_assembled_mounts.set(mounts)
                    mount_duration = time.monotonic() - mount_start
                    logger.info(f"[{handler_name}] Mount assembly completed in {mount_duration:.2f}s")
                    handler_span.set_attribute("fastacp.mount_assembly.duration_ms", int(mount_duration * 1000))
                except Exception as e:
                    logger.error(f"[{handler_name}] Failed to assemble mounts: {e}")
                    handler_span.set_attribute("fastacp.mount_assembly.failed", True)
                    _ctx_assembled_mounts.set(None)

                pre_handler_duration = time.monotonic() - handler_start
                logger.info(
                    f"[{handler_name}] Pre-handler setup completed in {pre_handler_duration:.2f}s, invoking handler"
                )
                handler_span.set_attribute("fastacp.pre_handler.duration_ms", int(pre_handler_duration * 1000))

                if not sync_up_after:
                    try:
                        with tracer.start_as_current_span(
                            "fastacp.handler.execute",
                            attributes={
                                "task.id": task_id,
                                "fastacp.handler.name": handler_name,
                            },
                        ):
                            return await fn(params)
                    finally:
                        _ctx_assembled_mounts.set(None)

                # With sync_up_after: run handler then sync up
                try:
                    with tracer.start_as_current_span(
                        "fastacp.handler.execute",
                        attributes={
                            "task.id": task_id,
                            "fastacp.handler.name": handler_name,
                        },
                    ):
                        return await fn(params)
                finally:
                    logger.info(
                        f"[{handler_name}] Syncing up filesystem {filesystem_id} "
                        f"(local_path={filesystem_local_path}) and system folder for task {task_id}"
                    )
                    sync_up_start = time.monotonic()
                    with tracer.start_as_current_span(
                        "fastacp.handler.sync_up",
                        attributes={
                            "task.id": task_id,
                            "filesystem.id": filesystem_id,
                            "fastacp.handler.name": handler_name,
                        },
                    ) as sync_up_span:
                        with tracer.start_as_current_span(
                            "fastacp.handler.sync_up.parallel_wait",
                            attributes={
                                "task.id": task_id,
                                "filesystem.id": filesystem_id,
                            },
                        ):
                            sync_results = await asyncio.gather(
                                _run_with_span(
                                    "fastacp.handler.sync_up.filesystem",
                                    filesystem_module.sync_up(filesystem_id, local_path=filesystem_local_path),
                                    {
                                        "task.id": task_id,
                                        "filesystem.id": filesystem_id,
                                    },
                                ),
                                _run_with_span(
                                    "fastacp.handler.sync_up.system_folder",
                                    task_module.sync_up_system_folder(task_id),
                                    {
                                        "task.id": task_id,
                                    },
                                ),
                                return_exceptions=True,
                            )
                        failed_count = sum(1 for sync_result in sync_results if isinstance(sync_result, Exception))
                        sync_up_span.set_attribute("sync.participants", 2)
                        sync_up_span.set_attribute("sync.failed_count", failed_count)
                    for i, sync_result in enumerate(sync_results):
                        if isinstance(sync_result, Exception):
                            sync_name = "filesystem" if i == 0 else "system_folder"
                            logger.error(f"[{handler_name}] Post-sync failed for {sync_name}: {sync_result}")
                    sync_up_duration = time.monotonic() - sync_up_start
                    handler_span.set_attribute("fastacp.sync_up.duration_ms", int(sync_up_duration * 1000))
                    _ctx_assembled_mounts.set(None)

        return handler_with_sync

    def _should_sandbox(self, method: RPCMethod) -> bool:
        """Check if a method should be executed in a sandbox."""
        return method in self._handler_refs

    def _wrap_with_maybe_sandbox(
        self,
        fn: Callable[[P], Awaitable[Any]],
        method: RPCMethod,
    ) -> Callable[[P], Awaitable[Any]]:
        """Wrap a handler to run sandboxed or in-process based on config."""

        async def maybe_sandboxed(params: P) -> Any:
            if self._should_sandbox(method):
                return await self._run_sandboxed_handler(method, params)
            return await fn(params)

        return maybe_sandboxed

    async def _run_sandboxed_handler(
        self,
        method: RPCMethod,
        params: BaseModel,
    ) -> dict[str, Any] | None:
        """Run a handler in a sandbox."""
        handler_ref = self._handler_refs[method]
        task = getattr(params, "task", None)
        task_id = getattr(task, "id", None)
        event = getattr(params, "event", None)
        event_id = getattr(event, "id", None)
        agent = getattr(params, "agent", None)
        agent_id = getattr(agent, "id", None)
        mounts = _ctx_assembled_mounts.get()

        logger.info(f"Running handler in sandbox: {handler_ref.module}.{handler_ref.function}")
        if mounts:
            logger.info(f"Passing {len(mounts)} mounts to sandbox")

        span_attributes: dict[str, Any] = {
            "rpc.method": method.value,
            "sandbox.handler.module": handler_ref.module,
            "sandbox.handler.function": handler_ref.function,
            "sandbox.mounts_count": len(mounts) if mounts else 0,
            "request.id": ctx_var_request_id.get(None) or "",
            "request_id": ctx_var_request_id.get(None) or "",
        }
        if task_id:
            span_attributes["task.id"] = task_id
        if event_id:
            span_attributes["event.id"] = event_id
        if agent_id:
            span_attributes["agent.id"] = agent_id

        with tracer.start_as_current_span(
            "fastacp.handler.sandbox_execute",
            attributes=span_attributes,
        ) as span:
            with tracer.start_as_current_span(
                "fastacp.handler.sandbox.invoke_runner",
                attributes=span_attributes,
            ):
                result = await run_handler_sandboxed(
                    method=method.value,
                    handler_ref=handler_ref,
                    params=params,
                    mounts=mounts,
                    task_id=task_id,
                )
            span.set_attribute("sandbox.success", result.success)
            span.set_attribute("sandbox.exit_code", result.exit_code)
            span.set_attribute("sandbox.timed_out", result.timed_out)

        # Log result - no truncation, full output is valuable for debugging sandbox issues
        if result.success:
            logger.info(f"Sandboxed handler completed: {handler_ref.module}.{handler_ref.function}")
            if result.stderr:
                logger.warning(f"Sandbox stderr (handler succeeded): {result.stderr}")
            return result.status_json
        else:
            logger.error(
                f"Sandboxed handler failed: {handler_ref.module}.{handler_ref.function}, "
                f"exit_code={result.exit_code}, timed_out={result.timed_out}"
            )
            if result.stdout:
                logger.error(f"Sandbox stdout: {result.stdout}")
            if result.stderr:
                logger.error(f"Sandbox stderr: {result.stderr}")

            raise SandboxedHandlerError(
                method=method,
                handler_ref=handler_ref,
                exit_code=result.exit_code,
                timed_out=result.timed_out,
                status_json=result.status_json,
            )

    @staticmethod
    def _extract_trace_context_headers(headers: Mapping[str, str]) -> dict[str, str]:
        """Collect W3C trace context headers from incoming request headers."""
        trace_headers: dict[str, str] = {}
        for key, value in headers.items():
            key_lower = key.lower()
            if key_lower in TRACE_CONTEXT_HEADERS and value:
                trace_headers[key_lower] = value
        return trace_headers

    @staticmethod
    def _resolve_background_context(trace_headers: Mapping[str, str]) -> contextvars.Context:
        """Resolve context for background RPC processing.

        Prefer existing request context when a valid span already exists
        (auto-instrumentation enabled). Otherwise, extract remote context
        from request headers so background handler spans continue Nucleus
        traces across the RPC boundary.
        """
        if trace.get_current_span().get_span_context().is_valid:
            return contextvars.copy_context()

        if not trace_headers:
            return contextvars.copy_context()

        extracted_context = extract(dict(trace_headers))
        token = otel_context.attach(extracted_context)
        try:
            return contextvars.copy_context()
        finally:
            otel_context.detach(token)

    async def _handle_jsonrpc(self, request: Request):
        """Main JSON-RPC endpoint handler"""
        rpc_request = None
        logger.info(f"[base_acp_server] received request: {datetime.now()}")
        try:
            data = await request.json()
            rpc_request = JSONRPCRequest(**data)

            # Check if the request is authenticated
            if refreshed_environment_variables and getattr(
                refreshed_environment_variables, "TERMINALUSE_AGENT_API_KEY", None
            ):
                authorization_header = request.headers.get("x-agent-api-key")
                if authorization_header != refreshed_environment_variables.TERMINALUSE_AGENT_API_KEY:
                    return JSONRPCResponse(
                        id=rpc_request.id,
                        error=JSONRPCError(code=-32601, message="Unauthorized"),
                    )

            # Check if method is valid first
            try:
                method = RPCMethod(rpc_request.method)
            except ValueError:
                logger.error(f"Method {rpc_request.method} was invalid")
                return JSONRPCResponse(
                    id=rpc_request.id,
                    error=JSONRPCError(code=-32601, message=f"Method {rpc_request.method} not found"),
                )

            if method not in self._handlers or self._handlers[method] is None:
                logger.error(f"Method {method} not found on existing ACP server")
                return JSONRPCResponse(
                    id=rpc_request.id,
                    error=JSONRPCError(code=-32601, message=f"Method {method} not found"),
                )

            # Extract application headers using allowlist approach (only x-* headers)
            # Matches gateway's security filtering rules
            # Forward filtered headers via params.request.headers to agent handlers
            custom_headers = {
                key: value
                for key, value in request.headers.items()
                if (
                    (
                        key.lower().startswith("x-")
                        and key.lower() not in FASTACP_HEADER_SKIP_EXACT
                        and not any(key.lower().startswith(p) for p in FASTACP_HEADER_SKIP_PREFIXES)
                    )
                    or key.lower() in TRACE_CONTEXT_HEADERS
                )
            }
            trace_headers = self._extract_trace_context_headers(request.headers)

            # Parse params into appropriate model based on method and include headers
            params_model = PARAMS_MODEL_BY_METHOD[method]
            params_data = dict(rpc_request.params) if rpc_request.params else {}

            # Add custom headers to the request structure if any headers were provided
            # Gateway sends filtered headers via HTTP, SDK extracts and populates params.request
            if custom_headers:
                params_data["request"] = {"headers": custom_headers}
            params = params_model.model_validate(params_data)

            if method in RPC_SYNC_METHODS:
                handler = self._handlers[method]
                result = await handler(params)

                if rpc_request.id is None:
                    # Seems like you should return None for notifications
                    return None
                else:
                    if isinstance(result, BaseModel):
                        result = result.model_dump()
                    return JSONRPCResponse(id=rpc_request.id, result=result)
            else:
                # Preserve request context when instrumented; otherwise fall back to
                # explicit W3C context extraction from incoming headers.
                ctx = self._resolve_background_context(trace_headers)

                if rpc_request.id is None:
                    task = asyncio.create_task(self._process_notification(method, params), context=ctx)
                    self._track_background_task(task)
                    return JSONRPCResponse(id=None)

                task = asyncio.create_task(self._process_request(rpc_request.id, method, params), context=ctx)
                self._track_background_task(task)
                return JSONRPCResponse(id=rpc_request.id, result={"status": "processing"})

        except Exception as e:
            logger.error(f"Error handling JSON-RPC request: {e}", exc_info=True)
            request_id = None
            if rpc_request is not None:
                request_id = rpc_request.id
            return JSONRPCResponse(
                id=request_id,
                error=JSONRPCError(code=-32603, message=str(e)),
            )

    async def _signal_handler_end(
        self, method: RPCMethod, params: Any, success: bool, error: str | None, duration_ms: int
    ):
        """Signal handler-end to Nucleus so the task transitions RUNNING → IDLE.

        Only fires for EVENT_SEND (the method that triggers IDLE → RUNNING).
        Retries indefinitely with exponential backoff until Nucleus accepts.
        """
        if method != RPCMethod.EVENT_SEND:
            return

        task_id = getattr(getattr(params, "task", None), "id", None)
        if not task_id:
            return
        event_id = getattr(getattr(params, "event", None), "id", None)

        retries = 0
        while True:
            attempt = retries + 1
            try:
                attempt_log = logger.debug if retries == 0 else logger.info
                attempt_log(
                    "Calling handler-end for task %s event %s (attempt=%s retries=%s)",
                    task_id,
                    event_id or "",
                    attempt,
                    retries,
                )
                client = self._get_terminaluse_client()
                with tracer.start_as_current_span(
                    "fastacp.request.signal_handler_end",
                    attributes={
                        "task.id": task_id,
                        "event.id": event_id or "",
                        "rpc.method": method.value,
                        "handler.success": success,
                        "handler.duration_ms": duration_ms,
                        "request.id": ctx_var_request_id.get(None) or "",
                        "handler_end.retry_attempt": retries,
                    },
                ):
                    await client.tasks.handler_end(
                        task_id=task_id,
                        event_id=event_id,
                        success=success,
                        error=error,
                        duration_ms=duration_ms,
                    )

                success_log = logger.debug if retries == 0 else logger.info
                success_log(
                    "Handler-end accepted for task %s event %s (attempt=%s retries=%s)",
                    task_id,
                    event_id or "",
                    attempt,
                    retries,
                )
                return
            except asyncio.CancelledError:
                raise
            except Exception as e:
                if not self._should_retry_handler_end_error(e):
                    status_code = e.status_code if isinstance(e, ApiError) else None
                    logger.error(
                        "Failed to signal handler-end for task %s event %s "
                        "(attempt=%s retries=%s): non-retryable error (status=%s), not retrying: %s",
                        task_id,
                        event_id or "",
                        attempt,
                        retries,
                        status_code,
                        e,
                        exc_info=True,
                    )
                    return
                delay_s = self._get_handler_end_retry_delay_seconds(error=e, retry_attempt=retries)
                retries += 1
                logger.error(
                    "Failed to signal handler-end for task %s event %s (attempt=%s retries=%s): %s. Retrying in %.2fs",
                    task_id,
                    event_id or "",
                    attempt,
                    retries,
                    e,
                    delay_s,
                    exc_info=True,
                )
                await asyncio.sleep(delay_s)

    @staticmethod
    def _is_retryable_handler_end_status_code(status_code: int | None) -> bool:
        """Return whether a handler-end HTTP status should be retried."""
        if status_code is None:
            return True
        if status_code >= 500:
            return True
        if status_code in {408, 429}:
            return True
        if 400 <= status_code < 500:
            return False
        return True

    @classmethod
    def _should_retry_handler_end_error(cls, error: Exception) -> bool:
        """Return whether a handler-end failure is retryable."""
        if isinstance(error, ApiError):
            return cls._is_retryable_handler_end_status_code(error.status_code)
        return True

    @staticmethod
    def _parse_retry_after_seconds(headers: Mapping[str, str] | None) -> float | None:
        """Parse Retry-After header as delta-seconds or HTTP-date."""
        if not headers:
            return None

        retry_after_value = next(
            (value for key, value in headers.items() if key.lower() == "retry-after"),
            None,
        )
        if not retry_after_value:
            return None

        value = retry_after_value.strip()
        if not value:
            return None

        try:
            seconds = float(value)
            if seconds >= 0:
                return seconds
            return None
        except ValueError:
            pass

        try:
            retry_after_datetime = parsedate_to_datetime(value)
        except (TypeError, ValueError, IndexError):
            return None

        if retry_after_datetime is None:
            return None

        if retry_after_datetime.tzinfo is None:
            retry_after_datetime = retry_after_datetime.replace(tzinfo=timezone.utc)

        seconds_until_retry = (retry_after_datetime - datetime.now(timezone.utc)).total_seconds()
        return max(0.0, seconds_until_retry)

    def _get_handler_end_retry_delay_seconds(self, error: Exception, retry_attempt: int) -> float:
        """Compute retry delay, honoring Retry-After on 429 responses when present."""
        computed_delay = self._compute_handler_end_retry_delay_seconds(retry_attempt)
        if not isinstance(error, ApiError) or error.status_code != 429:
            return computed_delay

        retry_after_delay = self._parse_retry_after_seconds(error.headers)
        if retry_after_delay is None:
            return computed_delay

        return retry_after_delay

    @staticmethod
    def _compute_handler_end_retry_delay_seconds(retry_attempt: int) -> float:
        """Compute exponential backoff delay with symmetric jitter for handler-end retries."""
        bounded_retry_attempt = max(0, retry_attempt)

        if HANDLER_END_INITIAL_RETRY_DELAY_SECONDS <= 0:
            base_delay = 0.0
        elif HANDLER_END_INITIAL_RETRY_DELAY_SECONDS >= HANDLER_END_MAX_RETRY_DELAY_SECONDS:
            base_delay = HANDLER_END_MAX_RETRY_DELAY_SECONDS
        else:
            # Saturate without evaluating huge exponents that can overflow.
            saturation_attempt = math.ceil(
                math.log2(HANDLER_END_MAX_RETRY_DELAY_SECONDS / HANDLER_END_INITIAL_RETRY_DELAY_SECONDS)
            )
            if bounded_retry_attempt >= saturation_attempt:
                base_delay = HANDLER_END_MAX_RETRY_DELAY_SECONDS
            else:
                base_delay = HANDLER_END_INITIAL_RETRY_DELAY_SECONDS * pow(2.0, bounded_retry_attempt)

        jitter_multiplier = 1 + (random.random() - 0.5) * HANDLER_END_RETRY_JITTER_FACTOR
        return max(0.0, base_delay * jitter_multiplier)

    async def _process_notification(self, method: RPCMethod, params: Any):
        """Process a notification (request with no ID) in the background"""
        start = time.monotonic()
        logger.info(f"[base_acp_server] Background task started for notification {method}")
        success = True
        error_msg: str | None = None
        task_id = getattr(getattr(params, "task", None), "id", "")
        with tracer.start_as_current_span(
            "fastacp.request.process_notification",
            attributes={
                "rpc.method": method.value,
                "task.id": task_id,
                "request.id": ctx_var_request_id.get(None) or "",
            },
        ) as span:
            try:
                await self._handlers[method](params)
                duration = time.monotonic() - start
                logger.info(f"[base_acp_server] Notification {method} completed in {duration:.2f}s")
            except Exception as e:
                success = False
                error_msg = str(e)
                logger.error(f"Error processing notification {method}: {e}", exc_info=True)
            finally:
                duration_ms = int((time.monotonic() - start) * 1000)
                span.set_attribute("handler.duration_ms", duration_ms)
                span.set_attribute("handler.success", success)
                await self._signal_handler_end(method, params, success, error_msg, duration_ms)

    async def _process_request(self, request_id: int | str, method: RPCMethod, params: Any):
        """Process a request in the background"""
        start = time.monotonic()
        logger.info(f"[base_acp_server] Background task started for request {request_id} method {method}")
        success = True
        error_msg: str | None = None
        task_id = getattr(getattr(params, "task", None), "id", "")
        with tracer.start_as_current_span(
            "fastacp.request.process",
            attributes={
                "rpc.method": method.value,
                "rpc.request_id": str(request_id),
                "task.id": task_id,
                "request.id": ctx_var_request_id.get(None) or "",
            },
        ) as span:
            try:
                await self._handlers[method](params)
                duration = time.monotonic() - start
                logger.info(f"Successfully processed request {request_id} for method {method} in {duration:.2f}s")
            except Exception as e:
                success = False
                error_msg = str(e)
                duration = time.monotonic() - start
                logger.error(
                    f"Error processing request {request_id} for method {method} after {duration:.2f}s: {e}",
                    exc_info=True,
                )
            finally:
                duration_ms = int((time.monotonic() - start) * 1000)
                span.set_attribute("handler.duration_ms", duration_ms)
                span.set_attribute("handler.success", success)
                await self._signal_handler_end(method, params, success, error_msg, duration_ms)

    def _track_background_task(self, task: asyncio.Task) -> None:
        """Track a background task to prevent GC and log uncaught exceptions.

        Adds the task to the internal set to keep a strong reference (preventing
        garbage collection), and registers a done callback to:
        1. Remove the task from the set when it completes
        2. Log any uncaught exceptions that weren't handled in the task itself

        This ensures fire-and-forget tasks don't silently swallow exceptions.
        """
        self._background_tasks.add(task)

        def _on_task_done(t: asyncio.Task) -> None:
            self._background_tasks.discard(t)
            # Check for exceptions that weren't raised/handled within the task
            if not t.cancelled():
                exc = t.exception()
                if exc is not None:
                    logger.error(
                        f"Uncaught exception in background task {t.get_name()}: {exc}",
                        exc_info=(type(exc), exc, exc.__traceback__),
                    )

        task.add_done_callback(_on_task_done)

    """
    Define all possible decorators to be overriden and implemented by each ACP implementation
    Then the users can override the default handlers by implementing their own handlers

    ACP Type: Async
    Decorators:
    - on_create
    - on_event
    - on_cancel
    """

    def _register_handler(
        self,
        fn: Callable[..., Awaitable[Any]],
        method: RPCMethod,
        handler_name: str,
        sync_up_after: bool = False,
        skip_sync: bool = False,
        sandbox_source_fn: Callable[..., Awaitable[Any]] | None = None,
        allow_in_process: bool = False,
    ) -> None:
        """Internal handler registration with sandboxing and sync wrapping.

        Args:
            fn: The handler function to register
            method: The RPC method to register the handler for
            handler_name: Name for logging (e.g., "on_create")
            sync_up_after: If True, sync filesystem up after handler completes
            skip_sync: If True, skip filesystem sync entirely (e.g., for cancel)
            sandbox_source_fn: Original handler for sandbox import path validation.
                Use this when fn is a framework wrapper around a user handler.
            allow_in_process: Allow explicit in-process execution when handler
                cannot be sandboxed.
        """
        validation_target = sandbox_source_fn or fn

        # Register for sandboxing
        try:
            handler_ref = validate_handler_for_sandbox(validation_target)
            self._handler_refs[method] = handler_ref
            logger.info(f"Handler registered for sandboxing: {handler_ref.module}.{handler_ref.function}")
        except HandlerValidationError as e:
            # Clear any prior sandbox ref for this method when registering
            # an unsandboxable replacement.
            self._handler_refs.pop(method, None)

            if allow_in_process:
                logger.info(f"Handler configured to run in-process (sandbox skipped): {e}")
            elif self._bool_env("TERMINALUSE_SANDBOX_FAIL_OPEN"):
                logger.warning(f"Handler not suitable for sandboxing, will run in-process: {e}")
            else:
                raise UnsandboxableHandlerError(
                    "Handler must be sandboxable (module-level function). "
                    "Set TERMINALUSE_SANDBOX_FAIL_OPEN=true to preserve legacy in-process fallback."
                ) from e

        maybe_sandboxed = self._wrap_with_maybe_sandbox(fn, method)

        if skip_sync:
            self._handlers[method] = maybe_sandboxed
        else:
            self._handlers[method] = self._wrap_with_sync(
                maybe_sandboxed, handler_name=handler_name, sync_up_after=sync_up_after
            )

    # ==========================================================================
    # TaskContext-based decorators (public API)
    # ==========================================================================

    def on_create(
        self,
        fn: Callable[..., Awaitable[Any]] | None = None,
        *,
        allow_in_process: bool = False,
    ):
        """Handle task/create with TaskContext injection.

        The decorated function receives:
        - ctx: TaskContext with pre-bound task/agent IDs and module wrappers
        - params: dict[str, Any] containing the user-provided parameters

        Example:
            @server.on_create
            async def handle_create(ctx: TaskContext, params: dict[str, Any]):
                agent_type = params.get("agent_type")
                await ctx.state.create(state={"agent_type": agent_type})
                await ctx.reply("Task created!")
        """

        def decorator(handler_fn: Callable[..., Awaitable[Any]]):
            @functools.wraps(handler_fn)
            async def with_context(params: CreateTaskParams):
                ctx = TaskContext(task=params.task, agent=params.agent, request=params.request)
                return await handler_fn(ctx, params.params or {})

            self._register_handler(
                with_context,
                RPCMethod.TASK_CREATE,
                handler_name="on_create",
                sync_up_after=True,
                sandbox_source_fn=handler_fn,
                allow_in_process=allow_in_process,
            )
            return handler_fn

        if fn is None:
            return decorator
        return decorator(fn)

    @override  # Intentionally shadows FastAPI's deprecated on_event lifecycle decorator
    def on_event(
        self,
        fn: Callable[..., Awaitable[Any]] | None = None,
        *,
        allow_in_process: bool = False,
    ):
        """Handle event/send with TaskContext injection.

        The decorated function receives:
        - ctx: TaskContext with pre-bound task/agent IDs and module wrappers
        - event: Event object containing the incoming event data

        Example:
            @server.on_event
            async def handle_event(ctx: TaskContext, event: Event):
                await ctx.reply(f"Received: {event.content}")
        """

        def decorator(handler_fn: Callable[..., Awaitable[Any]]):
            @functools.wraps(handler_fn)
            async def with_context(params: SendEventParams):
                ctx = TaskContext(task=params.task, agent=params.agent, request=params.request)
                return await handler_fn(ctx, params.event)

            self._register_handler(
                with_context,
                RPCMethod.EVENT_SEND,
                handler_name="on_event",
                sync_up_after=True,
                sandbox_source_fn=handler_fn,
                allow_in_process=allow_in_process,
            )
            return handler_fn

        if fn is None:
            return decorator
        return decorator(fn)

    def on_cancel(
        self,
        fn: Callable[..., Awaitable[Any]] | None = None,
        *,
        allow_in_process: bool = False,
    ):
        """Handle task/cancel with TaskContext injection.

        The decorated function receives:
        - ctx: TaskContext with pre-bound task/agent IDs and module wrappers

        Example:
            @server.on_cancel
            async def handle_cancel(ctx: TaskContext):
                await ctx.reply("Task cancelled")
        """

        def decorator(handler_fn: Callable[..., Awaitable[Any]]):
            @functools.wraps(handler_fn)
            async def with_context(params: CancelTaskParams):
                ctx = TaskContext(task=params.task, agent=params.agent, request=params.request)
                return await handler_fn(ctx)

            self._register_handler(
                with_context,
                RPCMethod.TASK_CANCEL,
                handler_name="on_cancel",
                skip_sync=True,
                sandbox_source_fn=handler_fn,
                allow_in_process=allow_in_process,
            )
            return handler_fn

        if fn is None:
            return decorator
        return decorator(fn)

    """
    End of Decorators
    """

    """
    ACP Server Lifecycle Methods
    """

    def run(self, host: str = "0.0.0.0", port: int = 8000, **kwargs):
        """Start the Uvicorn server for async handlers."""
        uvicorn.run(self, host=host, port=port, **kwargs)
