"""
Decorators for automatic tracking of function calls as tool events.
"""

from __future__ import annotations

import asyncio
import functools
import inspect
import json
import threading
import time
from typing import Any, Callable, Optional
from amplitude import Amplitude

import uuid

from .tracking import track_tool_call, track_span
from .privacy import PrivacyConfig
from ..context import SessionContext, get_active_context, _current_session
from ..utils.logger import get_logger


# ---------------------------------------------------------------------------
# JSON Schema generation from type hints
# ---------------------------------------------------------------------------

_TYPE_TO_JSON_SCHEMA: dict[type, str] = {
    str: "string",
    int: "integer",
    float: "number",
    bool: "boolean",
    list: "array",
    dict: "object",
    type(None): "null",
}


_STR_TYPE_TO_JSON_SCHEMA: dict[str, str] = {
    "str": "string",
    "int": "integer",
    "float": "number",
    "bool": "boolean",
    "list": "array",
    "dict": "object",
    "None": "null",
    "NoneType": "null",
}


def _resolve_annotation(ann: Any) -> str:
    """Map a type annotation to a JSON Schema type string.

    Handles both live types (``int``) and PEP 563 stringified annotations
    (``'int'``), plus generic aliases like ``list[str]``.
    """
    if ann is inspect.Parameter.empty:
        return "string"

    # PEP 563: annotations are strings when `from __future__ import annotations`
    if isinstance(ann, str):
        # Strip generics: "list[str]" → "list", "Optional[str]" → "Optional"
        base = ann.split("[", 1)[0].strip()
        return _STR_TYPE_TO_JSON_SCHEMA.get(base, "string")

    if ann in _TYPE_TO_JSON_SCHEMA:
        return _TYPE_TO_JSON_SCHEMA[ann]

    # Handle generic types like list[str], dict[str, int]
    origin = getattr(ann, "__origin__", None)
    if origin in _TYPE_TO_JSON_SCHEMA:
        return _TYPE_TO_JSON_SCHEMA[origin]

    return "string"


def _build_input_schema(fn: Callable[..., Any]) -> dict[str, Any]:
    """Build a JSON Schema-like dict from *fn*'s type annotations.

    Inspects ``inspect.signature(fn)`` and maps each parameter's annotation
    to a JSON Schema type.  Parameters without annotations are typed as
    ``"string"`` (safest default for display purposes).

    Handles both live types and PEP 563 stringified annotations.
    """
    sig = inspect.signature(fn)
    properties: dict[str, Any] = {}
    required: list[str] = []

    for pname, param in sig.parameters.items():
        if pname == "self":
            continue

        prop: dict[str, Any] = {"type": _resolve_annotation(param.annotation)}

        if param.default is not inspect.Parameter.empty:
            prop["default"] = param.default
        else:
            required.append(pname)

        properties[pname] = prop

    schema: dict[str, Any] = {"type": "object", "properties": properties}
    if required:
        schema["required"] = required
    return schema


