"""Span, Tracer, and TracerProvider with direct typed fields."""
from __future__ import annotations

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

from .context import get_current_span, reset_current_span, session_id_var, set_current_span, user_id_var
from .transport import BatchTransport
from .types import SpanKind, SpanStatus

logger = logging.getLogger("staso")


def _hex_to_uuid(hex_str: str) -> str:
    """Convert a hex string to UUID format."""
    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:
    """A span with direct typed fields matching the backend ClickHouse schema."""

    def __init__(
        self,
        name: str,
        trace_id: str,
        span_id: str,
        parent_span_id: str | None,
        transport: BatchTransport | None,
        agent_id: str,
        session_id: str,
        kind: SpanKind = SpanKind.CHAIN,
        user_id: str = "",
    ) -> None:
        self.name = name
        self.trace_id = trace_id
        self.span_id = span_id
        self.parent_span_id = parent_span_id
        self.agent_id = agent_id
        self.session_id = session_id
        self.user_id = user_id

        # Direct fields -> ClickHouse columns
        self.kind = kind
        self.status = SpanStatus.OK
        self.model: str = ""
        self.input_tokens: int = 0
        self.output_tokens: int = 0
        self.total_tokens: int = 0
        self.error_message: str | None = None

        # Rich dicts -> ClickHouse ZSTD JSON blobs
        self.input: dict[str, Any] = {}
        self.output: dict[str, Any] = {}
        self.metadata: dict[str, Any] = {}

        # Internal
        self._transport = transport
        self._start_time = time.time()
        self._end_time: float | None = None
        self._recording = True

    def set_status(self, status: SpanStatus, description: str | None = None) -> None:
        """Set the span status."""
        if self._recording:
            self.status = status
            if description:
                self.error_message = description

    def record_exception(self, exception: BaseException) -> None:
        """Record an exception on this span."""
        if self._recording:
            self.status = SpanStatus.ERROR
            self.error_message = str(exception)

    def annotate(
        self,
        name: str,
        *,
        value: float | str | bool,
        data_type: str = "numeric",
        comment: str | None = None,
    ) -> None:
        """Submit an annotation for this span's trace, scoped to this span."""
        from .annotation import annotate as _annotate

        _annotate(
            trace_id=self.trace_id,
            name=name,
            value=value,
            data_type=data_type,
            span_id=self.span_id,
            comment=comment,
        )

    def end(self) -> None:
        """Serialize and enqueue the span for transport."""
        try:
            if not self._recording:
                return
            self._recording = False
            self._end_time = time.time()

            if self._transport is None:
                return

            duration_ms = (self._end_time - self._start_time) * 1000
            ts = datetime.fromtimestamp(self._start_time, tz=timezone.utc).isoformat()

            span_dict: dict[str, Any] = {
                "trace_id": self.trace_id,
                "span_id": self.span_id,
                "parent_span_id": self.parent_span_id,
                "name": self.name,
                "kind": self.kind.value if isinstance(self.kind, SpanKind) else self.kind,
                "timestamp": ts,
                "duration_ms": round(duration_ms, 2),
                "status": self.status.value if isinstance(self.status, SpanStatus) else self.status,
                "model": self.model,
                "input_tokens": self.input_tokens,
                "output_tokens": self.output_tokens,
                "total_tokens": self.total_tokens,
                "input": _safe_json(self.input) if self.input else "",
                "output": _safe_json(self.output) if self.output else "",
                "attributes": _safe_json(self.metadata) if self.metadata else "{}",
                "error_message": self.error_message,
                "session_id": self.session_id,
                "agent_id": self.agent_id,
                "user_id": self.user_id,
            }

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


# Singleton noop span for fallback when span creation fails
_NOOP_SPAN = Span(
    name="noop",
    trace_id="00000000-0000-0000-0000-000000000000",
    span_id="00000000-0000-0000-0000-000000000000",
    parent_span_id=None,
    transport=None,
    agent_id="",
    session_id="",
)
_NOOP_SPAN._recording = False


class Tracer:
    """Creates Span instances with parent-child context propagation."""

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

    def start_span(self, name: str, kind: SpanKind = SpanKind.CHAIN) -> Span:
        """Create a new span, inheriting trace_id from the current parent."""
        try:
            parent = get_current_span()

            if parent is not None and not parent._recording:
                return _NOOP_SPAN

            if parent and parent._recording:
                trace_id_int = int(parent.trace_id.replace("-", ""), 16)
                parent_span_id = parent.span_id
            else:
                trace_id_int = random.getrandbits(128)
                parent_span_id = None

            span_id_int = random.getrandbits(64)
            trace_id = _hex_to_uuid(format(trace_id_int, "032x"))
            span_id = _hex_to_uuid(format(span_id_int, "016x"))
            sid = session_id_var.get("")
            uid = user_id_var.get("")

            return Span(
                name=name,
                trace_id=trace_id,
                span_id=span_id,
                parent_span_id=parent_span_id,
                transport=self._provider._transport,
                agent_id=self._provider._agent_id,
                session_id=sid,
                kind=kind,
                user_id=uid,
            )
        except Exception:
            logger.debug("staso: failed to start span", exc_info=True)
            return _NOOP_SPAN

    @contextlib.contextmanager
    def start_as_current_span(
        self,
        name: str,
        kind: SpanKind = SpanKind.CHAIN,
    ) -> Iterator[Span]:
        """Create a span, set it as current, yield it, and end it on exit."""
        span = self.start_span(name, kind=kind)
        token = set_current_span(span)
        try:
            yield span
        except Exception as exc:
            span.record_exception(exc)
            raise
        finally:
            span.end()
            reset_current_span(token)


class TracerProvider:
    """Creates Tracer instances backed by a 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) -> Tracer:
        return Tracer(self)

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