"""
Cross-service context propagation via HTTP headers.

When ``propagate_context=True`` is set on :class:`AIConfig`, provider
wrappers inject the active session/agent context into outgoing HTTP
headers.  Downstream services running the same SDK (or the ASGI
middleware) can pick up this context, linking the sessions into a
single distributed trace.

Headers emitted:

- ``traceparent`` — W3C Trace Context (version-trace_id-parent_id-flags).
- ``x-amplitude-session-id`` — the active session ID.
- ``x-amplitude-agent-id`` — the active agent ID (if set).

The :func:`extract_context` function parses these headers on the
receiving side.  :class:`AmplitudeAIMiddleware` already calls it
internally.
"""

from __future__ import annotations

import uuid
from typing import Any

from .context import get_active_context


def inject_context(headers: dict[str, str] | None = None) -> dict[str, str]:
    """Inject the active session context into *headers*.

    If no session is active, returns *headers* unchanged (or an empty dict).
    The returned dict is always a new copy — the original is never mutated.

    Args:
        headers: Existing headers to merge into (optional).

    Returns:
        A dict with context headers added.
    """
    result = dict(headers) if headers else {}

    ctx = get_active_context()
    if ctx is None:
        return result

    trace_id = ctx.trace_id or str(uuid.uuid4())
    # W3C traceparent: version-trace_id-parent_id-flags
    # Use a zero parent_id (this is the originating span).
    hex_trace = trace_id.replace("-", "")[:32].ljust(32, "0")
    parent_id = uuid.uuid4().hex[:16]
    result["traceparent"] = f"00-{hex_trace}-{parent_id}-01"

    if ctx.session_id:
        result["x-amplitude-session-id"] = ctx.session_id
    if ctx.agent_id:
        result["x-amplitude-agent-id"] = ctx.agent_id
    if ctx.user_id:
        result["x-amplitude-user-id"] = ctx.user_id

    return result


def extract_context(headers: dict[str, Any]) -> dict[str, Any]:
    """Parse context fields from incoming HTTP headers.

    Returns a dict with keys that can be passed to
    :class:`~amplitude_ai.context.SessionContext` or used directly.

    Handles:
    - ``traceparent`` (W3C) — extracts ``trace_id``.
    - ``x-trace-id`` — custom header fallback.
    - ``x-amplitude-session-id`` — session ID.
    - ``x-amplitude-agent-id`` — agent ID.
    - ``x-amplitude-user-id`` — user ID.
    """
    result: dict[str, Any] = {}

    # Trace ID: prefer traceparent, fall back to x-trace-id.
    traceparent = headers.get("traceparent", "")
    if traceparent:
        parts = traceparent.split("-")
        if len(parts) >= 2:
            result["trace_id"] = parts[1]
    if "trace_id" not in result:
        x_trace = headers.get("x-trace-id")
        if x_trace:
            result["trace_id"] = x_trace

    for header, key in [
        ("x-amplitude-session-id", "session_id"),
        ("x-amplitude-agent-id", "agent_id"),
        ("x-amplitude-user-id", "user_id"),
    ]:
        val = headers.get(header)
        if val:
            result[key] = val

    return result