class ToolCallTracker:
    """
    Global configuration for tool call tracking.

    Usage:
        amplitude = Amplitude(api_key="your_api_key")
        ToolCallTracker.set_amplitude(amplitude, user_id="user123")

        @tool
        def my_function(arg1, arg2):
            return result

    .. warning:: **Thread safety**

       All state is stored as class-level attributes (i.e., global).
       Calling :meth:`set_amplitude` in one thread mutates the config
       visible to *every* thread.  This is acceptable for single-user
       CLI scripts and notebooks but **not safe** for multi-tenant web
       servers where different requests serve different users.

       In multi-threaded / multi-request environments, pass
       ``amplitude``, ``user_id``, and other parameters directly to the
       ``@tool`` decorator instead of relying on the global tracker.
    """

    _amplitude: Optional[Amplitude] = None
    _user_id: Optional[str] = None
    _session_id: Optional[str] = None
    _trace_id: Optional[str] = None
    _turn_id: int = 1
    _env: Optional[str] = None
    _agent_id: Optional[str] = None
    _parent_agent_id: Optional[str] = None
    _customer_org_id: Optional[str] = None
    _agent_version: Optional[str] = None
    _context: Optional[dict[str, Any]] = None
    _privacy_config: Optional[PrivacyConfig] = None
    _event_properties: Optional[dict[str, Any]] = None
    _user_properties: Optional[dict[str, Any]] = None
    _groups: Optional[dict[str, Any]] = None

    @classmethod
    def set_amplitude(
        cls,
        amplitude: Amplitude,
        user_id: str,
        session_id: Optional[str] = None,
        trace_id: Optional[str] = None,
        turn_id: int = 1,
        env: Optional[str] = None,
        agent_id: Optional[str] = None,
        parent_agent_id: Optional[str] = None,
        customer_org_id: Optional[str] = None,
        agent_version: Optional[str] = None,
        context: Optional[dict[str, Any]] = None,
        privacy_config: Optional[PrivacyConfig] = None,
        event_properties: Optional[dict[str, Any]] = None,
        user_properties: Optional[dict[str, Any]] = None,
        groups: Optional[dict[str, Any]] = None,
    ) -> None:
        """
        Configure the global tracking settings for the @tool decorator.

        Args:
            amplitude: Amplitude client instance
            user_id: User identifier
            session_id: Session identifier for grouping related events
            trace_id: Trace identifier for linking events in a request chain
            turn_id: Turn number in conversation (default: 1)
            env: Environment (prod, staging, dev)
            agent_id: Agent/subproject identifier
            parent_agent_id: Delegating agent (multi-agent orchestration)
            customer_org_id: End-customer org (multi-tenant platforms)
            agent_version: Agent code version (e.g., "2.1.0")
            context: Arbitrary segmentation context dict
            privacy_config: Privacy configuration for content sanitization
            event_properties: Additional event properties
            user_properties: User properties to set
            groups: Group analytics properties
        """
        cls._amplitude = amplitude
        cls._user_id = user_id
        cls._session_id = session_id
        cls._trace_id = trace_id
        cls._turn_id = turn_id
        cls._env = env
        cls._agent_id = agent_id
        cls._parent_agent_id = parent_agent_id
        cls._customer_org_id = customer_org_id
        cls._agent_version = agent_version
        cls._context = context
        cls._privacy_config = privacy_config
        cls._event_properties = event_properties
        cls._user_properties = user_properties
        cls._groups = groups

    @classmethod
    def clear(cls) -> None:
        """Clear all tracking configuration."""
        cls._amplitude = None
        cls._user_id = None
        cls._session_id = None
        cls._trace_id = None
        cls._turn_id = 1
        cls._env = None
        cls._agent_id = None
        cls._parent_agent_id = None
        cls._customer_org_id = None
        cls._agent_version = None
        cls._context = None
        cls._privacy_config = None
        cls._event_properties = None
        cls._user_properties = None
        cls._groups = None

    @classmethod
    def is_configured(cls) -> bool:
        """Check if tracking is configured."""
        return cls._amplitude is not None and cls._user_id is not None


