from __future__ import annotations

from typing import Any, Mapping, Optional

from ..context import get_context, set_context
from ..middleware import get_request_id, get_session_id, set_request_id, set_session_id


def _header_get(headers: Any, key: str) -> Optional[str]:
    """Case-insensitive header fetch for mappings / HTTPMessage."""
    if headers is None:
        return None
    if hasattr(headers, "get"):
        value = headers.get(key)
        if value:
            return str(value)
        value = headers.get(key.lower())
        if value:
            return str(value)
    try:
        for k, v in headers.items():
            if str(k).lower() == key.lower():
                return str(v)
    except Exception:
        return None
    return None


def _tracking_headers(extra_headers: Optional[Mapping[str, str]] = None) -> dict[str, str]:
    """Build outbound headers with current tracking context."""
    ctx = get_context()
    request_id = get_request_id()
    session_id = get_session_id()
    trace_id = ctx.trace_id

    # Fallbacks so IDs still flow even if only one source was set.
    if not request_id:
        request_id = trace_id or ""
    if not trace_id:
        trace_id = request_id or ""
    if not session_id:
        session_id = ctx.session_id or ""

    headers: dict[str, str] = {}
    if request_id:
        headers["X-Request-ID"] = str(request_id)
    if session_id:
        headers["X-Session-ID"] = str(session_id)
    if trace_id:
        headers["X-Trace-ID"] = str(trace_id)

    if extra_headers:
        for k, v in extra_headers.items():
            headers[str(k)] = str(v)

    return headers


def _apply_response_ids(headers: Any) -> dict[str, Optional[str]]:
    """Read response IDs and backfill missing local context without replacing active request lineage."""
    ctx = get_context()

    current_request_id = get_request_id()
    current_session_id = get_session_id() or ctx.session_id
    current_trace_id = ctx.trace_id

    response_request_id = _header_get(headers, "X-Request-ID")
    response_trace_id = _header_get(headers, "X-Trace-ID")
    response_session_id = _header_get(headers, "X-Session-ID")

    # Preserve existing request/trace/session for end-to-end lineage; only backfill when missing.
    effective_request_id = current_request_id or response_request_id
    effective_trace_id = current_trace_id or response_trace_id or response_request_id
    effective_session_id = current_session_id or response_session_id

    if effective_request_id:
        set_request_id(effective_request_id)
    if effective_session_id:
        set_session_id(effective_session_id)
    if effective_trace_id or effective_session_id:
        set_context(trace_id=effective_trace_id, session_id=effective_session_id)

    return {
        "trace_id": effective_trace_id,
        "request_id": effective_request_id,
        "session_id": effective_session_id,
    }
