"""
Tracer for creating and managing spans.

The tracer is the main interface for creating spans. It handles:
- Span creation with proper parent linking
- Context propagation
- Span lifecycle management
- Export on span end
"""

from __future__ import annotations

import random
import threading
from contextlib import contextmanager
from contextvars import ContextVar
from typing import Any, Dict, Iterator, List, Optional, TYPE_CHECKING

from risicare_core import (
    Span,
    SpanKind,
    SpanStatus,
    SemanticPhase,
    LLMAttributes,
    generate_span_id,
    generate_trace_id,
    get_current_span,
    get_current_session,
    get_current_agent,
    get_current_phase,
    set_current_span,
    reset_current_span,
    register_span,
    unregister_span,
    utc_now,
)

if TYPE_CHECKING:
    from risicare.client import RisicareClient


# =============================================================================
# No-Op Span (for when tracing is disabled)
# =============================================================================

class _NoOpSpan:
    """
    A no-op span that does nothing.

    Used when tracing is disabled to avoid overhead.
    All methods are no-ops that return immediately.
    """

    __slots__ = ("trace_id", "span_id", "name", "kind", "status", "parent_span_id",
                 "session_id", "agent_id", "semantic_phase", "llm", "attributes",
                 "events", "links", "exceptions", "start_time", "end_time",
                 "status_message")

    def __init__(self) -> None:
        self.trace_id = ""
        self.span_id = ""
        self.name = ""
        self.kind = SpanKind.INTERNAL
        self.status = SpanStatus.UNSET
        self.parent_span_id: Optional[str] = None
        self.session_id: Optional[str] = None
        self.agent_id: Optional[str] = None
        self.semantic_phase: Optional[SemanticPhase] = None
        self.llm: Optional[LLMAttributes] = None
        self.attributes: Dict[str, Any] = {}
        self.events: List[Any] = []
        self.links: List[Any] = []
        self.exceptions: List[Any] = []
        self.start_time = utc_now()
        self.end_time = None
        self.status_message: Optional[str] = None

    @property
    def duration_ms(self) -> Optional[float]:
        return None

    @property
    def is_root(self) -> bool:
        return True

    @property
    def is_error(self) -> bool:
        return False

    @property
    def is_complete(self) -> bool:
        return True

    @property
    def has_exceptions(self) -> bool:
        return False

    def set_attribute(self, key: str, value: Any) -> None:
        pass

    def set_attributes(self, attributes: Dict[str, Any]) -> None:
        pass

    def get_attribute(self, key: str, default: Any = None) -> Any:
        return default

    def add_event(self, name: str, attributes: Optional[Dict[str, Any]] = None,
                  timestamp: Any = None) -> None:
        pass

    def add_link(self, trace_id: str, span_id: str,
                 attributes: Optional[Dict[str, Any]] = None) -> None:
        pass

    def record_exception(self, exception: BaseException, escaped: bool = True) -> None:
        pass

    def set_status(self, status: SpanStatus, message: Optional[str] = None) -> None:
        pass

    def mark_ok(self, message: Optional[str] = None) -> None:
        pass

    def mark_error(self, message: Optional[str] = None) -> None:
        pass

    def end(self, end_time: Any = None) -> None:
        pass

    def to_dict(self) -> Dict[str, Any]:
        return {}


# Singleton no-op span to avoid allocations
_NOOP_SPAN = _NoOpSpan()

# Trace-level sampling decision — propagated from root to all children
_trace_sampled: ContextVar[Optional[bool]] = ContextVar("_trace_sampled", default=None)


# =============================================================================
# Global Tracer
# =============================================================================

_tracer: Optional["Tracer"] = None
_tracer_lock = threading.Lock()


def get_tracer() -> "Tracer":
    """
    Get the global tracer instance.

    Returns a no-op tracer if SDK not initialized.
    """
    global _tracer
    with _tracer_lock:
        if _tracer is None:
            _tracer = Tracer()
        return _tracer


def _set_tracer(tracer: "Tracer") -> None:
    """Set the global tracer (internal use)."""
    global _tracer
    with _tracer_lock:
        _tracer = tracer


def _reset_for_testing() -> None:
    """
    Reset global tracer state for testing.

    This is intended for test fixtures only. Do not use in production.
    """
    global _tracer
    with _tracer_lock:
        _tracer = None


# =============================================================================
# Tracer Class
# =============================================================================