def tool(
    func: Optional[Callable[..., Any]] = None,
    *,
    name: Optional[str] = None,
    amplitude: Optional[Amplitude] = None,
    user_id: Optional[str] = None,
    session_id: Optional[str] = None,
    trace_id: Optional[str] = None,
    turn_id: Optional[int] = None,
    env: Optional[str] = None,
    agent_id: Optional[str] = None,
    parent_agent_id: Optional[str] = None,
    customer_org_id: Optional[str] = None,
    agent_version: Optional[str] = None,
    context: Optional[dict[str, Any]] = None,
    privacy_config: Optional[PrivacyConfig] = None,
    event_properties: Optional[dict[str, Any]] = None,
    user_properties: Optional[dict[str, Any]] = None,
    groups: Optional[dict[str, Any]] = None,
    input_schema: bool = False,
    timeout_ms: Optional[float] = None,
    on_error: Optional[Callable[[Exception, str], None]] = None,
) -> Callable[..., Any]:
    """
    Decorator to automatically track function calls as tool call events.

    Usage:

    1. With global configuration (recommended for multiple tools)::

        from amplitude import Amplitude
        from amplitude_ai import tool, ToolCallTracker

        amplitude = Amplitude(api_key="your_api_key")
        ToolCallTracker.set_amplitude(
            amplitude,
            user_id="user123",
            session_id="session_abc",
            trace_id="trace_xyz"
        )

        @tool
        def search_database(query: str) -> list:
            return ["result1", "result2"]

    2. With per-function configuration::

        @tool(amplitude=client, user_id="user123", session_id="session_abc")
        def my_function(arg1, arg2):
            return result

    3. With per-call runtime overrides (``amplitude_*`` kwargs)::

        @tool(amplitude=client)
        def search_docs(query: str) -> str:
            return "results..."

        # amplitude_user_id is extracted for tracking, not passed to search_docs
        result = search_docs(query="retention", amplitude_user_id="user-1")

       Supported runtime kwargs: ``amplitude_user_id``,
       ``amplitude_session_id``, ``amplitude_trace_id``,
       ``amplitude_turn_id``, ``amplitude_env``, ``amplitude_agent_id``,
       ``amplitude_parent_agent_id``, ``amplitude_customer_org_id``,
       ``amplitude_agent_version``.
       These override decorator params and global config (highest priority).

    4. With ``SessionContext`` auto-inheritance (recommended inside sessions)::

        Inside an active :class:`~amplitude_ai.client.Session`, the
        decorator automatically inherits ``session_id``, ``trace_id``,
        ``agent_id``, ``agent_version``, ``context``, and other fields
        from :func:`~amplitude_ai.context.get_active_context`.  Explicit
        decorator/runtime params still take priority.

    Args:
        func: The function to decorate (auto-provided when used as @tool)
        name: Custom name for the tool (defaults to function name)
        amplitude: Amplitude client instance (overrides global config)
        user_id: User identifier (overrides global config)
        session_id: Session identifier (overrides global config)
        trace_id: Trace identifier (overrides global config)
        turn_id: Turn number (overrides global config)
        env: Environment - prod, staging, dev (overrides global config)
        agent_id: Agent identifier (overrides global config)
        parent_agent_id: Delegating agent (overrides global config)
        customer_org_id: End-customer org (overrides global config)
        agent_version: Agent code version (overrides global config)
        context: Arbitrary segmentation context dict (overrides global config)
        privacy_config: Privacy configuration (overrides global config)
        event_properties: Additional event properties
        user_properties: User properties
        groups: Group properties
        input_schema: When *True*, auto-generate a JSON Schema from the
            function's type annotations and include it as
            ``[Agent] Tool Input Schema`` on the event.
        timeout_ms: Maximum execution time in milliseconds.  If exceeded,
            the function is interrupted, ``is_error=True`` is set, and
            ``error_message`` is ``"timeout"``.  Uses
            :func:`asyncio.wait_for` for async functions and
            :class:`threading.Timer` + :func:`_thread.interrupt_main`
            for sync functions.
        on_error: Callback ``(error, tool_name) -> None`` invoked when
            the tool raises an exception, *before* the error event is
            tracked.  Useful for alerting or fallback logic.

    Returns:
        Decorated function that tracks calls as tool events
    """

    def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
        tool_name = name or fn.__name__
        is_async = inspect.iscoroutinefunction(fn)

        # Pre-compute schema once at decoration time if requested.
        schema_json: Optional[str] = None
        if input_schema:
            try:
                schema_json = json.dumps(_build_input_schema(fn), default=str)
            except Exception:
                pass

        if is_async:
            @functools.wraps(fn)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                return await _execute_tracked_call_async(
                    fn, tool_name, args, kwargs,
                    amplitude, user_id, session_id, trace_id, turn_id,
                    env, agent_id, parent_agent_id, customer_org_id,
                    agent_version, context, privacy_config,
                    event_properties, user_properties, groups,
                    schema_json=schema_json,
                    timeout_ms=timeout_ms,
                    on_error=on_error,
                )
            return async_wrapper
        else:
            @functools.wraps(fn)
            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                return _execute_tracked_call_sync(
                    fn, tool_name, args, kwargs,
                    amplitude, user_id, session_id, trace_id, turn_id,
                    env, agent_id, parent_agent_id, customer_org_id,
                    agent_version, context, privacy_config,
                    event_properties, user_properties, groups,
                    schema_json=schema_json,
                    timeout_ms=timeout_ms,
                    on_error=on_error,
                )
            return sync_wrapper

    if func is not None:
        return decorator(func)
    return decorator


