"""AgentBill Distributed Tracing

Functions for cross-service trace propagation and correlation.
This is the SINGLE SOURCE OF TRUTH for trace context in the Python SDK.

P3-B: W3C Trace Context compliance
  - traceparent: 00-{32-hex-trace_id}-{16-hex-span_id}-{2-hex-flags}
  - tracestate: agentbill=customer_id:{cid},workflow_id:{wid}
  - All IDs are validated/normalized to correct hex lengths
  - [LEGACY_COMPAT:P3-B] agentbill-v1-* token format retained for backward compat

To find all P3-B legacy code: grep -rn "LEGACY_COMPAT:P3-B" sdks/python/
"""
import re
import os
import uuid
import secrets
import contextvars
from typing import Optional, Dict, Any, List

# Context variable to store the current tracing token
_tracing_token: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
    'agentbill_tracing_token', default=None
)

# Context variable to store trace context (trace_id, span_id, parent_span_id)
_trace_context: contextvars.ContextVar[Optional[Dict[str, str]]] = contextvars.ContextVar(
    'agentbill_trace_context', default=None
)

# Context variable for tracestate key-value pairs
_tracestate: contextvars.ContextVar[Optional[Dict[str, str]]] = contextvars.ContextVar(
    'agentbill_tracestate', default=None
)

# --- W3C Trace Context constants ---
_W3C_VERSION = "00"
_W3C_TRACE_FLAGS_SAMPLED = "01"
_W3C_TRACE_FLAGS_NOT_SAMPLED = "00"
_TRACEPARENT_RE = re.compile(
    r"^([0-9a-f]{2})-([0-9a-f]{32})-([0-9a-f]{16})-([0-9a-f]{2})$"
)
_HEX_RE = re.compile(r"^[0-9a-f]+$")


# --- ID generation & validation helpers ---

def _generate_trace_id() -> str:
    """Generate a W3C-compliant 32-hex-char trace ID."""
    return uuid.uuid4().hex  # 32 hex chars


def _generate_span_id() -> str:
    """Generate a W3C-compliant 16-hex-char span ID."""
    return secrets.token_hex(8)  # 16 hex chars


def _normalize_trace_id(raw: str) -> str:
    """Normalize a trace ID to 32 lowercase hex chars.
    
    If shorter, left-pads with zeros. If longer, truncates to 32.
    If not valid hex, generates a new ID.
    """
    raw = raw.lower().strip()
    if not _HEX_RE.match(raw):
        return _generate_trace_id()
    if len(raw) < 32:
        return raw.zfill(32)
    return raw[:32]


def _normalize_span_id(raw: str) -> str:
    """Normalize a span ID to 16 lowercase hex chars.
    
    If shorter, left-pads with zeros. If longer, truncates to 16.
    If not valid hex, generates a new ID.
    """
    raw = raw.lower().strip()
    if not _HEX_RE.match(raw):
        return _generate_span_id()
    if len(raw) < 16:
        return raw.zfill(16)
    return raw[:16]


def _is_valid_trace_id(trace_id: str) -> bool:
    """Check if a trace ID is valid (32 hex, not all zeros)."""
    return bool(_HEX_RE.match(trace_id)) and len(trace_id) == 32 and trace_id != "0" * 32


def _is_valid_span_id(span_id: str) -> bool:
    """Check if a span ID is valid (16 hex, not all zeros)."""
    return bool(_HEX_RE.match(span_id)) and len(span_id) == 16 and span_id != "0" * 16


# --- Tracestate helpers ---

def _encode_tracestate(baggage: Dict[str, str]) -> str:
    """Encode AgentBill baggage into a tracestate header value.
    
    Format: agentbill=key1:val1;key2:val2,other-vendor=value
    Only the agentbill vendor key is managed; other vendors are preserved.
    """
    if not baggage:
        return ""
    pairs = ";".join(f"{k}:{v}" for k, v in baggage.items() if v)
    return f"agentbill={pairs}" if pairs else ""


def _decode_tracestate(header: str) -> Dict[str, str]:
    """Decode a tracestate header, extracting AgentBill vendor values.
    
    Returns dict of AgentBill-specific key-value pairs.
    """
    result: Dict[str, str] = {}
    if not header:
        return result
    
    for member in header.split(","):
        member = member.strip()
        if member.startswith("agentbill="):
            payload = member[len("agentbill="):]
            for pair in payload.split(";"):
                pair = pair.strip()
                if ":" in pair:
                    k, v = pair.split(":", 1)
                    result[k.strip()] = v.strip()
    return result


