"""OpenTelemetry span wrappers for Mingx (MingxSpan, MingxGeneration, MingxEvent)."""

from datetime import datetime
from time import time_ns
from typing import TYPE_CHECKING, Any, Optional, Union

from opentelemetry import trace as otel_trace_api
from opentelemetry.trace.status import Status, StatusCode

from mingx._client.attributes import (
    MingxOtelSpanAttributes,
    create_generation_attributes,
    create_span_attributes,
    create_trace_attributes,
)

if TYPE_CHECKING:
    from mingx._client.client import MingxClient


def _format_trace_id(trace_id_int: int) -> str:
    return format(trace_id_int, "032x")


def _format_span_id(span_id_int: int) -> str:
    return format(span_id_int, "016x")


def _serialize_attr(obj: Any) -> Optional[str]:
    if obj is None or isinstance(obj, str):
        return obj
    if isinstance(obj, datetime):
        return obj.isoformat()
    try:
        return str(obj) if not isinstance(obj, (dict, list)) else __import__("json").dumps(obj, default=str)
    except (TypeError, ValueError):
        return str(obj)


class MingxObservationWrapper:
    """Base wrapper for Mingx span/observation types."""

    def __init__(
        self,
        *,
        otel_span: otel_trace_api.Span,
        client: "MingxClient",
        as_type: str,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        environment: Optional[str] = None,
        version: Optional[str] = None,
        level: Optional[str] = None,
        status_message: Optional[str] = None,
    ) -> None:
        self._otel_span = otel_span
        self._client = client
        self._observation_type = as_type
        self._environment = environment or getattr(client, "_environment", None)

        ctx = otel_span.get_span_context()
        self.trace_id = _format_trace_id(ctx.trace_id)
        self.span_id = _format_span_id(ctx.span_id)
        self.id = self.span_id

        otel_span.set_attribute(MingxOtelSpanAttributes.OBSERVATION_TYPE, as_type)
        if self._environment:
            otel_span.set_attribute(MingxOtelSpanAttributes.ENVIRONMENT, self._environment)

        attrs = create_span_attributes(
            observation_type=as_type,
            input=input,
            output=output,
            metadata=metadata,
            version=version,
            level=level,
            status_message=status_message,
        )
        otel_span.set_attributes({k: v for k, v in attrs.items() if v is not None})
        if level == "ERROR" and otel_span.is_recording():
            try:
                otel_span.set_status(Status(StatusCode.ERROR, description=status_message))
            except Exception:
                pass

    def end(self, *, end_time: Optional[int] = None) -> "MingxObservationWrapper":
        self._otel_span.end(end_time=end_time)
        return self

    def update_trace(
        self,
        *,
        name: Optional[str] = None,
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        version: Optional[str] = None,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        tags: Optional[list[str]] = None,
    ) -> "MingxObservationWrapper":
        if not self._otel_span.is_recording():
            return self
        attrs = create_trace_attributes(
            name=name,
            user_id=user_id,
            session_id=session_id,
            version=version,
            input=input,
            output=output,
            metadata=metadata,
            tags=tags,
        )
        self._otel_span.set_attributes(attrs)
        return self

    def update(
        self,
        *,
        name: Optional[str] = None,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[str] = None,
        status_message: Optional[str] = None,
        **kwargs: Any,
    ) -> "MingxObservationWrapper":
        if not self._otel_span.is_recording():
            return self
        if name:
            self._otel_span.update_name(name)
        attrs = create_span_attributes(
            observation_type=self._observation_type,
            input=input,
            output=output,
            metadata=metadata,
            version=version,
            level=level,
            status_message=status_message,
        )
        self._otel_span.set_attributes({k: v for k, v in attrs.items() if v is not None})
        if level == "ERROR":
            try:
                self._otel_span.set_status(Status(StatusCode.ERROR, description=status_message))
            except Exception:
                pass
        return self


class MingxSpan(MingxObservationWrapper):
    """Generic span for non-LLM operations."""

    def __init__(
        self,
        *,
        otel_span: otel_trace_api.Span,
        client: "MingxClient",
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        environment: Optional[str] = None,
        version: Optional[str] = None,
        level: Optional[str] = None,
        status_message: Optional[str] = None,
    ) -> None:
        super().__init__(
            otel_span=otel_span,
            client=client,
            as_type="span",
            input=input,
            output=output,
            metadata=metadata,
            environment=environment,
            version=version,
            level=level,
            status_message=status_message,
        )

    def start_observation(
        self,
        name: str,
        *,
        as_type: str = "span",
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        **kwargs: Any,
    ) -> Union["MingxSpan", "MingxGeneration", "MingxEvent"]:
        """Start a child observation. Caller must call .end()."""
        with otel_trace_api.use_span(self._otel_span):
            return self._client.start_observation(
                name=name,
                as_type=as_type,
                input=input,
                output=output,
                metadata=metadata,
                **kwargs,
            )

    def start_as_current_span(
        self,
        name: str,
        *,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        **kwargs: Any,
    ) -> Any:
        """Enter a child span as current (context manager). Delegates to client."""
        with otel_trace_api.use_span(self._otel_span):
            return self._client.start_as_current_observation(
                name=name,
                as_type="span",
                input=input,
                output=output,
                metadata=metadata,
                **kwargs,
            )

    def start_as_current_generation(
        self,
        name: str,
        *,
        model: Optional[str] = None,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        usage_details: Optional[dict[str, int]] = None,
        cost_details: Optional[dict[str, float]] = None,
        **kwargs: Any,
    ) -> Any:
        """Enter a child generation as current (context manager)."""
        with otel_trace_api.use_span(self._otel_span):
            return self._client.start_as_current_observation(
                name=name,
                as_type="generation",
                model=model,
                input=input,
                output=output,
                metadata=metadata,
                usage_details=usage_details,
                cost_details=cost_details,
                **kwargs,
            )