_AMPLITUDE_KWARG_PREFIX = "amplitude_"

# Mapping from runtime kwarg suffix → (decorator param name, supports None turn_id)
_RUNTIME_KWARG_MAP: dict[str, str] = {
    "user_id": "user_id",
    "session_id": "session_id",
    "trace_id": "trace_id",
    "turn_id": "turn_id",
    "env": "env",
    "agent_id": "agent_id",
    "parent_agent_id": "parent_agent_id",
    "customer_org_id": "customer_org_id",
    "agent_version": "agent_version",
}


def _extract_amplitude_kwargs(
    kwargs: dict[str, Any],
) -> tuple[dict[str, Any], dict[str, Any]]:
    """Separate ``amplitude_*`` runtime overrides from function kwargs.

    Returns ``(amplitude_overrides, clean_kwargs)`` where
    ``clean_kwargs`` can be safely passed to the decorated function.
    """
    overrides: dict[str, Any] = {}
    clean: dict[str, Any] = {}
    for key, value in kwargs.items():
        if key.startswith(_AMPLITUDE_KWARG_PREFIX):
            suffix = key[len(_AMPLITUDE_KWARG_PREFIX):]
            if suffix in _RUNTIME_KWARG_MAP:
                overrides[_RUNTIME_KWARG_MAP[suffix]] = value
            else:
                # Not a recognised override — pass through to the function
                clean[key] = value
        else:
            clean[key] = value
    return overrides, clean


def _resolve_context_fields(
    overrides: dict[str, Any],
    *,
    amplitude: Optional[Amplitude],
    user_id: Optional[str],
    session_id: Optional[str],
    trace_id: Optional[str],
    turn_id: Optional[int],
    env: Optional[str],
    agent_id: Optional[str],
    parent_agent_id: Optional[str],
    customer_org_id: Optional[str],
    agent_version: Optional[str],
    context: Optional[dict[str, Any]],
    privacy_config: Optional[PrivacyConfig],
    event_properties: Optional[dict[str, Any]],
    user_properties: Optional[dict[str, Any]],
    groups: Optional[dict[str, Any]],
) -> dict[str, Any]:
    """Resolve final tracking parameters from overrides > decorator > SessionContext > global.

    Priority (highest first): runtime ``amplitude_*`` overrides, decorator
    params, active :class:`SessionContext`, :class:`ToolCallTracker` globals.

    Returns a dict with all resolved values keyed by their logical names.
    """
    ctx = get_active_context()

    def _pick(key: str, decorator_val: Any, tracker_val: Any, ctx_val: Any = None) -> Any:
        return overrides.get(key) or decorator_val or ctx_val or tracker_val

    amp = amplitude or ToolCallTracker._amplitude
    uid = _pick("user_id", user_id, ToolCallTracker._user_id, getattr(ctx, "user_id", None))
    sess_id = _pick("session_id", session_id, ToolCallTracker._session_id, getattr(ctx, "session_id", None))
    trc_id = _pick("trace_id", trace_id, ToolCallTracker._trace_id, getattr(ctx, "trace_id", None))
    _raw_turn = overrides.get("turn_id")
    t_id: int = int(_raw_turn) if _raw_turn is not None else (turn_id if turn_id is not None else ToolCallTracker._turn_id)
    environment = _pick("env", env, ToolCallTracker._env, getattr(ctx, "env", None))
    agent = _pick("agent_id", agent_id, ToolCallTracker._agent_id, getattr(ctx, "agent_id", None))
    p_agent = _pick("parent_agent_id", parent_agent_id, ToolCallTracker._parent_agent_id, getattr(ctx, "parent_agent_id", None))
    cust_org = _pick("customer_org_id", customer_org_id, ToolCallTracker._customer_org_id, getattr(ctx, "customer_org_id", None))
    a_ver = _pick("agent_version", agent_version, ToolCallTracker._agent_version, getattr(ctx, "agent_version", None))
    ctx_dict = context or (ToolCallTracker._context if ToolCallTracker._context else None) or getattr(ctx, "context", None)
    priv_cfg = privacy_config or ToolCallTracker._privacy_config
    evt_props = event_properties or ToolCallTracker._event_properties
    usr_props = user_properties or ToolCallTracker._user_properties
    grps = groups or ToolCallTracker._groups or getattr(ctx, "groups", None)

    return {
        "amplitude": amp,
        "user_id": uid,
        "session_id": sess_id,
        "trace_id": trc_id,
        "turn_id": t_id,
        "env": environment,
        "agent_id": agent,
        "parent_agent_id": p_agent,
        "customer_org_id": cust_org,
        "agent_version": a_ver,
        "context": ctx_dict,
        "privacy_config": priv_cfg,
        "event_properties": evt_props,
        "user_properties": usr_props,
        "groups": grps,
    }


