from __future__ import annotations

from contextlib import asynccontextmanager, contextmanager
from typing import Any

from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter


class NexusTracer:
    """OpenTelemetry-compatible tracer with in-memory fallback exporter."""

    def __init__(self, service_name: str = "nexus", exporter_endpoint: str | None = None) -> None:
        self._in_memory = InMemorySpanExporter()
        provider = TracerProvider()
        if exporter_endpoint:
            provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter(endpoint=exporter_endpoint)))
        else:
            provider.add_span_processor(BatchSpanProcessor(self._in_memory))
        trace.set_tracer_provider(provider)
        self._tracer = trace.get_tracer(service_name)

    @contextmanager
    def span(self, name: str, attributes: dict | None = None):
        """Context manager for synchronous span recording."""
        with self._tracer.start_as_current_span(name) as sp:
            for key, value in (attributes or {}).items():
                sp.set_attribute(key, value)
            try:
                yield sp
            except Exception as exc:
                sp.record_exception(exc)
                raise

    @asynccontextmanager
    async def async_span(self, name: str, attributes: dict | None = None):
        """Async context manager for span recording."""
        with self._tracer.start_as_current_span(name) as sp:
            for key, value in (attributes or {}).items():
                sp.set_attribute(key, value)
            try:
                yield sp
            except Exception as exc:
                sp.record_exception(exc)
                raise

    def get_spans(self) -> list[dict[str, Any]]:
        """Return recorded in-memory spans."""
        out: list[dict[str, Any]] = []
        for sp in self._in_memory.get_finished_spans():
            out.append(
                {
                    "name": sp.name,
                    "start_time": sp.start_time,
                    "end_time": sp.end_time,
                    "attributes": dict(sp.attributes),
                    "status": str(sp.status.status_code),
                    "events": [
                        {"name": ev.name, "timestamp": ev.timestamp, "attributes": dict(ev.attributes)}
                        for ev in sp.events
                    ],
                }
            )
        return out

    def export_jaeger_format(self) -> list[dict[str, Any]]:
        """Convert in-memory spans to a Jaeger-like list representation."""
        return [
            {
                "operationName": s["name"],
                "startTime": s["start_time"],
                "duration": (s["end_time"] - s["start_time"]) if s["end_time"] and s["start_time"] else 0,
                "tags": [{"key": k, "value": v} for k, v in s["attributes"].items()],
                "logs": s["events"],
            }
            for s in self.get_spans()
        ]
