"""
Sandboxed handler entrypoint - runs inside bubblewrap sandbox.

Protocol:
  stdin:  JSON object with keys: method, handler, params
  stdout: Small status object as JSON (optional)
  stderr: All logging output (captured by runner)
  exit 0: success
  exit 1: failure
"""

from __future__ import annotations

import asyncio
import importlib
import json
import logging
import os
import sys
from typing import Any, Callable

from opentelemetry import context as otel_context
from opentelemetry import trace
from opentelemetry.propagate import extract

TRACE_CONTEXT_HEADERS = {"traceparent", "tracestate", "baggage"}
SANDBOX_SPAWN_TIME_KEY = "sandbox_spawn_time_ns"
tracer = trace.get_tracer(__name__)


class _StructuredLogFormatter(logging.Formatter):
    """JSON log formatter for structured log output.

    Outputs logs as JSON with separate fields for timestamp, level, logger, and message.
    This allows the LogSender to extract metadata without parsing text prefixes.
    """

    def format(self, record: logging.LogRecord) -> str:
        from datetime import datetime, timezone

        return json.dumps(
            {
                "ts": datetime.fromtimestamp(record.created, tz=timezone.utc).isoformat(),
                "level": record.levelname,
                "logger": record.name,
                "msg": record.getMessage(),
            },
            separators=(",", ":"),
        )


def _configure_logging() -> None:
    """Configure logging to stderr for sandbox environment.

    All logs go to stderr so they can be captured by the sandbox runner.
    Uses JSON format for structured log parsing by LogSender.
    The log level is controlled by TERMINALUSE_LOG_LEVEL (default: INFO).
    """
    log_level = os.getenv("TERMINALUSE_LOG_LEVEL", "INFO").upper()

    # Configure root logger with JSON formatter
    handler = logging.StreamHandler(sys.stderr)
    handler.setFormatter(_StructuredLogFormatter())

    # Clear any existing handlers and set up fresh
    logging.root.handlers = []
    logging.root.addHandler(handler)
    logging.root.setLevel(getattr(logging, log_level, logging.INFO))


def _timing_logs_enabled() -> bool:
    return os.getenv("TERMINALUSE_SANDBOX_TIMING_LOGS", "false").strip().lower() in {
        "1",
        "true",
        "yes",
        "on",
    }


def _log_timing_or_invocation(logger: logging.Logger, message: str) -> None:
    if _timing_logs_enabled():
        logger.info(message)
    else:
        logger.debug(message)


def _get_params_class(method: str):
    """Get the appropriate params class for a method."""
    # Import here to ensure it works in the jail
    from terminaluse.lib.types.acp import (
        CancelTaskParams,
        CreateTaskParams,
        SendEventParams,
    )

    method_to_class = {
        "task/create": CreateTaskParams,
        "event/send": SendEventParams,
        "task/cancel": CancelTaskParams,
    }

    cls = method_to_class.get(method)
    if cls is None:
        raise ValueError(f"Unsupported method: {method}")
    return cls


def _load_params(method: str, data: dict[str, Any]):
    """Load and validate params for the given method."""
    params_data = data.get("params", {})
    params_class = _get_params_class(method)
    return params_class.model_validate(params_data)


def _invoke_handler(method: str, params: Any, fn: Callable[..., Any]) -> None:
    """
    Invoke handler with TaskContext injection based on method type.

    The handler signatures are:
    - task/create: fn(ctx: TaskContext, params: dict[str, Any])
    - event/send: fn(ctx: TaskContext, event: Event)
    - task/cancel: fn(ctx: TaskContext)
    """
    from terminaluse.lib.types.task_context import TaskContext

    # Create TaskContext from params (all methods have task, agent, request)
    delivery_event_id = None
    if method == "event/send":
        delivery_event_id = getattr(getattr(params, "event", None), "id", None)
    ctx = TaskContext(
        task=params.task,
        agent=params.agent,
        request=params.request,
        delivery_event_id=delivery_event_id,
    )

    if method == "task/create":
        args = (ctx, params.params or {})
    elif method == "event/send":
        args = (ctx, params.event)
    elif method == "task/cancel":
        args = (ctx,)
    else:
        raise ValueError(f"Unsupported method: {method}")

    if asyncio.iscoroutinefunction(fn):
        asyncio.run(fn(*args))
    else:
        fn(*args)


