"""OpenTelemetry setup + trace capture (CONTRACT §2 ``traces``, §4 ``flow.span``).

Two responsibilities:
  1. :func:`setup_telemetry` configures an OTEL ``TracerProvider`` so the Strands
     agent loop emits spans. It is a **no-op when ``OTEL_TRACES_EXPORTER=none``**
     (or when the OTEL SDK is not installed).
  2. :class:`TraceCollector` captures spans into the ``traces`` table shape and
     emits ``flow.span`` SSE events. It works with or without the OTEL SDK: the
     deterministic pipeline records its own flow spans through :meth:`span`.

Keeping capture independent of the SDK means traces are produced even in the
offline/agent-less path.
"""

from __future__ import annotations

import logging
import time
import uuid
from contextlib import contextmanager
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Callable, Iterator, Optional

logger = logging.getLogger(__name__)

EmitFn = Callable[[str, dict[str, Any]], None]


@dataclass
class TraceRow:
    """One captured span, shaped for the ``traces`` table."""

    span_id: str
    name: str
    attributes_json: dict[str, Any]
    started_at: datetime
    duration_ms: int
    parent_span_id: Optional[str] = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "span_id": self.span_id,
            "parent_span_id": self.parent_span_id,
            "name": self.name,
            "attributes_json": self.attributes_json,
            "started_at": self.started_at.isoformat(sep=" ", timespec="microseconds"),
            "duration_ms": self.duration_ms,
        }


@dataclass
class TraceCollector:
    """Collects spans for one run and emits ``flow.span`` events."""

    emit: Optional[EmitFn] = None
    rows: list[TraceRow] = field(default_factory=list)

    @contextmanager
    def span(
        self, name: str, attributes: Optional[dict[str, Any]] = None,
        parent_span_id: Optional[str] = None,
    ) -> Iterator[str]:
        """Time a block as a span; records a :class:`TraceRow` and emits ``flow.span``."""
        span_id = uuid.uuid4().hex[:16]
        started_at = datetime.now(timezone.utc)
        start = time.perf_counter()
        try:
            yield span_id
        finally:
            duration_ms = int((time.perf_counter() - start) * 1000)
            row = TraceRow(
                span_id=span_id,
                name=name,
                attributes_json=attributes or {},
                started_at=started_at,
                duration_ms=duration_ms,
                parent_span_id=parent_span_id,
            )
            self.rows.append(row)
            if self.emit is not None:
                self.emit(
                    "flow.span",
                    {
                        "name": name,
                        "attributes": row.attributes_json,
                        "duration_ms": duration_ms,
                    },
                )

    def record(
        self, name: str, attributes: dict[str, Any], duration_ms: int,
        parent_span_id: Optional[str] = None,
    ) -> None:
        """Record a pre-measured span (e.g. forwarded from the OTEL SDK)."""
        row = TraceRow(
            span_id=uuid.uuid4().hex[:16],
            name=name,
            attributes_json=attributes,
            started_at=datetime.now(timezone.utc),
            duration_ms=duration_ms,
            parent_span_id=parent_span_id,
        )
        self.rows.append(row)
        if self.emit is not None:
            self.emit("flow.span", {"name": name, "attributes": attributes,
                                    "duration_ms": duration_ms})

    def as_trace_rows(self) -> list[dict[str, Any]]:
        return [r.to_dict() for r in self.rows]


def setup_telemetry(settings: "Any") -> bool:
    """Configure the OTEL tracer provider for the agent loop.

    Returns True if a real exporter was wired, False for the no-op path
    (``OTEL_TRACES_EXPORTER=none`` or the SDK is unavailable).
    """
    exporter = (settings.otel_traces_exporter or "none").lower()
    if exporter == "none":
        logger.info("OTEL disabled (OTEL_TRACES_EXPORTER=none); telemetry is a no-op.")
        return False
    try:
        from opentelemetry import trace
        from opentelemetry.sdk.resources import Resource
        from opentelemetry.sdk.trace import TracerProvider
        from opentelemetry.sdk.trace.export import BatchSpanProcessor

        resource = Resource.create({"service.name": settings.otel_service_name})
        provider = TracerProvider(resource=resource)

        if exporter == "otlp":
            from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
                OTLPSpanExporter,
            )

            endpoint = settings.otel_exporter_otlp_endpoint
            span_exporter = OTLPSpanExporter(endpoint=endpoint) if endpoint else OTLPSpanExporter()
        else:  # "console" or unknown -> console for visibility
            from opentelemetry.sdk.trace.export import ConsoleSpanExporter

            span_exporter = ConsoleSpanExporter()

        provider.add_span_processor(BatchSpanProcessor(span_exporter))
        trace.set_tracer_provider(provider)
        logger.info("OTEL tracer provider configured with exporter=%s", exporter)
        return True
    except Exception as exc:  # pragma: no cover - needs OTEL SDK
        logger.warning("OTEL setup failed (%s); continuing without telemetry.", exc)
        return False