"""Custom OpenTelemetry API implementations that feed spans into BatchTransport."""
from __future__ import annotations

import contextlib
import json
import logging
import random
import time
from collections.abc import Iterator, Sequence
from contextvars import ContextVar
from datetime import datetime, timezone
from typing import Any

from openinference.semconv.trace import SpanAttributes as OISpanAttributes
from opentelemetry import context as otel_context
from opentelemetry import trace
from opentelemetry.trace import (
    Link,
    NonRecordingSpan,
    SpanContext,
    SpanKind,
    StatusCode,
    TraceFlags,
)
from opentelemetry.trace import Span as SpanBase
from opentelemetry.trace import Tracer as TracerBase
from opentelemetry.trace import TracerProvider as TracerProviderBase
from opentelemetry.util.types import AttributeValue

from .transport import BatchTransport

logger = logging.getLogger("staso")

session_id_var: ContextVar[str] = ContextVar("_al_session_id", default="")

# Map OI span kinds to backend kinds
_OI_KIND_MAP: dict[str, str] = {
    "AGENT": "agent",
    "LLM": "llm",
    "TOOL": "tool",
    "CHAIN": "chain",
    "EMBEDDING": "llm",
    "RETRIEVER": "tool",
}


def _hex_to_uuid(hex_str: str) -> str:
    """Convert a hex string to UUID format (xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx)."""
    h = hex_str.zfill(32)
    return f"{h[:8]}-{h[8:12]}-{h[12:16]}-{h[16:20]}-{h[20:32]}"


def _safe_json(obj: Any) -> str:
    try:
        return json.dumps(obj, default=str)
    except (TypeError, ValueError, OverflowError):
        return "{}"


class Span(SpanBase):
    """A lightweight Span that serializes to the backend schema on end()."""

    def __init__(
        self,
        name: str,
        span_context: SpanContext,
        parent_span_id: int | None,
        transport: BatchTransport | None,
        agent_id: str,
        session_id: str,
        attributes: dict[str, AttributeValue] | None = None,
    ) -> None:
        self._name = name
        self._span_context = span_context
        self._parent_span_id = parent_span_id
        self._transport = transport
        self._agent_id = agent_id
        self._session_id = session_id
        self._start_time = time.time()
        self._end_time: float | None = None
        self._attributes: dict[str, AttributeValue] = dict(attributes or {})
        self._status_code = StatusCode.UNSET
        self._status_description: str | None = None
        self._recording = True

    def get_span_context(self) -> SpanContext:
        return self._span_context

    def is_recording(self) -> bool:
        return self._recording

    def set_attribute(self, key: str, value: AttributeValue) -> None:
        try:
            if self._recording:
                self._attributes[key] = value
        except Exception:
            logger.debug("staso: failed to set attribute", exc_info=True)

    def set_attributes(self, attributes: dict[str, AttributeValue]) -> None:
        try:
            if self._recording:
                self._attributes.update(attributes)
        except Exception:
            logger.debug("staso: failed to set attributes", exc_info=True)

    def add_event(
        self,
        name: str,
        attributes: trace.types.Attributes = None,
        timestamp: int | None = None,
    ) -> None:
        pass  # not needed for our backend

    def add_link(
        self,
        context: SpanContext,
        attributes: trace.types.Attributes = None,
    ) -> None:
        pass

    def update_name(self, name: str) -> None:
        if self._recording:
            self._name = name

    def set_status(self, status: StatusCode, description: str | None = None) -> None:
        try:
            if self._recording:
                self._status_code = status
                self._status_description = description
        except Exception:
            logger.debug("staso: failed to set status", exc_info=True)

    def record_exception(
        self,
        exception: BaseException,
        attributes: trace.types.Attributes = None,
        timestamp: int | None = None,
        escaped: bool = False,
    ) -> None:
        try:
            if self._recording:
                self._status_code = StatusCode.ERROR
                self._status_description = str(exception)
                self._attributes["exception.type"] = type(exception).__name__
                self._attributes["exception.message"] = str(exception)
        except Exception:
            logger.debug("staso: failed to record exception", exc_info=True)

    def end(self, end_time: int | None = None) -> None:
        try:
            if not self._recording:
                return
            self._recording = False
            self._end_time = time.time()

            if self._transport is None:
                return

            # Extract OI attributes for top-level span fields
            attrs = dict(self._attributes)
            oi_kind = str(attrs.pop(OISpanAttributes.OPENINFERENCE_SPAN_KIND, "CHAIN"))
            kind = _OI_KIND_MAP.get(oi_kind, "chain")

            model = str(attrs.pop(OISpanAttributes.LLM_MODEL_NAME, ""))
            provider = str(attrs.pop(OISpanAttributes.LLM_PROVIDER, ""))
            input_tokens = int(attrs.pop(OISpanAttributes.LLM_TOKEN_COUNT_PROMPT, 0))
            output_tokens = int(attrs.pop(OISpanAttributes.LLM_TOKEN_COUNT_COMPLETION, 0))
            total_tokens = int(
                attrs.pop(OISpanAttributes.LLM_TOKEN_COUNT_TOTAL, input_tokens + output_tokens)
            )
            input_val = str(attrs.pop(OISpanAttributes.INPUT_VALUE, ""))
            output_val = str(attrs.pop(OISpanAttributes.OUTPUT_VALUE, ""))
            tool_name = str(attrs.pop(OISpanAttributes.TOOL_NAME, ""))
            tags_raw = attrs.pop(OISpanAttributes.TAG_TAGS, None)
            tags: list[str] = list(tags_raw) if isinstance(tags_raw, (list, tuple)) else []

            # Remove exception attrs from remaining (already captured in error_message)
            attrs.pop("exception.type", None)
            attrs.pop("exception.message", None)
            attrs.pop(OISpanAttributes.SESSION_ID, None)

            duration_ms = (self._end_time - self._start_time) * 1000
            status = "error" if self._status_code == StatusCode.ERROR else "ok"
            ts = datetime.fromtimestamp(self._start_time, tz=timezone.utc).isoformat()

            # Remaining attributes as JSON
            remaining = {}
            for k, v in attrs.items():
                remaining[k] = v
            if tool_name:
                remaining["tool_name"] = tool_name
            if provider:
                remaining["provider"] = provider
            if tags:
                remaining["tags"] = tags

            span_dict: dict[str, Any] = {
                "trace_id": _hex_to_uuid(format(self._span_context.trace_id, "032x")),
                "span_id": _hex_to_uuid(format(self._span_context.span_id, "016x")),
                "parent_span_id": (
                    _hex_to_uuid(format(self._parent_span_id, "016x"))
                    if self._parent_span_id
                    else None
                ),
                "name": self._name,
                "kind": kind,
                "timestamp": ts,
                "duration_ms": round(duration_ms, 2),
                "status": status,
                "model": model,
                "input_tokens": input_tokens,
                "output_tokens": output_tokens,
                "total_tokens": total_tokens,
                "input": input_val,
                "output": output_val,
                "attributes": _safe_json(remaining) if remaining else "{}",
                "error_message": self._status_description if status == "error" else None,
                "session_id": self._session_id,
                "agent_id": self._agent_id,
            }

            self._transport.enqueue(span_dict)
        except Exception:
            logger.debug("staso: failed to end span", exc_info=True)