def _execute_tracked_call_sync(
    fn: Callable[..., Any],
    tool_name: str,
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
    amplitude: Optional[Amplitude],
    user_id: Optional[str],
    session_id: Optional[str],
    trace_id: Optional[str],
    turn_id: Optional[int],
    env: Optional[str],
    agent_id: Optional[str],
    parent_agent_id: Optional[str],
    customer_org_id: Optional[str],
    agent_version: Optional[str],
    context: Optional[dict[str, Any]],
    privacy_config: Optional[PrivacyConfig],
    event_properties: Optional[dict[str, Any]],
    user_properties: Optional[dict[str, Any]],
    groups: Optional[dict[str, Any]],
    *,
    schema_json: Optional[str] = None,
    timeout_ms: Optional[float] = None,
    on_error: Optional[Callable[[Exception, str], None]] = None,
) -> Any:
    """Execute a tracked synchronous function call."""
    overrides, clean_kwargs = _extract_amplitude_kwargs(kwargs)

    r = _resolve_context_fields(
        overrides,
        amplitude=amplitude, user_id=user_id, session_id=session_id,
        trace_id=trace_id, turn_id=turn_id, env=env, agent_id=agent_id,
        parent_agent_id=parent_agent_id, customer_org_id=customer_org_id,
        agent_version=agent_version, context=context,
        privacy_config=privacy_config, event_properties=event_properties,
        user_properties=user_properties, groups=groups,
    )

    amp = r["amplitude"]
    uid = r["user_id"]

    if amp is None or uid is None:
        get_logger(amp).warning(
            f"Tool '{tool_name}' called but tracking not configured. "
            "Call ToolCallTracker.set_amplitude() or pass amplitude/user_id to decorator."
        )
        return fn(*args, **clean_kwargs)

    sig = inspect.signature(fn)
    bound_args = sig.bind(*args, **clean_kwargs)
    bound_args.apply_defaults()
    tool_input = {
        k: v for k, v in bound_args.arguments.items() if k != "self"
    }

    start_time = time.time()
    success = True
    error_msg: Optional[str] = None
    result = None

    try:
        if timeout_ms is not None:
            result = _run_with_sync_timeout(fn, args, clean_kwargs, timeout_ms)
        else:
            result = fn(*args, **clean_kwargs)
        return result
    except Exception as e:
        success = False
        error_msg = str(e)
        if on_error is not None:
            try:
                on_error(e, tool_name)
            except Exception:
                pass
        get_logger(amp).error(f"Tool '{tool_name}' failed: {error_msg}")
        raise
    finally:
        end_time = time.time()
        latency_ms = (end_time - start_time) * 1000

        extra_props: dict[str, Any] = {}
        if schema_json is not None:
            extra_props["[Agent] Tool Input Schema"] = schema_json

        merged_evt_props = {**(r["event_properties"] or {}), **extra_props}

        try:
            track_tool_call(
                amplitude=amp,
                user_id=uid,
                tool_name=tool_name,
                success=success,
                latency_ms=latency_ms,
                session_id=r["session_id"],
                trace_id=r["trace_id"],
                turn_id=r["turn_id"],
                tool_input=tool_input,
                tool_output=result if success else None,
                error_message=error_msg,
                env=r["env"],
                agent_id=r["agent_id"],
                parent_agent_id=r["parent_agent_id"],
                customer_org_id=r["customer_org_id"],
                agent_version=r["agent_version"],
                context=r["context"],
                event_properties=merged_evt_props if extra_props else r["event_properties"],
                user_properties=r["user_properties"],
                groups=r["groups"],
                privacy_config=r["privacy_config"],
            )
        except Exception as track_error:
            get_logger(amp).error(f"Failed to track tool call '{tool_name}': {track_error}")


