"""Programmatic span builder for manual span creation."""

import json
from typing import Any

from opentelemetry import trace as trace_api
from opentelemetry.trace import Link, Status, StatusCode

from plyra_trace import _tracer
from plyra_trace.models import SpanLink
from plyra_trace.semconv import SpanAttributes, SpanKind


class SpanBuilder:
    """
    Programmatic span builder for cases where decorators aren't sufficient.

    Example:
        >>> with SpanBuilder("orchestrator-step", kind=SpanKind.AGENT) as span:
        ...     span.set_input({"query": "analyze competitors"})
        ...     span.set_agent(name="orchestrator", description="Routes to agents")
        ...     result = do_work()
        ...     span.set_output(result)
    """

    def __init__(
        self,
        name: str,
        kind: SpanKind = SpanKind.CHAIN,
        links: list[SpanLink] | None = None,
    ):
        """
        Create a span builder.

        Args:
            name: Span name
            kind: Span kind (default: CHAIN)
            links: Optional list of span links to other traces
        """
        self.name = name
        self.kind = kind
        self.links = links
        self._span: trace_api.Span | None = None
        self._tracer = _tracer.get_tracer()

    def __enter__(self) -> "SpanBuilder":
        """Start the span when entering the context."""
        # Convert SpanLink models to OTel Link objects
        otel_links = []
        if self.links:
            for link in self.links:
                otel_link = Link(
                    context=trace_api.SpanContext(
                        trace_id=int(link.trace_id, 16),
                        span_id=int(link.span_id, 16),
                        is_remote=True,
                        trace_flags=trace_api.TraceFlags(0x01),
                    ),
                    attributes=link.attributes,
                )
                otel_links.append(otel_link)

        # Start span
        self._span = self._tracer.start_span(
            name=self.name,
            links=otel_links if otel_links else None,
        )

        # Set span kind attributes
        self._span.set_attribute(SpanAttributes.SPAN_KIND, self.kind.value)
        self._span.set_attribute(SpanAttributes.OPENINFERENCE_SPAN_KIND, self.kind.value)

        # Make this span the current span
        self._context_token = trace_api.context_api.attach(trace_api.set_span_in_context(self._span))

        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """End the span when exiting the context."""
        if self._span is None:
            return

        # Set status based on exception
        if exc_type is not None:
            self._span.set_status(Status(StatusCode.ERROR, str(exc_val)))
            self._span.record_exception(exc_val)
        else:
            self._span.set_status(Status(StatusCode.OK))

        # End span and detach context
        self._span.end()
        trace_api.context_api.detach(self._context_token)

    def set_input(self, value: Any, mime_type: str = "application/json") -> "SpanBuilder":
        """
        Set input value on the span.

        Args:
            value: Input value (will be JSON-serialized)
            mime_type: MIME type of the input

        Returns:
            Self for chaining
        """
        if self._span:
            try:
                if isinstance(value, str):
                    serialized = value
                else:
                    serialized = json.dumps(value, default=str)
                self._span.set_attribute(SpanAttributes.INPUT_VALUE, serialized)
                self._span.set_attribute(SpanAttributes.INPUT_MIME_TYPE, mime_type)
            except (TypeError, ValueError):
                self._span.set_attribute(SpanAttributes.INPUT_VALUE, str(value))
                self._span.set_attribute(SpanAttributes.INPUT_MIME_TYPE, "text/plain")
        return self

    def set_output(self, value: Any, mime_type: str = "application/json") -> "SpanBuilder":
        """
        Set output value on the span.

        Args:
            value: Output value (will be JSON-serialized)
            mime_type: MIME type of the output

        Returns:
            Self for chaining
        """
        if self._span:
            try:
                if isinstance(value, str):
                    serialized = value
                else:
                    serialized = json.dumps(value, default=str)
                self._span.set_attribute(SpanAttributes.OUTPUT_VALUE, serialized)
                self._span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, mime_type)
            except (TypeError, ValueError):
                self._span.set_attribute(SpanAttributes.OUTPUT_VALUE, str(value))
                self._span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "text/plain")
        return self

    def set_status(self, status: str) -> "SpanBuilder":
        """
        Set span status.

        Args:
            status: "ok" or "error"

        Returns:
            Self for chaining
        """
        if self._span:
            if status.lower() == "ok":
                self._span.set_status(Status(StatusCode.OK))
            else:
                self._span.set_status(Status(StatusCode.ERROR))
        return self

    def set_agent(
        self,
        name: str,
        description: str | None = None,
        framework: str | None = None,
    ) -> "SpanBuilder":
        """
        Set agent attributes on the span.

        Args:
            name: Agent name
            description: Agent description
            framework: Framework name

        Returns:
            Self for chaining
        """
        if self._span:
            self._span.set_attribute(SpanAttributes.AGENT_NAME, name)
            if description:
                self._span.set_attribute(SpanAttributes.AGENT_DESCRIPTION, description)
            if framework:
                self._span.set_attribute(SpanAttributes.AGENT_FRAMEWORK, framework)
        return self

    def set_guard(
        self,
        policy: str,
        action: str,
        triggered: bool,
        confidence: float | None = None,
        details: dict | None = None,
    ) -> "SpanBuilder":
        """
        Set guard attributes on the span.

        Args:
            policy: Policy name
            action: Action taken ("allow", "block", etc.)
            triggered: Whether the guard was triggered
            confidence: Confidence score (0.0-1.0)
            details: Additional details

        Returns:
            Self for chaining
        """
        if self._span:
            self._span.set_attribute(SpanAttributes.GUARD_POLICY, policy)
            self._span.set_attribute(SpanAttributes.GUARD_ACTION, action)
            self._span.set_attribute(SpanAttributes.GUARD_TRIGGERED, triggered)
            if confidence is not None:
                self._span.set_attribute(SpanAttributes.GUARD_CONFIDENCE, confidence)
            if details is not None:
                self._span.set_attribute(SpanAttributes.GUARD_DETAILS, json.dumps(details))
        return self

    def set_llm(
        self,
        model: str,
        provider: str | None = None,
        token_counts: dict | None = None,
    ) -> "SpanBuilder":
        """
        Set LLM attributes on the span.

        Args:
            model: Model name
            provider: Provider name
            token_counts: Token count dictionary with keys like
                "prompt", "completion", "total"

        Returns:
            Self for chaining
        """
        if self._span:
            self._span.set_attribute(SpanAttributes.LLM_MODEL_NAME, model)
            if provider:
                self._span.set_attribute(SpanAttributes.LLM_PROVIDER, provider)
                self._span.set_attribute(SpanAttributes.LLM_SYSTEM, provider)
            if token_counts:
                if "prompt" in token_counts:
                    self._span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, token_counts["prompt"])
                if "completion" in token_counts:
                    self._span.set_attribute(
                        SpanAttributes.LLM_TOKEN_COUNT_COMPLETION,
                        token_counts["completion"],
                    )
                if "total" in token_counts:
                    self._span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, token_counts["total"])
        return self

    def set_attribute(self, key: str, value: Any) -> "SpanBuilder":
        """
        Set a custom attribute on the span.

        Args:
            key: Attribute key
            value: Attribute value

        Returns:
            Self for chaining
        """
        if self._span:
            if isinstance(value, (dict, list)):
                self._span.set_attribute(key, json.dumps(value))
            else:
                self._span.set_attribute(key, value)
        return self

    def add_event(self, name: str, attributes: dict | None = None) -> "SpanBuilder":
        """
        Add an event to the span.

        Args:
            name: Event name
            attributes: Event attributes

        Returns:
            Self for chaining
        """
        if self._span:
            self._span.add_event(name, attributes=attributes)
        return self
