"""OpenTelemetry Tracer for AgentBill

P3-A Migration: This module supports two modes controlled by config["otel_mode"]:
  - "legacy" (default): Manual OTLP JSON construction via httpx (current behavior)
  - "native": Uses native OTEL TracerProvider via otel_provider.py

All legacy-only code is annotated with [LEGACY_COMPAT:P3-A] for future removal.
To find all legacy code: grep -rn "LEGACY_COMPAT:P3-A" sdks/python/
"""
import json
import time
import uuid
import httpx
from typing import Dict, Any, Optional, Union
from .types import AgentBillConfig
from .span_processor import AgentBillSpanProcessor

# Resolve SDK version dynamically so service.version always matches the
# installed package (previously hardcoded to "7.16.1" — see bug report
# where v10.5.x installs were still reporting 7.16.1 in telemetry).
try:
    from importlib.metadata import version as _pkg_version, PackageNotFoundError
    try:
        _SDK_VERSION = _pkg_version("agentbill")
    except PackageNotFoundError:
        _SDK_VERSION = "unknown"
except Exception:
    _SDK_VERSION = "unknown"


class Span:
    """[LEGACY_COMPAT:P3-A] Legacy span class for manual OTLP mode.
    
    In native OTEL mode, NativeSpanWrapper (from otel_provider.py) is used instead.
    Both classes expose the same interface: trace_id, span_id, set_attributes(), 
    set_status(), end(). This ensures callers work unchanged regardless of mode.
    
    Target removal: after native OTEL mode is validated + 1 minor release.
    """
    
    def __init__(self, name: str, trace_id: str, span_id: str, parent_span_id: Optional[str], attributes: Dict[str, Any]):
        self.name = name
        self.trace_id = trace_id
        self.span_id = span_id
        self.parent_span_id = parent_span_id
        self.attributes = attributes
        self.start_time = time.time_ns()
        self.end_time: Optional[int] = None
        self.status = {"code": 0}
    
    def set_attributes(self, attributes: Dict[str, Any]):
        self.attributes.update(attributes)
    
    def set_status(self, code: int, message: str = ""):
        self.status = {"code": code, "message": message}
    
    def end(self):
        self.end_time = time.time_ns()