class Tracer(TracerBase):
    """Tracer that creates Span instances."""

    def __init__(self, provider: TracerProvider) -> None:
        self._provider = provider

    def start_span(
        self,
        name: str,
        context: otel_context.Context | None = None,
        kind: SpanKind = SpanKind.INTERNAL,
        attributes: trace.types.Attributes = None,
        links: Sequence[Link] | None = None,
        start_time: int | None = None,
        record_exception: bool = True,
        set_status_on_exception: bool = True,
    ) -> SpanBase:
        try:
            parent_ctx = context or otel_context.get_current()
            parent_span = trace.get_current_span(parent_ctx)
            parent_sc = parent_span.get_span_context() if parent_span else None

            # Determine trace_id: inherit from parent or generate new
            if parent_sc and parent_sc.is_valid:
                trace_id = parent_sc.trace_id
                parent_span_id = parent_sc.span_id
            else:
                trace_id = random.getrandbits(128)
                parent_span_id = None

            span_id = random.getrandbits(64)

            span_context = SpanContext(
                trace_id=trace_id,
                span_id=span_id,
                is_remote=False,
                trace_flags=TraceFlags(TraceFlags.SAMPLED),
            )

            sid = session_id_var.get("")

            return Span(
                name=name,
                span_context=span_context,
                parent_span_id=parent_span_id,
                transport=self._provider._transport,
                agent_id=self._provider._agent_id,
                session_id=sid,
                attributes=dict(attributes) if attributes else None,
            )
        except Exception:
            logger.debug("staso: failed to start span", exc_info=True)
            # Return a non-recording span so the caller never crashes
            return NonRecordingSpan(
                SpanContext(
                    trace_id=0,
                    span_id=0,
                    is_remote=False,
                    trace_flags=TraceFlags(TraceFlags.DEFAULT),
                )
            )

    @contextlib.contextmanager
    def start_as_current_span(
        self,
        name: str,
        context: otel_context.Context | None = None,
        kind: SpanKind = SpanKind.INTERNAL,
        attributes: trace.types.Attributes = None,
        links: Sequence[Link] | None = None,
        start_time: int | None = None,
        record_exception: bool = True,
        set_status_on_exception: bool = True,
        end_on_exit: bool = True,
    ) -> Iterator[SpanBase]:
        span = self.start_span(
            name,
            context=context,
            kind=kind,
            attributes=attributes,
            links=links,
            start_time=start_time,
            record_exception=record_exception,
            set_status_on_exception=set_status_on_exception,
        )
        token = otel_context.attach(trace.set_span_in_context(span, context))
        try:
            yield span
        except Exception as exc:
            if record_exception:
                span.record_exception(exc)
            if set_status_on_exception:
                span.set_status(StatusCode.ERROR, str(exc))
            raise
        finally:
            if end_on_exit:
                span.end()
            otel_context.detach(token)


class TracerProvider(TracerProviderBase):
    """TracerProvider that creates Tracer instances backed by BatchTransport."""

    def __init__(
        self,
        transport: BatchTransport | None,
        agent_id: str,
        enabled: bool = True,
    ) -> None:
        self._transport = transport
        self._agent_id = agent_id
        self._enabled = enabled

    def get_tracer(
        self,
        instrumenting_module_name: str,
        instrumenting_library_version: str | None = None,
        schema_url: str | None = None,
        attributes: dict[str, AttributeValue] | None = None,
    ) -> Tracer:
        return Tracer(self)

    def shutdown(self) -> None:
        if self._transport:
            self._transport.shutdown()