def _collect_trace_context_carrier(payload: dict[str, Any]) -> dict[str, str]:
    """Collect W3C trace context from payload.

    Priority:
    1) Headers explicitly forwarded in payload.params.request.headers
    2) Top-level payload.trace_context
    """
    trace_headers: dict[str, str] = {}

    payload_trace_context = payload.get("trace_context")
    if isinstance(payload_trace_context, dict):
        for key, value in payload_trace_context.items():
            if not isinstance(key, str):
                continue
            key_lower = key.lower()
            if key_lower in TRACE_CONTEXT_HEADERS:
                trace_headers[key_lower] = str(value)

    payload_params = payload.get("params")
    if isinstance(payload_params, dict):
        request_payload = payload_params.get("request")
        if isinstance(request_payload, dict):
            request_headers = request_payload.get("headers")
            if isinstance(request_headers, dict):
                for key, value in request_headers.items():
                    if not isinstance(key, str):
                        continue
                    key_lower = key.lower()
                    if key_lower in TRACE_CONTEXT_HEADERS:
                        trace_headers[key_lower] = str(value)

    return trace_headers


def _load_handler(handler_ref: dict[str, str]) -> Callable[..., Any]:
    """
    Load a user handler by import path.

    The handler must be importable inside the jail and use the TaskContext-based signature:
    - on_create: fn(ctx: TaskContext, params: dict[str, Any])
    - on_event: fn(ctx: TaskContext, event: Event)
    - on_cancel: fn(ctx: TaskContext)
    """
    module_name = handler_ref.get("module")
    function_name = handler_ref.get("function")

    if not module_name or not function_name:
        raise ValueError("handler.module and handler.function are required")

    try:
        with tracer.start_as_current_span(
            "sandbox.entrypoint.import_handler_module",
            attributes={
                "sandbox.handler.module": module_name,
                "sandbox.handler.function": function_name,
            },
        ):
            module = importlib.import_module(module_name)
    except ImportError as e:
        raise ValueError(f"Could not import handler module: {module_name}") from e

    fn = getattr(module, function_name, None)
    if fn is None:
        raise ValueError(f"Handler not found: {module_name}.{function_name}")

    if not callable(fn):
        raise ValueError(f"Handler is not callable: {module_name}.{function_name}")

    return fn


