"""
Core tracing primitives for Cascade SDK
"""
import os
import time
import logging
import threading
import contextvars
from typing import Optional, Dict, Any, List
from contextlib import contextmanager
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.sdk.resources import Resource
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.trace import Tracer, Status, StatusCode

logger = logging.getLogger(__name__)

# Global tracer instance
_tracer: Optional[Tracer] = None
_tracer_provider: Optional[TracerProvider] = None

# Global configuration for trace start events
_cascade_endpoint: Optional[str] = None
_cascade_api_key: Optional[str] = None
_cascade_project: Optional[str] = None

# Context variable for trace metadata
_trace_metadata: contextvars.ContextVar[Dict[str, Any]] = contextvars.ContextVar(
    "trace_metadata", default={}
)

# Context variable for current agent name
_current_agent: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
    "current_agent", default=None
)

# Context variable for current session id (multi-turn grouping)
_current_session_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
    "current_session_id", default=None
)


def _send_trace_start(
    trace_id: str,
    root_span_id: str,
    name: str,
    start_time: int,
    metadata: Optional[Dict[str, Any]] = None,
    session_id: Optional[str] = None,
) -> None:
    """
    Send trace start event to backend (non-blocking).
    
    PHASE 3, STEP 3.1 - Real-time trace updates
    
    This function sends an HTTP POST to /api/traces/start when a root span begins.
    It runs in a background thread to avoid blocking trace execution.
    
    Args:
        trace_id: Unique trace identifier (from OTEL span context)
        root_span_id: ID of the root span
        name: Trace/agent name
        start_time: Timestamp in nanoseconds
        metadata: Optional metadata dictionary
    """
    if not _cascade_endpoint or not _cascade_api_key or not _cascade_project:
        # Trace start events disabled (endpoint/API key/project not configured)
        return

    # NOTE: This used to run in a background thread (non-blocking). In practice, that can race
    # with the OTLP export of the root span: if the server receives the "start" placeholder
    # *after* it already stored the completed root span, a naive upsert can wipe end_time back
    # to NULL, leaving the trace "LIVE" forever in the dashboard polling UI.
    #
    # Sending synchronously with a short timeout makes the ordering deterministic:
    # placeholder root is created first, then the OTLP-exported span updates it.
    try:
        import requests

        base_url = _cascade_endpoint.replace("/v1/traces", "").rstrip("/")
        api_url = f"{base_url}/api/traces/start"

        payload = {
            "trace_id": trace_id,
            "root_span_id": root_span_id,
            "name": name,
            "project": _cascade_project,
            "start_time": start_time,
            "metadata": metadata,
            "session_id": session_id,
        }

        response = requests.post(
            api_url,
            json=payload,
            headers={"X-API-Key": _cascade_api_key},
            timeout=2,  # keep trace_run() snappy; failure is non-fatal
        )

        if response.status_code == 200:
            logger.debug(
                f"Trace start event sent successfully: {trace_id} (name='{name}', project='{_cascade_project}')"
            )
        else:
            logger.warning(
                f"Trace start event failed: {trace_id} "
                f"(status={response.status_code}, response={response.text})"
            )

    except Exception as e:
        # Don't crash if trace start fails - span export will create placeholder
        logger.warning(f"Failed to send trace start event for {trace_id}: {e}")