class AgentBillTracer:
    """OpenTelemetry tracer for AgentBill
    
    IMPORTANT: This tracer integrates with distributed.py for trace context.
    - If a trace context exists (via get_trace_context()), it reuses the trace_id
    - Always generates a new span_id for each span
    - Updates the global trace context so get_trace_context() returns correct values
    
    This ensures that:
    1. get_trace_context() always returns the trace_id that matches what's sent to the portal
    2. Distributed tracing works correctly across services
    3. Parent-child span relationships are properly maintained
    """
    
    def __init__(self, config: AgentBillConfig):
        self.config = config
        self.base_url = config.get("base_url", "https://api.agentbill.io")
        self.api_key = config["api_key"]
        self.customer_id = config.get("customer_id")
        self.debug = config.get("debug", False)
        
        # P3-A: OTEL mode switch — "legacy" (default) or "native"
        self._otel_mode = config.get("otel_mode", "legacy")
        
        if self._otel_mode == "native":
            # P3-A: Use native OTEL TracerProvider
            from .otel_provider import create_tracer_provider, OTEL_AVAILABLE
            if not OTEL_AVAILABLE:
                import warnings
                warnings.warn(
                    "otel_mode='native' requested but opentelemetry-sdk not installed. "
                    "Falling back to legacy mode. Install with: "
                    "pip install opentelemetry-sdk opentelemetry-exporter-otlp-proto-http",
                    RuntimeWarning,
                    stacklevel=2,
                )
                self._otel_mode = "legacy"
                self._provider = None
                self._native_tracer = None
            else:
                self._provider = create_tracer_provider(config)
                self._native_tracer = self._provider.get_tracer("agentbill")
        else:
            # [LEGACY_COMPAT:P3-A] No native provider in legacy mode
            self._provider = None
            self._native_tracer = None
        
        # [LEGACY_COMPAT:P3-A] Manual span buffer — only used in legacy mode.
        # In native mode, spans are managed by BatchSpanProcessor.
        self.spans = []
        
        # [SPAN_PROCESSOR:P3-D] Client-side span processor for identity injection,
        # PII redaction, deduplication, and sampling.
        self._span_processor = AgentBillSpanProcessor(
            customer_id=self.customer_id,
            agent_id=config.get("agent_id"),
            pii_redaction=config.get("pii_redaction"),
            dedup_window=config.get("dedup_window", 1000),
            sampling_rate=config.get("sampling_rate", 1.0),
            debug=self.debug,
        )
    
    def start_span(self, name: str, attributes: Dict[str, Any]) -> Union['Span', Any]:
        """Start a new span and return it with trace_id and span_id for correlation.
        
        Returns a Span (legacy mode) or NativeSpanWrapper (native mode).
        Both expose: trace_id, span_id, set_attributes(), set_status(), end().
        
        IMPORTANT: Uses distributed.py as the single source of truth for trace context.
        - Reuses existing trace_id if available (for span correlation within same trace)
        - Always generates new span_id (each span is unique)
        - Updates global trace context so get_trace_context() returns correct values
        """
        if self._otel_mode == "native":
            return self._start_span_native(name, attributes)
        
        # [LEGACY_COMPAT:P3-A] Manual span creation — entire method below.
        # Remove this branch after native OTEL mode is validated + 1 minor release.
        return self._start_span_legacy(name, attributes)
    
    def _start_span_native(self, name: str, attributes: Dict[str, Any]):
        """P3-A: Start span using native OTEL tracer."""
        from .distributed import get_trace_context, set_trace_context
        from .otel_provider import NativeSpanWrapper
        from .context import get_context_attributes  # [CONTEXT_UPGRADE:P3-C]
        from opentelemetry import trace as otel_trace, context as otel_context
        
        # [CONTEXT_UPGRADE:P3-C] Inject rich context attributes into span
        ctx_attrs = get_context_attributes()
        merged_attrs = {**ctx_attrs, **attributes}  # User attrs override context
        
        # Use OTEL context propagation — native tracer handles parent linking
        otel_span = self._native_tracer.start_span(name=name, attributes=merged_attrs)
        
        # Extract IDs for backward-compatible TraceContext
        span_ctx = otel_span.get_span_context()
        trace_id = format(span_ctx.trace_id, '032x')
        span_id = format(span_ctx.span_id, '016x')
        
        # Get parent info from distributed context
        existing_ctx = get_trace_context()
        parent_span_id = existing_ctx.get("span_id") if existing_ctx else None
        
        # Update distributed context for cross-service propagation
        set_trace_context(trace_id, span_id)
        
        if self.debug:
            print(f"[AgentBill Tracer] Started native span: {name}, trace_id={trace_id[:8]}..., span_id={span_id[:8]}...")
        
        return NativeSpanWrapper(otel_span, trace_id, span_id, parent_span_id)
    
    def _start_span_legacy(self, name: str, attributes: Dict[str, Any]) -> Span:
        """[LEGACY_COMPAT:P3-A] Manual span creation — remove after native OTEL validated."""
        from .distributed import get_trace_context, set_trace_context
        from .context import get_context_attributes  # [CONTEXT_UPGRADE:P3-C]
        
        # [CONTEXT_UPGRADE:P3-C] Inject rich context attributes into span
        ctx_attrs = get_context_attributes()
        attributes = {**ctx_attrs, **attributes}  # User attrs override context
        
        # [SPAN_PROCESSOR:P3-D] Identity injection + sampling decision
        attributes = self._span_processor.on_start(attributes)
        
        # Check for existing trace context
        ctx = get_trace_context()
        parent_span_id = None
        
        if ctx:
            # Reuse existing trace_id for correlation
            trace_id = ctx.get("trace_id")
            # Store parent span_id for proper OTEL hierarchy
            parent_span_id = ctx.get("span_id")
            if self.debug:
                print(f"[AgentBill Tracer] Using existing trace context: trace_id={trace_id[:8]}...")
        else:
            # Generate new trace_id for new trace
            trace_id = uuid.uuid4().hex
            if self.debug:
                print(f"[AgentBill Tracer] Creating new trace: trace_id={trace_id[:8]}...")
        
        # Always generate new span_id for this span
        span_id = uuid.uuid4().hex[:16]
        
        # Update global trace context so get_trace_context() returns correct values
        set_trace_context(trace_id, span_id)
        
        attributes["service.name"] = "agentbill-python-sdk"
        if self.customer_id:
            attributes["customer.id"] = self.customer_id
        
        span = Span(name, trace_id, span_id, parent_span_id, attributes)
        
        # [SPAN_PROCESSOR:P3-D] Process on_end is deferred to flush;
        # we still add to buffer here. on_end filtering happens in flush/export.
        self.spans.append(span)
        
        if self.debug:
            print(f"[AgentBill Tracer] Started span: {name}, trace_id={trace_id[:8]}..., span_id={span_id[:8]}...")
        
        # Return span which has trace_id and span_id as properties
        return span
    
    def flush_sync(self):
        """Synchronous flush for non-async contexts"""
        import asyncio
        try:
            loop = asyncio.get_event_loop()
            if loop.is_running():
                # If loop is running, schedule the task
                asyncio.create_task(self.flush())
            else:
                # If no loop, run it
                loop.run_until_complete(self.flush())
        except RuntimeError:
            # No event loop exists, create one
            asyncio.run(self.flush())
    
    async def flush(self):
        """Flush spans to AgentBill."""
        if self._otel_mode == "native":
            # P3-A: Native flush — BatchSpanProcessor handles export automatically.
            # force_flush ensures all buffered spans are sent immediately.
            if self._provider:
                self._provider.force_flush()
                if self.debug:
                    print("AgentBill: ✅ Native OTEL flush completed")
            return
        
        # [LEGACY_COMPAT:P3-A] Manual HTTP flush — remove after native OTEL validated.
        if not self.spans:
            if self.debug:
                print("AgentBill: No spans to flush")
            return
        
        # [SPAN_PROCESSOR:P3-D] Apply on_end processing (PII redaction, dedup, sampling)
        processed_spans = []
        for span in self.spans:
            span_dict = {
                "trace_id": span.trace_id,
                "span_id": span.span_id,
                "attributes": span.attributes,
            }
            result = self._span_processor.on_end(span_dict)
            if result is not None:
                # Write back possibly-redacted attributes
                span.attributes = result["attributes"]
                processed_spans.append(span)
            elif self.debug:
                print(f"AgentBill: Span {span.span_id[:8]}... dropped by processor")
        
        self.spans = processed_spans
        if not self.spans:
            if self.debug:
                print("AgentBill: All spans filtered by processor")
            return
        
        payload = self._build_otlp_payload()
        url = f"{self.base_url}/functions/v1/otel-collector"
        headers = {
            "X-API-Key": self.api_key,
            "Content-Type": "application/json",
            # PHASE3-R3-SDK-REQUESTID — client-generated by contract.
            # Must be stable across transport retries so the server WAL can dedupe.
            # Server-generated IDs would change per retry and break idempotency.
            # Contract: docs/internal/integration/TELEMETRY_CONTRACT.md.
            "X-Request-Id": __import__("uuid").uuid4().hex,
        }
        
        if self.debug:
            print(f"AgentBill: Flushing {len(self.spans)} spans to {url}")
            print(f"AgentBill: API Key: {self.api_key[:12]}...")
            print(f"AgentBill: Full headers being sent: {headers}")
            print(f"AgentBill: Payload preview: {str(payload)[:200]}...")
        
        async with httpx.AsyncClient() as client:
            try:
                response = await client.post(url, json=payload, headers=headers)
                
                if self.debug:
                    print(f"AgentBill flush response code: {response.status_code}")
                    print(f"AgentBill response headers: {dict(response.headers)}")
                    print(f"AgentBill response body: {response.text[:500]}")
                
                # Accept 200 (sync ingest) and 202 (async WAL ingest).
                # /otel-collector returns 202 with a wal_id when WAL is enabled.
                if response.status_code in (200, 202):
                    self.spans.clear()
                    if self.debug:
                        print("AgentBill: ✅ Spans successfully flushed")
                else:
                    if self.debug:
                        print(f"AgentBill: ❌ Flush failed with status {response.status_code}")
            except Exception as e:
                if self.debug:
                    print(f"AgentBill flush error: {type(e).__name__}: {e}")
    
    def _build_otlp_payload(self) -> Dict[str, Any]:
        """[LEGACY_COMPAT:P3-A] Build manual OTLP export payload.
        
        Remove after native OTEL mode is validated + 1 minor release.
        In native mode, OTLPSpanExporter handles serialization automatically.
        
        IMPORTANT: agent_id and customer_id must be sent as resource attributes
        for the otel-collector to extract them correctly.
        
        v6.8.6 FIX: Check signal config for agent_id set by @agentbill_traced decorator,
        falling back to AgentBill config. This fixes the issue where agent_id passed
        to the decorator wasn't being propagated to wrap_openai() spans.
        """
        from .signals import get_signal_config
        
        resource_attributes = [
            {"key": "service.name", "value": {"stringValue": "agentbill-python-sdk"}},
            {"key": "service.version", "value": {"stringValue": _SDK_VERSION}}
        ]
        
        # Get signal config which may have agent_id from @agentbill_traced decorator
        signal_config = get_signal_config() or {}
        
        # Add customer_id: prefer signal config, then check external_customer_id as alias, fallback to tracer config
        customer_id = (
            signal_config.get("customer_id") or 
            signal_config.get("external_customer_id") or 
            self.customer_id or 
            self.config.get("external_customer_id")
        )
        if customer_id:
            resource_attributes.append(
                {"key": "customer.id", "value": {"stringValue": customer_id}}
            )
        
        # Add external_customer_id as separate attribute for lookup (always send if configured)
        external_customer_id = (
            signal_config.get("external_customer_id") or 
            self.config.get("external_customer_id")
        )
        if external_customer_id:
            resource_attributes.append(
                {"key": "customer.external_id", "value": {"stringValue": external_customer_id}}
            )
        
        # Add agent_id: prefer signal config (from decorator), fallback to tracer config
        # This is the CRITICAL fix - decorator's agent_id now propagates to wrap_openai spans
        agent_id = signal_config.get("agent_id") or self.config.get("agent_id")
        if agent_id:
            resource_attributes.append(
                {"key": "agent.id", "value": {"stringValue": agent_id}}
            )
        
        return {
            "resourceSpans": [{
                "resource": {
                    "attributes": resource_attributes
                },
                "scopeSpans": [{
                    "scope": {"name": "agentbill", "version": _SDK_VERSION},
                    "spans": [self._span_to_otlp(span) for span in self.spans]
                }]
            }]
        }
    
    def _span_to_otlp(self, span: Span) -> Dict[str, Any]:
        """[LEGACY_COMPAT:P3-A] Convert span to OTLP format — remove after native OTEL validated."""
        otlp_span = {
            "traceId": span.trace_id,
            "spanId": span.span_id,
            "name": span.name,
            "kind": 1,  # CLIENT
            "startTimeUnixNano": str(span.start_time),
            "endTimeUnixNano": str(span.end_time or time.time_ns()),
            "attributes": [
                {"key": k, "value": self._value_to_otlp(v)}
                for k, v in span.attributes.items()
            ],
            "status": span.status
        }
        
        # Add parent span ID if this is a child span
        if span.parent_span_id:
            otlp_span["parentSpanId"] = span.parent_span_id
        
        return otlp_span
    
    def _value_to_otlp(self, value: Any) -> Dict[str, Any]:
        """[LEGACY_COMPAT:P3-A] Convert value to OTLP format — remove after native OTEL validated."""
        if isinstance(value, str):
            return {"stringValue": value}
        elif isinstance(value, (int, float)):
            return {"intValue": int(value)}
        elif isinstance(value, bool):
            return {"boolValue": value}
        else:
            return {"stringValue": str(value)}