class MingxGeneration(MingxObservationWrapper):
    """Span for LLM/model generations with model, usage, cost."""

    def __init__(
        self,
        *,
        otel_span: otel_trace_api.Span,
        client: "MingxClient",
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        environment: Optional[str] = None,
        version: Optional[str] = None,
        level: Optional[str] = None,
        status_message: Optional[str] = None,
        completion_start_time: Optional[datetime] = None,
        model: Optional[str] = None,
        model_parameters: Optional[dict[str, Any]] = None,
        usage_details: Optional[dict[str, int]] = None,
        cost_details: Optional[dict[str, float]] = None,
    ) -> None:
        self._otel_span = otel_span
        self._client = client
        self._observation_type = "generation"
        self._environment = environment or getattr(client, "_environment", None)

        ctx = otel_span.get_span_context()
        self.trace_id = _format_trace_id(ctx.trace_id)
        self.span_id = _format_span_id(ctx.span_id)
        self.id = self.span_id

        otel_span.set_attribute(MingxOtelSpanAttributes.OBSERVATION_TYPE, "generation")
        if self._environment:
            otel_span.set_attribute(MingxOtelSpanAttributes.ENVIRONMENT, self._environment)

        ct = _serialize_attr(completion_start_time) if completion_start_time else None
        attrs = create_generation_attributes(
            input=input,
            output=output,
            metadata=metadata,
            version=version,
            level=level,
            status_message=status_message,
            model=model,
            model_parameters=model_parameters,
            usage_details=usage_details,
            cost_details=cost_details,
            completion_start_time=ct,
        )
        otel_span.set_attributes({k: v for k, v in attrs.items() if v is not None})
        if level == "ERROR" and otel_span.is_recording():
            try:
                otel_span.set_status(Status(StatusCode.ERROR, description=status_message))
            except Exception:
                pass

    def update(
        self,
        *,
        name: Optional[str] = None,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[str] = None,
        status_message: Optional[str] = None,
        completion_start_time: Optional[datetime] = None,
        model: Optional[str] = None,
        model_parameters: Optional[dict[str, Any]] = None,
        usage_details: Optional[dict[str, int]] = None,
        cost_details: Optional[dict[str, float]] = None,
        **kwargs: Any,
    ) -> "MingxGeneration":
        if not self._otel_span.is_recording():
            return self
        if name:
            self._otel_span.update_name(name)
        ct = _serialize_attr(completion_start_time) if completion_start_time else None
        attrs = create_generation_attributes(
            input=input,
            output=output,
            metadata=metadata,
            version=version,
            level=level,
            status_message=status_message,
            model=model,
            model_parameters=model_parameters,
            usage_details=usage_details,
            cost_details=cost_details,
            completion_start_time=ct,
        )
        self._otel_span.set_attributes({k: v for k, v in attrs.items() if v is not None})
        if level == "ERROR":
            try:
                self._otel_span.set_status(Status(StatusCode.ERROR, description=status_message))
            except Exception:
                pass
        return self

    def start_observation(
        self,
        name: str,
        *,
        as_type: str = "span",
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        **kwargs: Any,
    ) -> Union["MingxSpan", "MingxGeneration", "MingxEvent"]:
        """Start a child observation. Caller must call .end()."""
        with otel_trace_api.use_span(self._otel_span):
            return self._client.start_observation(
                name=name,
                as_type=as_type,
                input=input,
                output=output,
                metadata=metadata,
                **kwargs,
            )


class MingxEvent(MingxObservationWrapper):
    """Point-in-time event (start_time = end_time)."""

    def __init__(
        self,
        *,
        otel_span: otel_trace_api.Span,
        client: "MingxClient",
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        environment: Optional[str] = None,
        version: Optional[str] = None,
        level: Optional[str] = None,
        status_message: Optional[str] = None,
    ) -> None:
        super().__init__(
            otel_span=otel_span,
            client=client,
            as_type="event",
            input=input,
            output=output,
            metadata=metadata,
            environment=environment,
            version=version,
            level=level,
            status_message=status_message,
        )

    def update(
        self,
        *,
        name: Optional[str] = None,
        input: Optional[Any] = None,
        output: Optional[Any] = None,
        metadata: Optional[Any] = None,
        version: Optional[str] = None,
        level: Optional[str] = None,
        status_message: Optional[str] = None,
        **kwargs: Any,
    ) -> "MingxEvent":
        # Events are typically immutable; allow update for consistency but log or no-op if preferred
        return super().update(
            name=name,
            input=input,
            output=output,
            metadata=metadata,
            version=version,
            level=level,
            status_message=status_message,
            **kwargs,
        )