class Tracer:
    """
    Tracer for creating and managing spans.

    The tracer provides methods to create spans that represent
    operations in your agent system. Spans are automatically
    linked to form traces.

    Example:
        tracer = get_tracer()

        with tracer.start_span("research") as span:
            span.set_attribute("query", query)
            result = llm.invoke(query)
            span.set_attribute("result_length", len(result))
    """

    def __init__(self, client: Optional["RisicareClient"] = None):
        """
        Initialize the tracer.

        Args:
            client: Optional client for span export.
        """
        self._client = client

    @property
    def is_enabled(self) -> bool:
        """Check if tracing is enabled."""
        if self._client:
            return self._client.is_enabled
        # Check for global client
        from risicare.client import get_client
        client = get_client()
        return client.is_enabled if client else False

    def create_span(
        self,
        name: str,
        kind: SpanKind = SpanKind.INTERNAL,
        attributes: Optional[Dict[str, Any]] = None,
    ) -> Span:
        """
        Create a span with context linking but manual lifecycle.

        Unlike start_span(), this does NOT use a context manager and does NOT
        set the span as current context. The caller is responsible for:
        1. Calling span.end() when the operation completes
        2. Calling tracer.export_span(span) after ending

        Returns _NOOP_SPAN if tracing is disabled or the trace is not sampled.

        Primary use case: streaming responses where the span must outlive its
        creation scope (the generator pattern).

        Args:
            name: Name describing the operation.
            kind: Kind of span (INTERNAL, AGENT, LLM_CALL, etc.).
            attributes: Initial span attributes.

        Returns:
            The new span, or _NOOP_SPAN if disabled/not sampled.
        """
        if not self.is_enabled:
            return _NOOP_SPAN  # type: ignore

        # Check sampling
        parent = get_current_span()
        if parent is None:
            # Root span: make sampling decision
            trace_id = generate_trace_id()
            sampled = self._should_sample(trace_id)
            if not sampled:
                return _NOOP_SPAN  # type: ignore
        else:
            trace_id = None  # _create_span will inherit from parent
            sampled = _trace_sampled.get()
            if sampled is False:
                return _NOOP_SPAN  # type: ignore

        return self._create_span(name, kind, attributes, trace_id=trace_id)

    def export_span(self, span: "Span") -> None:
        """
        Export a manually-managed span.

        Call this after span.end() for spans created with create_span().
        No-op if the span has no span_id (i.e. is a _NoOpSpan).
        """
        if not span.span_id:
            return
        self._export_span(span)

    def _should_sample(self, trace_id: Optional[str] = None) -> bool:
        """Check if this trace should be sampled.

        Uses deterministic hashing of trace_id for consistent decisions
        across distributed systems. Falls back to random if no trace_id.
        """
        client = self._client
        if client is None:
            from risicare.client import get_client
            client = get_client()

        if client is None:
            return True  # No client, sample everything

        sample_rate = client.config.sample_rate
        if sample_rate >= 1.0:
            return True
        if sample_rate <= 0.0:
            return False

        # Deterministic: hash trace_id for consistent decision
        if trace_id:
            hash_val = int(trace_id[:8], 16) / 0xFFFFFFFF
            return hash_val < sample_rate

        return random.random() < sample_rate

    @contextmanager
    def start_span(
        self,
        name: str,
        kind: SpanKind = SpanKind.INTERNAL,
        attributes: Optional[Dict[str, Any]] = None,
        links: Optional[list] = None,
        register: bool = False,
    ) -> Iterator[Span]:
        """
        Start a new span as a context manager.

        The span automatically:
        - Links to parent span if one exists
        - Inherits trace ID from parent or generates new one
        - Sets session_id and agent_id from context
        - Ends when context exits
        - Exports on end

        Args:
            name: Name describing the operation.
            kind: Kind of span (INTERNAL, AGENT, LLM_CALL, etc.).
            attributes: Initial span attributes.
            links: Links to related spans.
            register: Whether to register span for ID-based lookup.

        Yields:
            The active span (or no-op span if disabled/not sampled).

        Example:
            with tracer.start_span("llm_call", kind=SpanKind.LLM_CALL) as span:
                span.set_attribute("model", "gpt-4")
                response = llm.invoke(prompt)
                span.llm = LLMAttributes(
                    model="gpt-4",
                    prompt_tokens=100,
                    completion_tokens=50,
                )
        """
        # Check if enabled (fast path for disabled)
        if not self.is_enabled:
            yield _NOOP_SPAN  # type: ignore
            return

        # Head-based trace-coherent sampling
        parent = get_current_span()
        if parent is None:
            # Root span: make sampling decision and store it
            trace_id = generate_trace_id()
            sampled = self._should_sample(trace_id)
            sampling_token = _trace_sampled.set(sampled)
            if not sampled:
                yield _NOOP_SPAN  # type: ignore
                _trace_sampled.reset(sampling_token)
                return
        else:
            trace_id = None  # _create_span will inherit from parent
            sampling_token = None
            # Child span: inherit parent's sampling decision
            sampled = _trace_sampled.get()
            if sampled is False:
                yield _NOOP_SPAN  # type: ignore
                return

        # Create span with context (pass trace_id to avoid double generation)
        span = self._create_span(name, kind, attributes, trace_id=trace_id)

        # Add links
        if links:
            for link in links:
                span.add_link(link.trace_id, link.span_id, link.attributes)

        # Register if requested (for async generator workaround)
        if register:
            register_span(span)

        # Set as current span
        token = set_current_span(span)

        try:
            yield span
            # Mark OK if no exception and status not set
            if span.status == SpanStatus.UNSET:
                span.mark_ok()
        except Exception as e:
            # Record exception
            span.record_exception(e, escaped=True)
            raise
        finally:
            # End span
            span.end()

            # Reset context
            reset_current_span(token)

            # Reset sampling context for root spans
            if sampling_token is not None:
                _trace_sampled.reset(sampling_token)

            # Unregister if registered
            if register:
                unregister_span(span.span_id)

            # Export
            self._export_span(span)

    def start_span_no_context(
        self,
        name: str,
        kind: SpanKind = SpanKind.INTERNAL,
        attributes: Optional[Dict[str, Any]] = None,
        parent_span_id: Optional[str] = None,
        trace_id: Optional[str] = None,
    ) -> Span:
        """
        Start a span without setting it as current context.

        Use this when you need manual span lifecycle management.
        You must call span.end() and export manually.

        Args:
            name: Span name.
            kind: Span kind.
            attributes: Initial attributes.
            parent_span_id: Optional parent span ID.
            trace_id: Optional trace ID.

        Returns:
            The new span (not set as current).
        """
        span = Span(
            trace_id=trace_id or generate_trace_id(),
            span_id=generate_span_id(),
            name=name,
            kind=kind,
            parent_span_id=parent_span_id,
        )

        if attributes:
            span.set_attributes(attributes)

        return span

    def _create_span(
        self,
        name: str,
        kind: SpanKind,
        attributes: Optional[Dict[str, Any]],
        trace_id: Optional[str] = None,
    ) -> Span:
        """Create a span with proper context linking.

        Args:
            name: Span name.
            kind: Span kind.
            attributes: Initial attributes.
            trace_id: Pre-generated trace ID for root spans. If provided and
                this is a root span, this ID is used instead of generating a new
                one. This ensures the same trace_id used for the sampling
                decision is assigned to the span.
        """
        # Get current context
        parent = get_current_span()
        session = get_current_session()
        agent = get_current_agent()
        phase = get_current_phase()

        # Determine trace ID and parent span ID
        if parent:
            actual_trace_id = parent.trace_id
            parent_span_id = parent.span_id
        else:
            actual_trace_id = trace_id or generate_trace_id()
            parent_span_id = None

        # Create span
        span = Span(
            trace_id=actual_trace_id,
            span_id=generate_span_id(),
            name=name,
            kind=kind,
            parent_span_id=parent_span_id,
            session_id=session.session_id if session else None,
            agent_id=agent.agent_id if agent else None,
            agent_name=agent.agent_name if agent else None,
            agent_type=agent.agent_type if agent else None,
            semantic_phase=phase,
        )

        # Set initial attributes
        if attributes:
            span.set_attributes(attributes)

        # Add agent metadata as attributes (for backward compat with worker extraction)
        if agent:
            if agent.agent_role:
                span.set_attribute("agent.role", agent.agent_role)
            if agent.agent_type:
                span.set_attribute("agent.type", agent.agent_type)
            if agent.agent_name:
                span.set_attribute("agent.name", agent.agent_name)
            if agent.parent_agent_id:
                span.set_attribute("parent_agent_id", agent.parent_agent_id)

        # Propagate user_id from session context into span attributes
        # so workers can extract it for session aggregation
        if session and session.user_id:
            span.set_attribute("session.user_id", session.user_id)

        return span

    def _export_span(self, span: Span) -> None:
        """Export span via client."""
        client = self._client
        if client is None:
            from risicare.client import get_client
            client = get_client()

        if client:
            client.export_span(span)


# =============================================================================
# Convenience Functions
# =============================================================================

def start_span(
    name: str,
    kind: SpanKind = SpanKind.INTERNAL,
    attributes: Optional[Dict[str, Any]] = None,
    register: bool = False,
) -> Iterator[Span]:
    """
    Start a new span using the global tracer.

    Convenience function equivalent to get_tracer().start_span(...).

    Args:
        name: Span name.
        kind: Span kind.
        attributes: Initial attributes.
        register: Whether to register for ID-based lookup.

    Returns:
        Context manager yielding the span.
    """
    return get_tracer().start_span(
        name=name,
        kind=kind,
        attributes=attributes,
        register=register,
    )