def _merge_tracestate(incoming: str, agentbill_baggage: Dict[str, str]) -> str:
    """Merge AgentBill baggage into an existing tracestate header.
    
    AgentBill entry is placed first (highest priority per W3C spec).
    Other vendor entries are preserved.
    """
    other_members: List[str] = []
    if incoming:
        for member in incoming.split(","):
            member = member.strip()
            if not member.startswith("agentbill="):
                other_members.append(member)
    
    ab_entry = _encode_tracestate(agentbill_baggage)
    parts = [ab_entry] + other_members if ab_entry else other_members
    return ",".join(parts)


# --- Public API ---

def generate_tracing_token() -> str:
    """[LEGACY_COMPAT:P3-B] Generate an agentbill-v1-* tracing token.
    
    Retained for backward compatibility with services that use the
    proprietary token format. New integrations should use
    propagate_trace_headers() which emits standard W3C headers.
    
    Target removal: after all consumers migrate to W3C headers.
    
    Returns:
        str: A tracing token for backward-compatible propagation
    """
    ctx = _trace_context.get()
    random_suffix = secrets.token_hex(4)
    
    if ctx:
        trace_id = ctx.get("trace_id", "")
        span_id = ctx.get("span_id", "")
        token = f"agentbill-v1-{trace_id}-{span_id}-{random_suffix}"
    else:
        trace_id = _generate_trace_id()
        span_id = _generate_span_id()
        _trace_context.set({
            "trace_id": trace_id,
            "span_id": span_id
        })
        token = f"agentbill-v1-{trace_id}-{span_id}-{random_suffix}"
    
    _tracing_token.set(token)
    return token


def set_tracing_token(token: str) -> bool:
    """[LEGACY_COMPAT:P3-B] Set context from an agentbill-v1-* token.
    
    Retained for backward compatibility. Normalizes extracted IDs
    to W3C-compliant lengths.
    
    Target removal: after all consumers migrate to W3C headers.
    """
    if not token or not token.startswith("agentbill-v1-"):
        return False
    
    try:
        parts = token.split("-")
        if len(parts) >= 5:
            trace_id = _normalize_trace_id(parts[2])
            span_id = _normalize_span_id(parts[3])
            
            _trace_context.set({
                "trace_id": trace_id,
                "span_id": span_id,
                "parent_token": token
            })
            _tracing_token.set(token)
            return True
    except Exception:
        pass
    
    return False


def get_tracing_token() -> Optional[str]:
    """Get the current tracing token.
    
    Returns:
        Optional[str]: The current tracing token or None
    """
    return _tracing_token.get()


def get_trace_context() -> Optional[Dict[str, str]]:
    """Get the current trace context (trace_id, span_id).
    
    IDs are guaranteed to be W3C-compliant:
    - trace_id: 32 lowercase hex characters
    - span_id: 16 lowercase hex characters
    
    Returns:
        Optional[Dict[str, str]]: Dict with trace_id and span_id, or None
    """
    return _trace_context.get()


def clear_trace_context() -> None:
    """Clear the current trace context, tracing token, and tracestate."""
    _trace_context.set(None)
    _tracing_token.set(None)
    _tracestate.set(None)


def set_trace_context(trace_id: str, span_id: str, parent_span_id: Optional[str] = None) -> None:
    """Set the trace context with W3C-normalized IDs.
    
    Args:
        trace_id: The trace ID (will be normalized to 32 hex chars)
        span_id: The span ID (will be normalized to 16 hex chars)
        parent_span_id: Optional parent span ID (normalized to 16 hex chars)
    """
    ctx: Dict[str, str] = {
        "trace_id": _normalize_trace_id(trace_id),
        "span_id": _normalize_span_id(span_id)
    }
    if parent_span_id:
        ctx["parent_span_id"] = _normalize_span_id(parent_span_id)
    _trace_context.set(ctx)


def set_tracestate_baggage(baggage: Dict[str, str]) -> None:
    """Set AgentBill-specific tracestate baggage.
    
    This baggage is propagated in the W3C tracestate header under
    the 'agentbill' vendor key.
    
    Common keys: customer_id, workflow_id, scope_id, correlation_id
    
    Args:
        baggage: Key-value pairs to propagate
        
    Example:
        >>> set_tracestate_baggage({
        ...     "customer_id": "cust_abc",
        ...     "workflow_id": "wf_123"
        ... })
    """
    _tracestate.set(baggage)


def get_tracestate_baggage() -> Optional[Dict[str, str]]:
    """Get the current AgentBill tracestate baggage.
    
    Returns:
        Optional[Dict[str, str]]: Current baggage or None
    """
    return _tracestate.get()


