"""
Observation wrapper classes for enhanced span management.

Wraps OpenTelemetry spans with type-specific methods and attributes.
"""

import json
import time
from dataclasses import dataclass
from typing import Any, Dict, List, Optional

from opentelemetry.trace import Span, Status, StatusCode

from ..types import Attrs
from .types import ObservationType


@dataclass
class UsageDetails:
    """
    Flexible usage details for token tracking.

    Supports multiple token types including cache, audio, reasoning, etc.
    """

    input_tokens: int = 0
    output_tokens: int = 0
    total_tokens: int = 0

    # Extended token types
    cache_read_tokens: int = 0
    cache_creation_tokens: int = 0
    audio_input_tokens: int = 0
    audio_output_tokens: int = 0
    reasoning_tokens: int = 0
    image_tokens: int = 0

    def to_dict(self) -> Dict[str, int]:
        """Convert to dictionary, excluding zero values."""
        result = {}
        if self.input_tokens > 0:
            result["input_tokens"] = self.input_tokens
        if self.output_tokens > 0:
            result["output_tokens"] = self.output_tokens
        if self.total_tokens > 0:
            result["total_tokens"] = self.total_tokens
        if self.cache_read_tokens > 0:
            result["cache_read_tokens"] = self.cache_read_tokens
        if self.cache_creation_tokens > 0:
            result["cache_creation_tokens"] = self.cache_creation_tokens
        if self.audio_input_tokens > 0:
            result["audio_input_tokens"] = self.audio_input_tokens
        if self.audio_output_tokens > 0:
            result["audio_output_tokens"] = self.audio_output_tokens
        if self.reasoning_tokens > 0:
            result["reasoning_tokens"] = self.reasoning_tokens
        if self.image_tokens > 0:
            result["image_tokens"] = self.image_tokens
        return result