def main() -> int:
    """Main entrypoint for jailed handler execution."""
    import time

    t_start = time.perf_counter()

    # Configure logging first so all handler logs go to stderr
    _configure_logging()
    logger = logging.getLogger(__name__)

    t_logging = time.perf_counter()
    context_token = None

    try:
        _log_timing_or_invocation(
            logger,
            f"[TIMING] sandbox_entrypoint_start: {(t_logging - t_start) * 1000:.0f}ms (logging setup)",
        )

        # Read payload from stdin
        with tracer.start_as_current_span("sandbox.entrypoint.read_payload"):
            raw = sys.stdin.read()
            if not raw:
                raise ValueError("No input received on stdin")
            payload = json.loads(raw)
        t_payload = time.perf_counter()
        _log_timing_or_invocation(
            logger,
            f"[TIMING] payload_read: {(t_payload - t_logging) * 1000:.0f}ms",
        )

        # Continue upstream trace context across sandbox boundary before expensive work.
        trace_context_carrier = _collect_trace_context_carrier(payload)
        if trace_context_carrier:
            extracted_context = extract(trace_context_carrier)
            context_token = otel_context.attach(extracted_context)

        with tracer.start_as_current_span(
            "sandbox.entrypoint.extract_trace_context",
            attributes={"sandbox.entrypoint.trace_context.headers": len(trace_context_carrier)},
        ) as extract_context_span:
            extract_context_span.set_attribute(
                "sandbox.entrypoint.trace_context.applied",
                bool(trace_context_carrier),
            )

        # Extract and validate
        with tracer.start_as_current_span("sandbox.entrypoint.validate_payload"):
            method = payload.get("method")
            if not method:
                raise ValueError("Missing 'method' in payload")

            handler_ref = payload.get("handler", {})
            if not handler_ref:
                raise ValueError("Missing 'handler' in payload")

            payload_params = payload.get("params")
            request_id = ""
            task_id = ""
            event_id = ""
            spawn_to_main_ms: int | None = None
            if isinstance(payload_params, dict):
                request_payload = payload_params.get("request")
                if isinstance(request_payload, dict):
                    request_id_value = request_payload.get("id")
                    if isinstance(request_id_value, str):
                        request_id = request_id_value
                task_payload = payload_params.get("task")
                if isinstance(task_payload, dict):
                    task_id_value = task_payload.get("id")
                    if isinstance(task_id_value, str):
                        task_id = task_id_value
                event_payload = payload_params.get("event")
                if isinstance(event_payload, dict):
                    event_id_value = event_payload.get("id")
                    if isinstance(event_id_value, str):
                        event_id = event_id_value

            spawn_time_value = payload.get(SANDBOX_SPAWN_TIME_KEY)
            if isinstance(spawn_time_value, int) and spawn_time_value > 0:
                delta_ns = time.time_ns() - spawn_time_value
                if delta_ns >= 0:
                    spawn_to_main_ms = int(delta_ns / 1_000_000)

        main_span_attributes: dict[str, Any] = {
            "sandbox.method": method,
            "sandbox.handler.module": str(handler_ref.get("module") or ""),
            "sandbox.handler.function": str(handler_ref.get("function") or ""),
            "sandbox.entrypoint.payload.bytes": len(raw),
            "sandbox.entrypoint.startup.logging_setup_ms": int((t_logging - t_start) * 1000),
            "sandbox.entrypoint.startup.payload_read_ms": int((t_payload - t_logging) * 1000),
            "sandbox.entrypoint.trace_context.headers": len(trace_context_carrier),
        }
        if task_id:
            main_span_attributes["task.id"] = task_id
        if event_id:
            main_span_attributes["event.id"] = event_id
        if request_id:
            main_span_attributes["request.id"] = request_id
            main_span_attributes["request_id"] = request_id
        if spawn_to_main_ms is not None:
            main_span_attributes["sandbox.entrypoint.spawn_to_main_ms"] = spawn_to_main_ms
            _log_timing_or_invocation(
                logger,
                f"[TIMING] spawn_to_main: {spawn_to_main_ms}ms",
            )

        with tracer.start_as_current_span(
            "sandbox.entrypoint.main",
            attributes=main_span_attributes,
        ) as main_span:
            _log_timing_or_invocation(
                logger,
                f"Invoking handler: {handler_ref.get('module')}.{handler_ref.get('function')} for method={method}",
            )

            # Load params
            t_before_params = time.perf_counter()
            with tracer.start_as_current_span(
                "sandbox.entrypoint.load_params",
                attributes={"sandbox.method": method},
            ):
                params = _load_params(method, payload)
            t_after_params = time.perf_counter()
            _log_timing_or_invocation(
                logger,
                f"[TIMING] load_params: {(t_after_params - t_before_params) * 1000:.0f}ms",
            )

            # Load handler (includes module import span inside _load_handler)
            with tracer.start_as_current_span(
                "sandbox.entrypoint.load_handler",
                attributes={
                    "sandbox.handler.module": str(handler_ref.get("module") or ""),
                    "sandbox.handler.function": str(handler_ref.get("function") or ""),
                },
            ):
                fn = _load_handler(handler_ref)
            t_after_handler_load = time.perf_counter()
            _log_timing_or_invocation(
                logger,
                f"[TIMING] load_handler (imports): {(t_after_handler_load - t_after_params) * 1000:.0f}ms",
            )

            # Execute handler with TaskContext injection
            with tracer.start_as_current_span(
                "sandbox.entrypoint.invoke_handler",
                attributes={
                    "sandbox.method": method,
                    "sandbox.handler.module": str(handler_ref.get("module") or ""),
                    "sandbox.handler.function": str(handler_ref.get("function") or ""),
                },
            ):
                _invoke_handler(method, params, fn)

            t_after_invoke = time.perf_counter()
            _log_timing_or_invocation(
                logger,
                f"[TIMING] handler_execution: {(t_after_invoke - t_after_handler_load) * 1000:.0f}ms",
            )
            _log_timing_or_invocation(
                logger,
                f"[TIMING] total_sandbox_time: {(t_after_invoke - t_start) * 1000:.0f}ms",
            )
            _log_timing_or_invocation(logger, "Handler completed successfully")

            # Success - write status to stdout
            with tracer.start_as_current_span("sandbox.entrypoint.serialize_response"):
                out = {"status": "ok"}
                response_body = json.dumps(out)
                sys.stdout.write(response_body)

            main_span.set_attribute("sandbox.success", True)
            main_span.set_attribute("sandbox.entrypoint.total_ms", int((t_after_invoke - t_start) * 1000))
        return 0

    except Exception as exc:
        # Log full exception with traceback
        logger.exception(f"Sandboxed handler failed: {exc}")

        # Try to write error status to stdout (ACP wrapper may use this)
        try:
            with tracer.start_as_current_span("sandbox.entrypoint.serialize_error_response"):
                error_out = {"status": "error", "error": str(exc)}
                sys.stdout.write(json.dumps(error_out))
        except Exception:
            pass

        return 1
    finally:
        if context_token is not None:
            otel_context.detach(context_token)


if __name__ == "__main__":
    raise SystemExit(main())
