from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict

from entropic_core.core.runtime_events import RuntimeEvent


def _require_otel() -> None:
    try:
        import opentelemetry  # noqa: F401
    except Exception as e:  # pragma: no cover
        raise ImportError(
            "OpenTelemetry extras not installed. Install with: pip install \"entropic-core[otel]\""
        ) from e


@dataclass
class OpenTelemetryEventSink:
    """Export RuntimeEvents to OpenTelemetry spans (best-effort).

    Usage:
        sink = OpenTelemetryEventSink(service_name="entropic-core")
        cfg = MiddlewareConfig(event_sink=sink, run_id="run_123")

    Notes:
    - This sink configures a local TracerProvider by default (not global).
    - For production, prefer passing your own tracer/provider from your app.
    """

    service_name: str = "entropic-core"
    tracer: Any = None

    def __post_init__(self) -> None:
        _require_otel()

        if self.tracer is not None:
            return

        # Local provider: avoid mutating global state in a library by default.
        from opentelemetry.sdk.resources import Resource
        from opentelemetry.sdk.trace import TracerProvider
        from opentelemetry.sdk.trace.export import (
            BatchSpanProcessor,
            ConsoleSpanExporter,
            SimpleSpanProcessor,
        )
        from opentelemetry.trace import get_tracer

        res = Resource.create({"service.name": self.service_name})
        provider = TracerProvider(resource=res)
        provider.add_span_processor(BatchSpanProcessor(ConsoleSpanExporter()))
        # Also add a simple processor for immediate feedback in local/dev.
        provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter()))

        # Create a tracer bound to this provider without setting the global provider.
        self._provider = provider
        self.tracer = get_tracer("entropic-core", tracer_provider=provider)

        self._spans_by_request: Dict[str, Any] = {}

    def emit(self, event: RuntimeEvent) -> None:
        # Never throw from an exporter.
        try:
            et = str(event.type or "")
            rid = str(event.request_id or "")
            if not rid:
                return

            if et == "request":
                self._on_request(event)
            elif et == "response":
                self._on_response(event)
            elif et == "error":
                self._on_error(event)
        except Exception:
            return

    def _start_span(self, event: RuntimeEvent) -> Any:
        name = "entropic.llm"
        span = self.tracer.start_span(name)
        span.set_attribute("entropic.run_id", event.run_id)
        span.set_attribute("entropic.request_id", event.request_id)
        if event.provider:
            span.set_attribute("llm.provider", event.provider)
        if event.model:
            span.set_attribute("llm.model", event.model)
        return span

    def _on_request(self, event: RuntimeEvent) -> None:
        span = self._start_span(event)
        payload = event.payload or {}
        if "messages_count" in payload:
            span.set_attribute("entropic.messages_count", int(payload["messages_count"] or 0))
        if payload.get("entropy_before") is not None:
            span.set_attribute("entropic.entropy_before", float(payload["entropy_before"]))
        self._spans_by_request[event.request_id] = span

    def _on_response(self, event: RuntimeEvent) -> None:
        span = self._spans_by_request.pop(event.request_id, None)
        if span is None:
            span = self._start_span(event)

        p = event.payload or {}
        if p.get("latency_ms") is not None:
            span.set_attribute("entropic.latency_ms", float(p["latency_ms"]))
        if p.get("entropy_after") is not None:
            span.set_attribute("entropic.entropy_after", float(p["entropy_after"]))
        if p.get("is_refusal") is not None:
            span.set_attribute("entropic.is_refusal", bool(p["is_refusal"]))
        if p.get("tool_calls_count") is not None:
            span.set_attribute("entropic.tool_calls_count", int(p["tool_calls_count"] or 0))
        if p.get("retry_count") is not None:
            span.set_attribute("entropic.retry_count", int(p["retry_count"] or 0))

        usage = p.get("usage")
        if isinstance(usage, dict):
            span.set_attribute("llm.tokens.prompt", int(usage.get("prompt_tokens", 0) or 0))
            span.set_attribute("llm.tokens.completion", int(usage.get("completion_tokens", 0) or 0))
            span.set_attribute("llm.tokens.total", int(usage.get("total_tokens", 0) or 0))

        interventions = p.get("interventions")
        if isinstance(interventions, list):
            span.set_attribute(
                "entropic.interventions", ",".join(str(x) for x in interventions[:32])
            )

        span.end()

    def _on_error(self, event: RuntimeEvent) -> None:
        span = self._spans_by_request.pop(event.request_id, None)
        if span is None:
            span = self._start_span(event)
        p = event.payload or {}
        span.set_attribute("error", True)
        if p.get("error_type"):
            span.set_attribute("exception.type", str(p.get("error_type")))
        if p.get("error"):
            span.set_attribute("exception.message", str(p.get("error"))[:1000])
        span.end()

    def shutdown(self) -> None:
        prov = getattr(self, "_provider", None)
        if prov is not None:
            try:
                prov.shutdown()
            except Exception:
                pass