class BrokleObservation:
    """
    Base observation wrapper for OpenTelemetry spans.

    Provides enhanced methods for updating span attributes and
    managing observation lifecycle.

    Example:
        >>> with client.start_observation("my-operation") as obs:
        ...     obs.update(metadata={"key": "value"})
        ...     obs.set_output(result)
    """

    def __init__(
        self,
        span: Span,
        observation_type: ObservationType = ObservationType.SPAN,
    ):
        """
        Initialize observation wrapper.

        Args:
            span: Underlying OpenTelemetry span
            observation_type: Type of observation
        """
        self._span = span
        self._observation_type = observation_type
        self._start_time = time.perf_counter()
        self._completion_start_time: Optional[float] = None
        self._ended = False

    @property
    def span(self) -> Span:
        """Access the underlying OpenTelemetry span."""
        return self._span

    @property
    def observation_type(self) -> ObservationType:
        """Get the observation type."""
        return self._observation_type

    @property
    def trace_id(self) -> str:
        """Get the trace ID as hex string."""
        ctx = self._span.get_span_context()
        return format(ctx.trace_id, "032x") if ctx else ""

    @property
    def span_id(self) -> str:
        """Get the span ID as hex string."""
        ctx = self._span.get_span_context()
        return format(ctx.span_id, "016x") if ctx else ""

    def update(
        self,
        name: Optional[str] = None,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Dict[str, Any]] = None,
        level: Optional[str] = None,
        version: Optional[str] = None,
        **extra_attributes,
    ) -> "BrokleObservation":
        """
        Update observation attributes.

        Args:
            name: Update span name
            input: Set input value
            output: Set output value
            metadata: Update metadata (merged with existing)
            level: Set span level (DEBUG, DEFAULT, WARNING, ERROR)
            version: Set version
            **extra_attributes: Additional span attributes

        Returns:
            Self for chaining
        """
        if self._ended:
            return self

        if name:
            # Note: OpenTelemetry spans don't support name updates after creation
            # Store as attribute for backend to use
            self._span.set_attribute(Attrs.BROKLE_SPAN_NAME, name)

        if input is not None:
            input_str = self._serialize(input)
            self._span.set_attribute(Attrs.INPUT_VALUE, input_str)
            self._span.set_attribute(Attrs.INPUT_MIME_TYPE, "application/json")

        if output is not None:
            output_str = self._serialize(output)
            self._span.set_attribute(Attrs.OUTPUT_VALUE, output_str)
            self._span.set_attribute(Attrs.OUTPUT_MIME_TYPE, "application/json")

        if metadata:
            self._span.set_attribute(Attrs.BROKLE_TRACE_METADATA, json.dumps(metadata))

        if level:
            self._span.set_attribute(Attrs.BROKLE_SPAN_LEVEL, level)

        if version:
            self._span.set_attribute(Attrs.BROKLE_VERSION, version)

        for key, value in extra_attributes.items():
            self._span.set_attribute(key, value)

        return self

    def set_input(self, input: Any) -> "BrokleObservation":
        """Set input value."""
        return self.update(input=input)

    def set_output(self, output: Any) -> "BrokleObservation":
        """Set output value."""
        return self.update(output=output)

    def set_error(self, error: Exception) -> "BrokleObservation":
        """
        Record an error on this observation.

        Args:
            error: Exception to record

        Returns:
            Self for chaining
        """
        if not self._ended:
            self._span.set_status(Status(StatusCode.ERROR, str(error)))
            self._span.record_exception(error)
            self._span.set_attribute(Attrs.BROKLE_SPAN_LEVEL, "ERROR")
        return self

    def update_trace(
        self,
        session_id: Optional[str] = None,
        user_id: Optional[str] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> "BrokleObservation":
        """
        Update trace-level attributes (propagated to all spans in trace).

        Args:
            session_id: Session grouping identifier
            user_id: User identifier
            tags: Trace tags (additive - merged with existing)
            metadata: Trace metadata (merged with existing)

        Returns:
            Self for chaining
        """
        if self._ended:
            return self

        if session_id:
            self._span.set_attribute(Attrs.SESSION_ID, session_id)
        if user_id:
            self._span.set_attribute(Attrs.USER_ID, user_id)
        if tags:
            self._span.set_attribute(Attrs.BROKLE_TRACE_TAGS, json.dumps(tags))
        if metadata:
            self._span.set_attribute(Attrs.BROKLE_TRACE_METADATA, json.dumps(metadata))

        return self

    def end(self) -> None:
        """End the observation (span)."""
        if not self._ended:
            self._span.end()
            self._ended = True

    def __enter__(self) -> "BrokleObservation":
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        if exc_val:
            self.set_error(exc_val)
        self.end()
        return False

    @staticmethod
    def _serialize(value: Any) -> str:
        """Serialize value to JSON string."""
        try:
            if value is None:
                return "null"
            if isinstance(value, str):
                return json.dumps(value)
            if hasattr(value, "model_dump"):
                return json.dumps(value.model_dump(exclude_none=True))
            if hasattr(value, "__dataclass_fields__"):
                import dataclasses

                return json.dumps(dataclasses.asdict(value))
            return json.dumps(value, default=str)
        except Exception:
            return json.dumps(str(value))


class BrokleGeneration(BrokleObservation):
    """
    Generation observation for LLM calls.

    Provides additional methods for tracking:
    - Model and parameters
    - Token usage
    - Completion start time (for TTFT)
    - Cost calculation hints

    Example:
        >>> with client.start_generation(
        ...     name="chat",
        ...     model="gpt-4",
        ... ) as gen:
        ...     response = openai.chat.completions.create(...)
        ...     gen.mark_completion_start()  # Mark first token
        ...     gen.set_usage(input_tokens=100, output_tokens=50)
        ...     gen.set_output(response.choices[0].message.content)
    """

    def __init__(
        self,
        span: Span,
        model: Optional[str] = None,
        provider: Optional[str] = None,
    ):
        super().__init__(span, ObservationType.GENERATION)
        self._model = model
        self._provider = provider

        # Set generation-specific attributes
        if model:
            span.set_attribute(Attrs.GEN_AI_REQUEST_MODEL, model)
        if provider:
            span.set_attribute(Attrs.GEN_AI_PROVIDER_NAME, provider)

    def mark_completion_start(self) -> "BrokleGeneration":
        """
        Mark the completion start time (first token received).

        Call this when the first token of a streaming response arrives
        to accurately track time-to-first-token (TTFT).

        Returns:
            Self for chaining
        """
        if self._completion_start_time is None:
            self._completion_start_time = time.perf_counter()

            # Calculate TTFT
            ttft_ms = (self._completion_start_time - self._start_time) * 1000
            self._span.set_attribute(Attrs.GEN_AI_RESPONSE_TTFT, ttft_ms)

        return self

    def set_usage(
        self,
        input_tokens: int = 0,
        output_tokens: int = 0,
        total_tokens: Optional[int] = None,
        **extended_usage,
    ) -> "BrokleGeneration":
        """
        Set token usage for this generation.

        Args:
            input_tokens: Number of input/prompt tokens
            output_tokens: Number of output/completion tokens
            total_tokens: Total tokens (calculated if not provided)
            **extended_usage: Extended token types (cache_read_tokens, audio_tokens, etc.)

        Returns:
            Self for chaining
        """
        if self._ended:
            return self

        # Calculate total if not provided
        if total_tokens is None:
            total_tokens = input_tokens + output_tokens

        # Set standard usage attributes
        if input_tokens > 0:
            self._span.set_attribute(Attrs.GEN_AI_USAGE_INPUT_TOKENS, input_tokens)
        if output_tokens > 0:
            self._span.set_attribute(Attrs.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens)
        if total_tokens > 0:
            self._span.set_attribute(Attrs.BROKLE_USAGE_TOTAL_TOKENS, total_tokens)

        # Set extended usage attributes
        if "cache_read_tokens" in extended_usage:
            self._span.set_attribute(
                Attrs.GEN_AI_USAGE_INPUT_TOKENS_CACHE_READ,
                extended_usage["cache_read_tokens"],
            )
        if "cache_creation_tokens" in extended_usage:
            self._span.set_attribute(
                Attrs.GEN_AI_USAGE_INPUT_TOKENS_CACHE_CREATION,
                extended_usage["cache_creation_tokens"],
            )
        if "audio_input_tokens" in extended_usage:
            self._span.set_attribute(
                Attrs.GEN_AI_USAGE_INPUT_AUDIO_TOKENS,
                extended_usage["audio_input_tokens"],
            )
        if "audio_output_tokens" in extended_usage:
            self._span.set_attribute(
                Attrs.GEN_AI_USAGE_OUTPUT_AUDIO_TOKENS,
                extended_usage["audio_output_tokens"],
            )
        if "reasoning_tokens" in extended_usage:
            self._span.set_attribute(
                Attrs.GEN_AI_USAGE_REASONING_TOKENS, extended_usage["reasoning_tokens"]
            )

        return self

    def set_model_parameters(
        self,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        top_p: Optional[float] = None,
        frequency_penalty: Optional[float] = None,
        presence_penalty: Optional[float] = None,
        **extra_params,
    ) -> "BrokleGeneration":
        """
        Set model parameters for this generation.

        Args:
            temperature: Sampling temperature
            max_tokens: Maximum tokens to generate
            top_p: Nucleus sampling parameter
            frequency_penalty: Frequency penalty
            presence_penalty: Presence penalty
            **extra_params: Additional model-specific parameters

        Returns:
            Self for chaining
        """
        if self._ended:
            return self

        if temperature is not None:
            self._span.set_attribute(Attrs.GEN_AI_REQUEST_TEMPERATURE, temperature)
        if max_tokens is not None:
            self._span.set_attribute(Attrs.GEN_AI_REQUEST_MAX_TOKENS, max_tokens)
        if top_p is not None:
            self._span.set_attribute(Attrs.GEN_AI_REQUEST_TOP_P, top_p)
        if frequency_penalty is not None:
            self._span.set_attribute(
                Attrs.GEN_AI_REQUEST_FREQUENCY_PENALTY, frequency_penalty
            )
        if presence_penalty is not None:
            self._span.set_attribute(
                Attrs.GEN_AI_REQUEST_PRESENCE_PENALTY, presence_penalty
            )

        return self

    def set_finish_reason(self, reason: str) -> "BrokleGeneration":
        """Set the finish/stop reason."""
        if not self._ended:
            self._span.set_attribute(Attrs.GEN_AI_RESPONSE_FINISH_REASON, reason)
        return self


class BrokleEvent(BrokleObservation):
    """
    Event observation for point-in-time occurrences.

    Events have no meaningful duration - they mark specific moments.

    Example:
        >>> with client.start_event("user-action") as event:
        ...     event.update(metadata={"action": "click", "target": "submit"})
    """

    def __init__(self, span: Span):
        super().__init__(span, ObservationType.EVENT)


class BrokleAgent(BrokleObservation):
    """
    Agent observation for autonomous AI agent operations.

    Tracks agent reasoning, tool usage, and multi-step operations.

    Example:
        >>> with client.start_agent("research-agent") as agent:
        ...     agent.update(metadata={"objective": "Find information about X"})
        ...     # Agent loop
        ...     agent.set_output(final_result)
    """

    def __init__(self, span: Span):
        super().__init__(span, ObservationType.AGENT)

    def set_tools(self, tools: List[str]) -> "BrokleAgent":
        """Set available tools for this agent."""
        if not self._ended:
            self._span.set_attribute("brokle.agent.tools", json.dumps(tools))
        return self

    def set_iterations(self, count: int) -> "BrokleAgent":
        """Set the number of agent iterations/steps."""
        if not self._ended:
            self._span.set_attribute("brokle.agent.iterations", count)
        return self


class BrokleTool(BrokleObservation):
    """
    Tool observation for function/tool calls.

    Tracks tool name, input parameters, and results.

    Example:
        >>> with client.start_tool("calculator") as tool:
        ...     tool.set_input({"operation": "add", "a": 1, "b": 2})
        ...     result = calculator.add(1, 2)
        ...     tool.set_output(result)
    """

    def __init__(self, span: Span, tool_name: Optional[str] = None):
        super().__init__(span, ObservationType.TOOL)
        if tool_name:
            span.set_attribute("brokle.tool.name", tool_name)


class BrokleRetrieval(BrokleObservation):
    """
    Retrieval observation for RAG operations.

    Tracks retrieval queries, document counts, and relevance scores.

    Example:
        >>> with client.start_retrieval("vector-search") as retrieval:
        ...     retrieval.set_query(query_text)
        ...     docs = vector_db.search(query_text, k=5)
        ...     retrieval.set_documents(docs)
    """

    def __init__(self, span: Span):
        super().__init__(span, ObservationType.RETRIEVAL)

    def set_query(self, query: str) -> "BrokleRetrieval":
        """Set the retrieval query."""
        if not self._ended:
            self._span.set_attribute("brokle.retrieval.query", query)
        return self

    def set_documents(
        self,
        documents: List[Dict[str, Any]],
        count: Optional[int] = None,
    ) -> "BrokleRetrieval":
        """
        Set retrieved documents.

        Args:
            documents: List of retrieved documents
            count: Document count (calculated from list if not provided)

        Returns:
            Self for chaining
        """
        if not self._ended:
            doc_count = count if count is not None else len(documents)
            self._span.set_attribute("brokle.retrieval.document_count", doc_count)
            # Store first few documents for debugging (limit to avoid huge spans)
            preview = documents[:5] if len(documents) > 5 else documents
            self._span.set_attribute(
                "brokle.retrieval.documents", json.dumps(preview, default=str)
            )
        return self

    def set_top_k(self, k: int) -> "BrokleRetrieval":
        """Set the top-k parameter used for retrieval."""
        if not self._ended:
            self._span.set_attribute("brokle.retrieval.top_k", k)
        return self