async def _execute_tracked_call_async(
    fn: Callable[..., Any],
    tool_name: str,
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
    amplitude: Optional[Amplitude],
    user_id: Optional[str],
    session_id: Optional[str],
    trace_id: Optional[str],
    turn_id: Optional[int],
    env: Optional[str],
    agent_id: Optional[str],
    parent_agent_id: Optional[str],
    customer_org_id: Optional[str],
    agent_version: Optional[str],
    context: Optional[dict[str, Any]],
    privacy_config: Optional[PrivacyConfig],
    event_properties: Optional[dict[str, Any]],
    user_properties: Optional[dict[str, Any]],
    groups: Optional[dict[str, Any]],
    *,
    schema_json: Optional[str] = None,
    timeout_ms: Optional[float] = None,
    on_error: Optional[Callable[[Exception, str], None]] = None,
) -> Any:
    """Execute a tracked asynchronous function call."""
    overrides, clean_kwargs = _extract_amplitude_kwargs(kwargs)

    r = _resolve_context_fields(
        overrides,
        amplitude=amplitude, user_id=user_id, session_id=session_id,
        trace_id=trace_id, turn_id=turn_id, env=env, agent_id=agent_id,
        parent_agent_id=parent_agent_id, customer_org_id=customer_org_id,
        agent_version=agent_version, context=context,
        privacy_config=privacy_config, event_properties=event_properties,
        user_properties=user_properties, groups=groups,
    )

    amp = r["amplitude"]
    uid = r["user_id"]

    if amp is None or uid is None:
        get_logger(amp).warning(
            f"Tool '{tool_name}' called but tracking not configured. "
            "Call ToolCallTracker.set_amplitude() or pass amplitude/user_id to decorator."
        )
        return await fn(*args, **clean_kwargs)

    sig = inspect.signature(fn)
    bound_args = sig.bind(*args, **clean_kwargs)
    bound_args.apply_defaults()
    tool_input = {
        k: v for k, v in bound_args.arguments.items() if k != "self"
    }

    start_time = time.time()
    success = True
    error_msg: Optional[str] = None
    result = None

    try:
        coro = fn(*args, **clean_kwargs)
        if timeout_ms is not None:
            result = await asyncio.wait_for(coro, timeout=timeout_ms / 1000.0)
        else:
            result = await coro
        return result
    except asyncio.TimeoutError:
        success = False
        error_msg = "timeout"
        timeout_exc = TimeoutError(f"Tool '{tool_name}' exceeded {timeout_ms}ms timeout")
        if on_error is not None:
            try:
                on_error(timeout_exc, tool_name)
            except Exception:
                pass
        get_logger(amp).error(f"Tool '{tool_name}' timed out after {timeout_ms}ms")
        raise timeout_exc from None
    except Exception as e:
        success = False
        error_msg = str(e)
        if on_error is not None:
            try:
                on_error(e, tool_name)
            except Exception:
                pass
        get_logger(amp).error(f"Tool '{tool_name}' failed: {error_msg}")
        raise
    finally:
        end_time = time.time()
        latency_ms = (end_time - start_time) * 1000

        extra_props: dict[str, Any] = {}
        if schema_json is not None:
            extra_props["[Agent] Tool Input Schema"] = schema_json

        merged_evt_props = {**(r["event_properties"] or {}), **extra_props}

        try:
            track_tool_call(
                amplitude=amp,
                user_id=uid,
                tool_name=tool_name,
                success=success,
                latency_ms=latency_ms,
                session_id=r["session_id"],
                trace_id=r["trace_id"],
                turn_id=r["turn_id"],
                tool_input=tool_input,
                tool_output=result if success else None,
                error_message=error_msg,
                env=r["env"],
                agent_id=r["agent_id"],
                parent_agent_id=r["parent_agent_id"],
                customer_org_id=r["customer_org_id"],
                agent_version=r["agent_version"],
                context=r["context"],
                event_properties=merged_evt_props if extra_props else r["event_properties"],
                user_properties=r["user_properties"],
                groups=r["groups"],
                privacy_config=r["privacy_config"],
            )
        except Exception as track_error:
            get_logger(amp).error(f"Failed to track tool call '{tool_name}': {track_error}")