def init_tracing(
    project: Optional[str] = None,
    endpoint: Optional[str] = None,
    api_key: Optional[str] = None,
    metadata: Optional[Dict[str, str]] = None,
    environment: Optional[str] = None,
    version: Optional[str] = None,
    user_id: Optional[str] = None,
    json_export_path: Optional[str] = None,
    evals: Optional[List[str]] = None,
    session_evals: Optional[List[str]] = None,
) -> None:
    """
    Initialize OpenTelemetry tracing for Cascade SDK.
    
    **Minimal Usage (Recommended):**
    ```python
    from cascade import init_tracing
    
    # Option 1: Just use environment variable (simplest)
    # Set CASCADE_API_KEY=csk_live_xxx in your .env
    init_tracing()
    
    # Option 2: Specify project name to organize traces
    init_tracing(project="customer_support_chatbot")
    
    # Option 3: Pass API key directly
    init_tracing(api_key="csk_live_xxx...")
    
    # Option 4: Enable auto-evaluation (zero extra code!)
    init_tracing(project="my_agent", evals=["helpfulness", "hallucination"])
    ```
    
    Args:
        project: Project name to organize traces (optional - defaults to auto-generated name from API key)
        endpoint: OTLP endpoint URL (optional - defaults to https://api.runcascade.com/v1/traces)
        api_key: API key for authentication (optional - reads from CASCADE_API_KEY env var if not provided)
        metadata: Optional metadata dictionary (e.g., {"user_id": "123", "version": "v2"})
        environment: Optional environment name (e.g., 'production', 'staging')
        version: Optional application version
        user_id: Optional user identifier
        json_export_path: Optional path to export traces locally for debugging
        evals: Optional list of scorer names to auto-evaluate every trace.
               Accepts built-in short names (e.g. ``"helpfulness"``, ``"hallucination"``),
               full keys (``"trace_helpfulness"``), or custom scorer names.
               When provided, a scheduled evaluation task is created automatically.
        session_evals: Optional list of session-scoped scorer names.  These are
               evaluated when ``end_session()`` is called after a multi-turn
               conversation ends.
    
    **Environment Variables:**
    - CASCADE_API_KEY: Your organization's API key (get from dashboard settings)
    - CASCADE_ENDPOINT: Optional custom endpoint (defaults to production)
    """
    global _tracer, _tracer_provider, _cascade_endpoint, _cascade_api_key, _cascade_project

    try:
        # Use environment variable or parameter for endpoint
        endpoint = endpoint or os.getenv("CASCADE_ENDPOINT", "https://api.runcascade.com/v1/traces")
        
        # Use environment variable or parameter for API key
        api_key = api_key or os.getenv("CASCADE_API_KEY")
        
        # If project not provided, derive from API key or use default
        if not project:
            if api_key:
                project = api_key.split('_')[0:3]
                project = '_'.join(project)[:20]
            else:
                project = "default"
        
        # Store in global variables for trace start events (Phase 3)
        _cascade_endpoint = endpoint
        _cascade_api_key = api_key
        _cascade_project = project
        
        # Detect Modal environment
        is_modal = os.getenv("MODAL_ENVIRONMENT") is not None
        if is_modal and endpoint == "http://localhost:8000/v1/traces":
            import warnings
            warnings.warn(
                "Running in Modal environment but no CASCADE_ENDPOINT configured. "
                "Traces will not be sent. Set CASCADE_ENDPOINT to your tunnel URL.",
                UserWarning
            )
        
        # Build resource attributes
        resource_attributes = {
            "service.name": project,
            "cascade.project": project,
            "cascade.environment": environment or "local",
        }
        
        if environment:
            resource_attributes["deployment.environment"] = environment
        if version:
            resource_attributes["service.version"] = version
        if user_id:
            resource_attributes["cascade.user_id"] = user_id
        
        if metadata:
            for key, value in metadata.items():
                resource_attributes[f"cascade.{key}"] = str(value)
        
        resource = Resource.create(resource_attributes)
        
        # Build headers for API key authentication
        headers = None
        if api_key:
            headers = {"authorization": f"Bearer {api_key}"}
        
        otlp_exporter = OTLPSpanExporter(
            endpoint=endpoint,
            headers=headers,
        )
        
        span_processor = BatchSpanProcessor(
            otlp_exporter,
            max_queue_size=8192,
            max_export_batch_size=512,
        )

        # If another SDK (e.g. Sentry, Braintrust) already owns the global
        # TracerProvider, attach Cascade's exporter to it instead of replacing
        # it. This prevents "signal eating" where one SDK silently drops spans
        # meant for the other. If no real provider exists yet, create one.
        existing_provider = trace.get_tracer_provider()
        if isinstance(existing_provider, TracerProvider):
            existing_provider.add_span_processor(span_processor)
            _tracer_provider = existing_provider
            logger.debug("Cascade: attaching to existing TracerProvider")
        else:
            _tracer_provider = TracerProvider(resource=resource)
            _tracer_provider.add_span_processor(span_processor)
            trace.set_tracer_provider(_tracer_provider)

        # Add JSON file exporter if requested (optional, parallel export)
        if json_export_path:
            try:
                from cascade.exporters.json_exporter import JSONFileExporter
                json_exporter = JSONFileExporter(json_export_path)
                json_span_processor = BatchSpanProcessor(json_exporter)
                _tracer_provider.add_span_processor(json_span_processor)
                logger.info(f"JSON file export enabled: {json_export_path}")
            except Exception as e:
                logger.error(f"Failed to initialize JSON file exporter: {e}")
                logger.warning("Continuing without JSON file export")
        
        _tracer = trace.get_tracer(__name__)

    except Exception as e:
        logger.warning(f"Cascade tracing initialization failed: {e}. Tracing will be disabled.")
        _tracer = None
        return

    # ── Tier 1: auto-evaluation setup ────────────────────────────────────
    if evals:
        from cascade.evaluations import _setup_scheduled_evals
        thread = threading.Thread(
            target=_setup_scheduled_evals,
            args=(project, evals),
            daemon=True,
        )
        thread.start()

    # ── Session-level auto-evaluation setup ──────────────────────────────
    if session_evals:
        from cascade.evaluations import _setup_session_evals
        thread = threading.Thread(
            target=_setup_session_evals,
            args=(project, session_evals),
            daemon=True,
        )
        thread.start()


def get_tracer() -> Optional[Tracer]:
    """Get the global tracer instance."""
    return _tracer


def get_current_agent() -> Optional[str]:
    """Get the currently active agent name from context."""
    return _current_agent.get()


