"""
Span types for distributed tracing.

Spans represent operations with a start time, duration, and metadata.
They form the building blocks of traces.
"""

from __future__ import annotations

import traceback
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

from risicare_core.types.base import (
    AgentID,
    Metadata,
    SemanticPhase,
    SpanID,
    SpanKind,
    SpanStatus,
    Timestamp,
    TraceID,
    generate_span_id,
    utc_now,
    validate_trace_id,
    validate_span_id,
)


# =============================================================================
# Span Events
# =============================================================================

@dataclass(frozen=True)
class SpanEvent:
    """
    An event that occurred during a span's lifetime.

    Events are timestamped annotations that provide additional
    context about what happened during the span.

    Attributes:
        name: Human-readable event name.
        timestamp: When the event occurred.
        attributes: Additional event metadata (immutable after creation).
    """
    name: str
    timestamp: Timestamp = field(default_factory=utc_now)
    attributes: Metadata = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "name": self.name,
            "timestamp": self.timestamp.isoformat(),
            "attributes": dict(self.attributes),  # Copy to ensure immutability
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "SpanEvent":
        """Create SpanEvent from dictionary."""
        from datetime import datetime
        timestamp = data.get("timestamp")
        if isinstance(timestamp, str):
            timestamp = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
        elif timestamp is None:
            timestamp = utc_now()
        return cls(
            name=data["name"],
            timestamp=timestamp,
            attributes=dict(data.get("attributes", {})),
        )


# =============================================================================
# Span Links
# =============================================================================

@dataclass(frozen=True)
class SpanLink:
    """
    A link to another span, possibly in a different trace.

    Links are used to connect causally related spans that are
    not in a parent-child relationship.

    Attributes:
        trace_id: The linked span's trace ID.
        span_id: The linked span's ID.
        attributes: Additional link metadata (immutable after creation).
    """
    trace_id: TraceID
    span_id: SpanID
    attributes: Metadata = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization (camelCase)."""
        return {
            "traceId": self.trace_id,
            "spanId": self.span_id,
            "attributes": dict(self.attributes),
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "SpanLink":
        """Create SpanLink from dictionary (supports both camelCase and snake_case)."""
        return cls(
            trace_id=data.get("traceId") or data.get("trace_id", ""),
            span_id=data.get("spanId") or data.get("span_id", ""),
            attributes=dict(data.get("attributes", {})),
        )


# =============================================================================
# Exception Info
# =============================================================================

@dataclass(frozen=True)
class ExceptionInfo:
    """
    Information about an exception that occurred during a span.

    Attributes:
        type: Exception class name.
        message: Exception message.
        stacktrace: Full stack trace as string.
        timestamp: When the exception occurred.
        escaped: Whether the exception escaped the span.
    """
    type: str
    message: str
    stacktrace: str
    timestamp: Timestamp = field(default_factory=utc_now)
    escaped: bool = True

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "type": self.type,
            "message": self.message,
            "stacktrace": self.stacktrace,
            "timestamp": self.timestamp.isoformat(),
            "escaped": self.escaped,
        }

    @classmethod
    def from_exception(
        cls,
        exception: BaseException,
        escaped: bool = True,
    ) -> "ExceptionInfo":
        """Create ExceptionInfo from an exception."""
        return cls(
            type=type(exception).__name__,
            message=str(exception),
            stacktrace=traceback.format_exc(),
            escaped=escaped,
        )

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ExceptionInfo":
        """Create ExceptionInfo from dictionary."""
        from datetime import datetime
        timestamp = data.get("timestamp")
        if isinstance(timestamp, str):
            timestamp = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
        elif timestamp is None:
            timestamp = utc_now()
        return cls(
            type=data["type"],
            message=data["message"],
            stacktrace=data["stacktrace"],
            timestamp=timestamp,
            escaped=data.get("escaped", True),
        )


# =============================================================================
# LLM-Specific Attributes
# =============================================================================

@dataclass
class LLMAttributes:
    """
    LLM-specific attributes for a span.

    These attributes capture details about LLM API calls.
    """
    model: Optional[str] = None
    provider: Optional[str] = None
    prompt_tokens: Optional[int] = None
    completion_tokens: Optional[int] = None
    total_tokens: Optional[int] = None
    cost_usd: Optional[float] = None
    temperature: Optional[float] = None
    max_tokens: Optional[int] = None
    stop_reason: Optional[str] = None
    confidence: Optional[float] = None  # Extracted from logprobs

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary, excluding None values."""
        return {k: v for k, v in {
            "model": self.model,
            "provider": self.provider,
            "prompt_tokens": self.prompt_tokens,
            "completion_tokens": self.completion_tokens,
            "total_tokens": self.total_tokens,
            "cost_usd": self.cost_usd,
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
            "stop_reason": self.stop_reason,
            "confidence": self.confidence,
        }.items() if v is not None}

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "LLMAttributes":
        """Create LLMAttributes from dictionary."""
        return cls(
            model=data.get("model"),
            provider=data.get("provider"),
            prompt_tokens=data.get("prompt_tokens"),
            completion_tokens=data.get("completion_tokens"),
            total_tokens=data.get("total_tokens"),
            cost_usd=data.get("cost_usd"),
            temperature=data.get("temperature"),
            max_tokens=data.get("max_tokens"),
            stop_reason=data.get("stop_reason"),
            confidence=data.get("confidence"),
        )