def format_traceparent() -> Optional[str]:
    """Format the current trace context as a W3C traceparent header.
    
    Format: {version}-{trace_id}-{span_id}-{trace_flags}
    Example: 00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01
    
    Returns:
        Optional[str]: The traceparent header value, or None if no context
    """
    ctx = _trace_context.get()
    if not ctx:
        return None
    
    trace_id = ctx.get("trace_id", "")
    span_id = ctx.get("span_id", "")
    
    if not _is_valid_trace_id(trace_id) or not _is_valid_span_id(span_id):
        return None
    
    return f"{_W3C_VERSION}-{trace_id}-{span_id}-{_W3C_TRACE_FLAGS_SAMPLED}"


def parse_traceparent(header: str) -> Optional[Dict[str, str]]:
    """Parse a W3C traceparent header into its components.
    
    Args:
        header: The traceparent header value
        
    Returns:
        Optional[Dict[str, str]]: Dict with version, trace_id, span_id, 
        trace_flags or None if invalid
    """
    if not header:
        return None
    
    match = _TRACEPARENT_RE.match(header.strip().lower())
    if not match:
        return None
    
    version, trace_id, span_id, trace_flags = match.groups()
    
    # Per W3C spec: all-zero trace_id or span_id is invalid
    if trace_id == "0" * 32 or span_id == "0" * 16:
        return None
    
    return {
        "version": version,
        "trace_id": trace_id,
        "span_id": span_id,
        "trace_flags": trace_flags
    }


def propagate_trace_headers() -> Dict[str, str]:
    """Get HTTP headers for trace propagation (W3C + legacy).
    
    Returns both W3C standard headers and the legacy AgentBill token.
    
    Headers returned:
    - traceparent: W3C Trace Context (always, if context exists)
    - tracestate: W3C tracestate with AgentBill baggage (if baggage set)
    - X-AgentBill-Trace: [LEGACY_COMPAT:P3-B] legacy token (if set)
    
    [CONTEXT_UPGRADE:P3-C] Automatically merges rich context baggage into tracestate.
    
    Returns:
        Dict[str, str]: Headers to add to outgoing HTTP requests
    """
    headers: Dict[str, str] = {}
    
    # [CONTEXT_UPGRADE:P3-C] Merge rich context baggage into tracestate
    try:
        from .context import get_context_baggage
        ctx_baggage = get_context_baggage()
        if ctx_baggage:
            # Merge with any existing tracestate baggage
            existing = _tracestate.get() or {}
            merged = {**existing, **ctx_baggage}
            _tracestate.set(merged)
    except ImportError:
        pass  # context module not available (shouldn't happen)
    
    # W3C traceparent (primary)
    traceparent = format_traceparent()
    if traceparent:
        headers["traceparent"] = traceparent
    
    # W3C tracestate (if baggage exists)
    baggage = _tracestate.get()
    if baggage:
        tracestate = _encode_tracestate(baggage)
        if tracestate:
            headers["tracestate"] = tracestate
    
    # [LEGACY_COMPAT:P3-B] Legacy AgentBill token header
    # Target removal: after all consumers migrate to W3C headers
    token = get_tracing_token()
    if token:
        headers["X-AgentBill-Trace"] = token
    
    return headers


def extract_trace_from_headers(headers: Dict[str, str]) -> bool:
    """Extract and set trace context from incoming HTTP headers.
    
    Priority order:
    1. W3C traceparent (preferred, standards-compliant)
    2. [LEGACY_COMPAT:P3-B] AgentBill token (backward compat)
    
    Also extracts tracestate baggage if present.
    
    Args:
        headers: HTTP headers dict (case-insensitive lookup)
        
    Returns:
        bool: True if trace context was extracted and set
    """
    # Normalize header keys to lowercase for case-insensitive lookup
    norm = {k.lower(): v for k, v in headers.items()}
    
    # 1. Try W3C traceparent first (preferred)
    traceparent = norm.get("traceparent")
    if traceparent:
        parsed = parse_traceparent(traceparent)
        if parsed:
            set_trace_context(parsed["trace_id"], parsed["span_id"])
            
            # Extract tracestate if present
            tracestate = norm.get("tracestate", "")
            if tracestate:
                baggage = _decode_tracestate(tracestate)
                if baggage:
                    _tracestate.set(baggage)
            
            return True
    
    # 2. [LEGACY_COMPAT:P3-B] Fall back to AgentBill token
    # Target removal: after all consumers migrate to W3C headers
    token = norm.get("x-agentbill-trace")
    if token and set_tracing_token(token):
        return True
    
    return False