def set_session_id(session_id: Optional[str]) -> None:
    """
    Set the current session ID in context.

    All subsequent ``trace_run()`` calls in this context inherit the session ID
    unless explicitly overridden in ``metadata``.
    """
    _current_session_id.set(session_id)


def get_session_id() -> Optional[str]:
    """Get the current session ID from context."""
    return _current_session_id.get()


@contextmanager
def trace_session(session_id: str):
    """
    Context manager to scope multiple traces to one session.

    Example:
        with trace_session("chat-123"):
            with trace_run("Turn", metadata={"turn": 1}):
                ...
            with trace_run("Turn", metadata={"turn": 2}):
                ...
    """
    token = _current_session_id.set(session_id)
    try:
        yield
    finally:
        _current_session_id.reset(token)


@contextmanager
def trace_run(
    name: str,
    metadata: Optional[Dict[str, Any]] = None,
):
    """
    Context manager for tracing agent execution.
    
    Creates a root span for the agent run and maintains trace context.
    
    Args:
        name: Name of the agent/function being traced
        metadata: Additional metadata to attach to the span (e.g., session_id, user_id, request_id)
    
    Example:
        ```python
        with trace_run("MyAgent", metadata={"task_id": "123"}):
            # Your agent code here
            pass
        ```
    """
    if _tracer is None:
        logger.debug("Cascade tracing not initialized — trace_run() is a no-op")
        yield None
        return
    
    metadata = dict(metadata or {})

    # Inherit session id from context for multi-turn conversations
    # unless caller explicitly sets metadata["session_id"].
    if "session_id" not in metadata or not metadata.get("session_id"):
        session_id = _current_session_id.get()
        if session_id:
            metadata["session_id"] = session_id
    
    # Store metadata in context
    _trace_metadata.set(metadata)
    
    # Start root span - start_as_current_span returns a context manager
    with _tracer.start_as_current_span(name, kind=trace.SpanKind.SERVER) as span_context:
        # Get the actual span from the current context
        span = trace.get_current_span()
        
        # Determine root status before try/except so it's always
        # available in the finally block.
        parent = getattr(span, 'parent', None)
        is_root = parent is None or not getattr(parent, 'is_valid', lambda: False)()
        
        try:
            # Add metadata as span attributes
            for key, value in metadata.items():
                if value is not None:
                    span.set_attribute(f"cascade.{key}", str(value))
            
            span.set_attribute("cascade.span_type", "function")
            
            span_ctx = span.get_span_context()
            trace_id = format(span_ctx.trace_id, '032x')
            span_id = format(span_ctx.span_id, '016x')
            start_time = time.time_ns()
            if is_root:
                _send_trace_start(
                    trace_id=trace_id,
                    root_span_id=span_id,
                    name=name,
                    start_time=start_time,
                    metadata=metadata,
                    session_id=metadata.get("session_id"),
                )
            
            yield span
            
            span.set_status(Status(StatusCode.OK))
            
        except Exception as e:
            span.set_status(Status(StatusCode.ERROR, str(e)))
            span.record_exception(e)
            raise
        finally:
            # Force flush after root span ends so the root span is exported
            # before the process can terminate.  Critical for serverless
            # (Lambda, Cloud Functions, etc.) where the process is killed
            # immediately after the handler returns.
            if is_root and _tracer_provider is not None:
                try:
                    _tracer_provider.force_flush(timeout_millis=5000)
                except Exception:
                    pass


@contextmanager
def trace_agent(
    agent_name: str,
    metadata: Optional[Dict[str, Any]] = None,
):
    """
    Context manager for tracing agent execution.
    
    Creates a span for the agent and sets it as the current agent context.
    All tools and LLM calls within this context will be tagged with the agent name.
    
    Args:
        agent_name: Name of the agent (e.g., "ExplorerAgent", "EvaluatorAgent")
        metadata: Additional metadata to attach to the agent span (e.g., {"loop": 1})
    
    Example:
        ```python
        with trace_agent("EvaluatorAgent", metadata={"loop": 1}):
            # Run evaluator logic
            result = evaluator.run()
            # All tool calls and LLM calls here will be tagged with agent="EvaluatorAgent"
        ```
    """
    if _tracer is None:
        logger.debug("Cascade tracing not initialized — trace_agent() is a no-op")
        yield None
        return
    
    metadata = metadata or {}
    
    # Set current agent in context
    token = _current_agent.set(agent_name)
    
    try:
        with _tracer.start_as_current_span(
            agent_name, 
            kind=trace.SpanKind.INTERNAL
        ) as span_context:
            span = trace.get_current_span()
            
            try:
                span.set_attribute("cascade.span_type", "agent")
                span.set_attribute("cascade.agent_name", agent_name)
                
                for key, value in metadata.items():
                    if value is not None:
                        span.set_attribute(f"agent.{key}", str(value))
                
                yield span
                
                span.set_status(Status(StatusCode.OK))
                
            except Exception as e:
                span.set_status(Status(StatusCode.ERROR, str(e)))
                span.record_exception(e)
                raise
    finally:
        _current_agent.reset(token)