# =============================================================================
# Base Span
# =============================================================================

@dataclass
class Span:
    """
    A span represents a single operation within a trace.

    This is the core unit of observability in Risicare. Spans can
    be nested to form a tree structure representing the execution flow.

    Attributes:
        trace_id: The trace this span belongs to.
        span_id: Unique identifier for this span.
        parent_span_id: Parent span's ID, or None for root spans.
        name: Human-readable name describing the operation.
        kind: The kind of span (see SpanKind).
        start_time: When the span started.
        end_time: When the span ended, or None if still active.
        status: The span's status (see SpanStatus).
        status_message: Optional message explaining the status.
        attributes: Key-value pairs with span metadata.
        events: List of events that occurred during the span.
        links: Links to related spans.
        exceptions: List of exceptions that occurred.
        session_id: Session this span belongs to.
        agent_id: Agent that created this span.
        semantic_phase: The semantic phase (THINK/DECIDE/ACT/etc).
        llm: LLM-specific attributes if this is an LLM call.
    """
    # Required fields
    trace_id: TraceID
    span_id: SpanID
    name: str

    # Optional fields with defaults
    parent_span_id: Optional[SpanID] = None
    kind: SpanKind = SpanKind.INTERNAL
    start_time: Timestamp = field(default_factory=utc_now)
    end_time: Optional[Timestamp] = None
    status: SpanStatus = SpanStatus.UNSET
    status_message: Optional[str] = None

    # Collections
    attributes: Dict[str, Any] = field(default_factory=dict)
    events: List[SpanEvent] = field(default_factory=list)
    links: List[SpanLink] = field(default_factory=list)
    exceptions: List[ExceptionInfo] = field(default_factory=list)

    # Risicare-specific fields
    session_id: Optional[str] = None
    agent_id: Optional[AgentID] = None
    agent_name: Optional[str] = None
    agent_type: Optional[str] = None
    semantic_phase: Optional[SemanticPhase] = None
    llm: Optional[LLMAttributes] = None

    def __post_init__(self) -> None:
        """Validate span IDs after initialization."""
        if not validate_trace_id(self.trace_id):
            raise ValueError(
                f"Invalid trace_id: {self.trace_id!r}. "
                f"Must be a 32-character lowercase hex string."
            )
        if not validate_span_id(self.span_id):
            raise ValueError(
                f"Invalid span_id: {self.span_id!r}. "
                f"Must be a 16-character lowercase hex string."
            )
        if self.parent_span_id is not None and not validate_span_id(self.parent_span_id):
            raise ValueError(
                f"Invalid parent_span_id: {self.parent_span_id!r}. "
                f"Must be a 16-character lowercase hex string."
            )

    # ==========================================================================
    # Properties
    # ==========================================================================

    @property
    def duration_ms(self) -> Optional[float]:
        """Calculate span duration in milliseconds."""
        if self.end_time is None:
            return None
        delta = self.end_time - self.start_time
        return delta.total_seconds() * 1000

    @property
    def is_root(self) -> bool:
        """Check if this is a root span (no parent)."""
        return self.parent_span_id is None

    @property
    def is_error(self) -> bool:
        """Check if the span has an error status."""
        return self.status == SpanStatus.ERROR

    @property
    def is_complete(self) -> bool:
        """Check if the span has ended."""
        return self.end_time is not None

    @property
    def has_exceptions(self) -> bool:
        """Check if any exceptions were recorded."""
        return len(self.exceptions) > 0

    # ==========================================================================
    # Attribute Methods
    # ==========================================================================

    def set_attribute(self, key: str, value: Any) -> None:
        """
        Set a span attribute.

        Args:
            key: Attribute key (should use semantic conventions).
            value: Attribute value (must be JSON-serializable).
        """
        self.attributes[key] = value

    def set_attributes(self, attributes: Dict[str, Any]) -> None:
        """
        Set multiple attributes at once.

        Args:
            attributes: Dictionary of key-value pairs.
        """
        self.attributes.update(attributes)

    def get_attribute(self, key: str, default: Any = None) -> Any:
        """
        Get a span attribute.

        Args:
            key: Attribute key.
            default: Default value if key not found.

        Returns:
            The attribute value or default.
        """
        return self.attributes.get(key, default)

    # ==========================================================================
    # Event Methods
    # ==========================================================================

    def add_event(
        self,
        name: str,
        attributes: Optional[Metadata] = None,
        timestamp: Optional[Timestamp] = None,
    ) -> None:
        """
        Add an event to the span.

        Args:
            name: Event name.
            attributes: Optional event attributes.
            timestamp: Optional event timestamp (defaults to now).
        """
        event = SpanEvent(
            name=name,
            timestamp=timestamp or utc_now(),
            attributes=attributes or {},
        )
        self.events.append(event)

    # ==========================================================================
    # Link Methods
    # ==========================================================================

    def add_link(
        self,
        trace_id: TraceID,
        span_id: SpanID,
        attributes: Optional[Metadata] = None,
    ) -> None:
        """
        Add a link to another span.

        Args:
            trace_id: Linked span's trace ID.
            span_id: Linked span's ID.
            attributes: Optional link attributes.
        """
        link = SpanLink(
            trace_id=trace_id,
            span_id=span_id,
            attributes=attributes or {},
        )
        self.links.append(link)

    # ==========================================================================
    # Exception Methods
    # ==========================================================================

    def record_exception(
        self,
        exception: BaseException,
        escaped: bool = True,
    ) -> None:
        """
        Record an exception that occurred during the span.

        Args:
            exception: The exception that occurred.
            escaped: Whether the exception escaped the span.
        """
        exc_info = ExceptionInfo.from_exception(exception, escaped=escaped)
        self.exceptions.append(exc_info)

        # Set error status if exception escaped
        if escaped:
            self.set_status(SpanStatus.ERROR, str(exception))

    # ==========================================================================
    # Status Methods
    # ==========================================================================

    def set_status(
        self,
        status: SpanStatus,
        message: Optional[str] = None,
    ) -> None:
        """
        Set the span's status.

        Args:
            status: The new status.
            message: Optional status message.
        """
        self.status = status
        self.status_message = message

    def mark_ok(self, message: Optional[str] = None) -> None:
        """Mark the span as OK."""
        self.set_status(SpanStatus.OK, message)

    def mark_error(self, message: Optional[str] = None) -> None:
        """Mark the span as error."""
        self.set_status(SpanStatus.ERROR, message)

    # ==========================================================================
    # Lifecycle Methods
    # ==========================================================================

    def end(self, end_time: Optional[Timestamp] = None) -> None:
        """
        End the span.

        Args:
            end_time: Optional end time (defaults to now).
        """
        if self.end_time is None:
            self.end_time = end_time or utc_now()

    # ==========================================================================
    # Serialization
    # ==========================================================================

    def to_dict(self) -> Dict[str, Any]:
        """Convert span to dictionary for serialization.

        Uses camelCase for field names to match gateway API contract.
        """
        result: Dict[str, Any] = {
            "traceId": self.trace_id,
            "spanId": self.span_id,
            "name": self.name,
            "kind": self.kind.value,
            "startTime": self.start_time.isoformat(),
            "status": self.status.value,
            # Always include these for consistent schema
            "attributes": dict(self.attributes),
            "events": [e.to_dict() for e in self.events],
            "links": [link.to_dict() for link in self.links],
            "exceptions": [e.to_dict() for e in self.exceptions],
        }

        # Optional fields
        if self.parent_span_id:
            result["parentSpanId"] = self.parent_span_id
        if self.end_time:
            result["endTime"] = self.end_time.isoformat()
            result["durationMs"] = self.duration_ms
        if self.status_message:
            result["statusMessage"] = self.status_message

        # Risicare-specific
        if self.session_id:
            result["sessionId"] = self.session_id
        if self.agent_id:
            result["agentId"] = self.agent_id
        if self.agent_name:
            result["agentName"] = self.agent_name
        if self.agent_type:
            result["agentType"] = self.agent_type
        if self.semantic_phase:
            result["semanticPhase"] = self.semantic_phase.value
        # LLM fields — output as flat camelCase to match gateway schema
        llm = self.llm
        if llm:
            if llm.provider:
                result["llmProvider"] = llm.provider
            if llm.model:
                result["llmModel"] = llm.model
            if llm.prompt_tokens is not None:
                result["llmPromptTokens"] = llm.prompt_tokens
            if llm.completion_tokens is not None:
                result["llmCompletionTokens"] = llm.completion_tokens
            if llm.total_tokens is not None:
                result["llmTotalTokens"] = llm.total_tokens
            if llm.cost_usd is not None:
                result["llmCostUsd"] = llm.cost_usd
        else:
            # Fallback: extract from gen_ai.* attributes set by providers
            attrs = self.attributes
            model = attrs.get("gen_ai.response.model") or attrs.get("gen_ai.request.model")
            if model:
                result["llmModel"] = model
            provider = attrs.get("gen_ai.system")
            if provider:
                result["llmProvider"] = provider
            pt = attrs.get("gen_ai.usage.prompt_tokens")
            if pt is not None:
                result["llmPromptTokens"] = pt
            ct = attrs.get("gen_ai.usage.completion_tokens")
            if ct is not None:
                result["llmCompletionTokens"] = ct
            tt = attrs.get("gen_ai.usage.total_tokens")
            if tt is not None:
                result["llmTotalTokens"] = tt

        return result

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Span":
        """
        Create a Span from a dictionary.

        Supports both camelCase (gateway format) and snake_case (internal format).

        Args:
            data: Dictionary with span data (e.g., from JSON).

        Returns:
            A new Span instance.
        """
        from datetime import datetime

        # Helper to get value supporting both camelCase and snake_case
        def get(camel: str, snake: str, default: Any = None) -> Any:
            return data.get(camel) or data.get(snake, default)

        # Parse timestamps
        start_time = get("startTime", "start_time")
        if isinstance(start_time, str):
            start_time = datetime.fromisoformat(start_time.replace("Z", "+00:00"))
        elif start_time is None:
            start_time = utc_now()

        end_time = get("endTime", "end_time")
        if isinstance(end_time, str):
            end_time = datetime.fromisoformat(end_time.replace("Z", "+00:00"))

        # Parse enums
        kind = data.get("kind", "internal")
        if isinstance(kind, str):
            kind = SpanKind(kind)

        status = data.get("status", "unset")
        if isinstance(status, str):
            status = SpanStatus(status)

        semantic_phase = get("semanticPhase", "semantic_phase")
        if isinstance(semantic_phase, str):
            semantic_phase = SemanticPhase(semantic_phase)

        # Parse nested objects
        events = [
            SpanEvent.from_dict(e) if isinstance(e, dict) else e
            for e in data.get("events", [])
        ]
        links = [
            SpanLink.from_dict(link) if isinstance(link, dict) else link
            for link in data.get("links", [])
        ]
        exceptions = [
            ExceptionInfo.from_dict(e) if isinstance(e, dict) else e
            for e in data.get("exceptions", [])
        ]

        llm_data = data.get("llm")
        llm = LLMAttributes.from_dict(llm_data) if isinstance(llm_data, dict) else None

        return cls(
            trace_id=get("traceId", "trace_id"),
            span_id=get("spanId", "span_id"),
            name=data["name"],
            parent_span_id=get("parentSpanId", "parent_span_id"),
            kind=kind,
            start_time=start_time,
            end_time=end_time,
            status=status,
            status_message=get("statusMessage", "status_message"),
            attributes=dict(data.get("attributes", {})),
            events=events,
            links=links,
            exceptions=exceptions,
            session_id=get("sessionId", "session_id"),
            agent_id=get("agentId", "agent_id"),
            agent_name=get("agentName", "agent_name"),
            agent_type=get("agentType", "agent_type"),
            semantic_phase=semantic_phase,
            llm=llm,
        )

    # ==========================================================================
    # Factory Methods
    # ==========================================================================

    @classmethod
    def create(
        cls,
        name: str,
        trace_id: Optional[TraceID] = None,
        parent_span_id: Optional[SpanID] = None,
        kind: SpanKind = SpanKind.INTERNAL,
        **kwargs: Any,
    ) -> "Span":
        """
        Factory method to create a new span.

        Args:
            name: Span name.
            trace_id: Optional trace ID (generated if not provided).
            parent_span_id: Optional parent span ID.
            kind: Span kind.
            **kwargs: Additional span attributes.

        Returns:
            A new Span instance.
        """
        from risicare_core.types.base import generate_trace_id

        return cls(
            trace_id=trace_id or generate_trace_id(),
            span_id=generate_span_id(),
            name=name,
            parent_span_id=parent_span_id,
            kind=kind,
            **kwargs,
        )