def _run_with_sync_timeout(
    fn: Callable[..., Any],
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
    timeout_ms: float,
) -> Any:
    """Run *fn* synchronously with a timeout.

    Uses a background thread to detect when the timeout expires and raises
    ``TimeoutError`` in the calling thread via ``_thread.interrupt_main()``.
    Falls back to a simpler approach if interrupt is not available.
    """
    result_box: list[Any] = []
    error_box: list[Exception] = []
    done = threading.Event()

    def _target() -> None:
        try:
            result_box.append(fn(*args, **kwargs))
        except Exception as exc:
            error_box.append(exc)
        finally:
            done.set()

    worker = threading.Thread(target=_target, daemon=True)
    worker.start()
    finished = done.wait(timeout=timeout_ms / 1000.0)

    if not finished:
        raise TimeoutError(f"Tool exceeded {timeout_ms}ms timeout")

    if error_box:
        raise error_box[0]
    return result_box[0] if result_box else None


# ======================================================================
# @observe decorator
# ======================================================================


def observe(
    func: Callable[..., Any] | None = None,
    *,
    name: str | None = None,
    amplitude: Amplitude | None = None,
    user_id: str | None = None,
    agent_id: str | None = None,
    env: str | None = None,
    privacy_config: PrivacyConfig | None = None,
) -> Any:
    """Decorator that tracks a function as a ``[Agent] Span`` event.

    Session-boundary rules:

    1. If a session context exists (from ``BoundAgent.session()``,
       middleware, or an outer ``@observe``), attaches to it --
       emits only a ``[Agent] Span`` event.
    2. If no session exists, creates a session scoped to the
       function call (session_id = new UUID), emits
       ``[Agent] Span`` on return, then ``[Agent] Session End``.
    3. Nested ``@observe`` calls attach as nested spans to the
       outer session -- no session explosion.

    Content-mode interaction: in ``metadata_only``, captures function
    name, latency, and error status only (no ``input_state`` / ``output_state``).
    """
    def _decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
        span_name = name or fn.__name__
        is_async = asyncio.iscoroutinefunction(fn)

        def _resolve_params() -> dict[str, Any]:
            """Merge decorator args → ContextVar → ToolCallTracker."""
            ctx = get_active_context()
            amp = amplitude or (ToolCallTracker._amplitude if ToolCallTracker._amplitude else None)
            uid = user_id or (ctx.user_id if ctx else None) or ToolCallTracker._user_id or ""
            aid = agent_id or (ctx.agent_id if ctx else None) or ToolCallTracker._agent_id
            env_ = env or (ctx.env if ctx else None) or ToolCallTracker._env
            pc = privacy_config or ToolCallTracker._privacy_config
            return {
                "amplitude": amp,
                "user_id": uid,
                "agent_id": aid,
                "env": env_,
                "privacy_config": pc,
                "context": ctx,
            }

        def _serialize_state(
            value: Any, pc: PrivacyConfig | None,
        ) -> dict[str, Any] | None:
            """Serialize function args/return for span state, respecting content_mode."""
            if value is None:
                return None
            if pc is not None:
                mode = pc.content_mode
                if mode == "metadata_only" or (mode is None and pc.privacy_mode):
                    return None
            if isinstance(value, dict):
                return value
            return {"value": str(value)}

        def _emit_span(
            params: dict[str, Any],
            ctx: SessionContext | None,
            session_id: str,
            latency_ms: float,
            is_error: bool,
            error_message: str | None,
            input_state: dict[str, Any] | None,
            output_state: dict[str, Any] | None,
        ) -> None:
            amp = params["amplitude"]
            if amp is None:
                return
            try:
                track_span(
                    amplitude=amp,
                    user_id=params["user_id"],
                    span_name=span_name,
                    trace_id=(ctx.trace_id if ctx and ctx.trace_id else None) or str(uuid.uuid4()),
                    latency_ms=latency_ms,
                    input_state=input_state,
                    output_state=output_state,
                    is_error=is_error,
                    error_message=error_message,
                    session_id=session_id,
                    agent_id=params["agent_id"] or span_name,
                    env=params["env"],
                    privacy_config=params["privacy_config"],
                )
            except Exception as e:
                get_logger(amp).error(f"Failed to track @observe span '{span_name}': {e}")

        def _emit_session_end(
            params: dict[str, Any], session_id: str,
        ) -> None:
            amp = params["amplitude"]
            if amp is None:
                return
            try:
                from .tracking import track_session_end
                track_session_end(
                    amplitude=amp,
                    user_id=params["user_id"],
                    session_id=session_id,
                    env=params["env"],
                    agent_id=params["agent_id"] or span_name,
                    privacy_config=params["privacy_config"],
                )
            except Exception as e:
                get_logger(amp).error(f"Failed to end @observe session '{span_name}': {e}")

        @functools.wraps(fn)
        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
            import uuid as _uuid
            params = _resolve_params()
            ctx = get_active_context()
            owns_session = ctx is None
            session_id = ctx.session_id if ctx else str(_uuid.uuid4())

            token = None
            if owns_session:
                new_ctx = SessionContext(
                    session_id=session_id,
                    user_id=params["user_id"],
                    agent_id=params["agent_id"] or span_name,
                    env=params["env"],
                )
                token = _current_session.set(new_ctx)
                ctx = new_ctx

            input_state = _serialize_state(
                dict(zip(inspect.signature(fn).parameters.keys(), args), **kwargs),
                params["privacy_config"],
            )

            start = time.time()
            is_error = False
            error_message = None
            result = None
            try:
                result = fn(*args, **kwargs)
                return result
            except Exception as exc:
                is_error = True
                error_message = str(exc)
                raise
            finally:
                latency_ms = (time.time() - start) * 1000.0
                output_state = _serialize_state(result, params["privacy_config"])
                _emit_span(
                    params, ctx, session_id, latency_ms,
                    is_error, error_message, input_state, output_state,
                )
                if owns_session:
                    _emit_session_end(params, session_id)
                    if token is not None:
                        _current_session.reset(token)

        @functools.wraps(fn)
        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
            import uuid as _uuid
            params = _resolve_params()
            ctx = get_active_context()
            owns_session = ctx is None
            session_id = ctx.session_id if ctx else str(_uuid.uuid4())

            token = None
            if owns_session:
                new_ctx = SessionContext(
                    session_id=session_id,
                    user_id=params["user_id"],
                    agent_id=params["agent_id"] or span_name,
                    env=params["env"],
                )
                token = _current_session.set(new_ctx)
                ctx = new_ctx

            input_state = _serialize_state(
                dict(zip(inspect.signature(fn).parameters.keys(), args), **kwargs),
                params["privacy_config"],
            )

            start = time.time()
            is_error = False
            error_message = None
            result = None
            try:
                result = await fn(*args, **kwargs)
                return result
            except Exception as exc:
                is_error = True
                error_message = str(exc)
                raise
            finally:
                latency_ms = (time.time() - start) * 1000.0
                output_state = _serialize_state(result, params["privacy_config"])
                _emit_span(
                    params, ctx, session_id, latency_ms,
                    is_error, error_message, input_state, output_state,
                )
                if owns_session:
                    _emit_session_end(params, session_id)
                    if token is not None:
                        _current_session.reset(token)

        return async_wrapper if is_async else sync_wrapper

    if func is not None:
        return _decorator(func)
    return _decorator


__all__ = ["tool", "ToolCallTracker", "observe"]
