"""
OTLP Export — Send Aigie traces to external OpenTelemetry collectors.

Converts Aigie's internal trace/span format to OTel-compatible spans and
exports them via OTLP/HTTP (protobuf or JSON). No grpcio dependency.

Usage:
    from aigie.otlp_export import OTLPExportBridge

    bridge = OTLPExportBridge(
        endpoint="http://localhost:4318",  # OTLP/HTTP endpoint
        service_name="my-agent-service",
    )
    bridge.attach(aigie_client)
    # All Aigie traces will now also be exported to your OTel collector.

Install:
    pip install aigie[otlp]
"""

import logging
from typing import Any, Dict, List

logger = logging.getLogger("aigie.otlp_export")


def _check_otlp_deps() -> bool:
    """Check if OTLP export dependencies are available."""
    try:
        from opentelemetry.sdk.trace import TracerProvider
        from opentelemetry.sdk.trace.export import BatchSpanProcessor

        return True
    except ImportError:
        return False


def _check_otlp_http_exporter() -> bool:
    """Check if OTLP HTTP exporter is available."""
    try:
        from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
            OTLPSpanExporter,
        )

        return True
    except ImportError:
        return False


class OTLPExportBridge:
    """Bridge that taps into Aigie's EventBuffer and exports spans via OTLP/HTTP.

    This is a non-intrusive add-on: it hooks into the existing buffer's
    flush callback to capture spans as they flow through, converts them to
    OTel format, and exports via the standard OTLP HTTP exporter.

    Does NOT require grpcio. Uses opentelemetry-exporter-otlp-proto-http.
    """

    def __init__(
        self,
        endpoint: str = "http://localhost:4318",
        service_name: str = "aigie-agent",
        headers: Dict[str, str] | None = None,
        resource_attributes: Dict[str, str] | None = None,
        export_timeout_ms: int = 10000,
        max_queue_size: int = 2048,
        max_export_batch_size: int = 512,
    ):
        self._endpoint = endpoint
        self._service_name = service_name
        self._headers = headers or {}
        self._resource_attributes = resource_attributes or {}
        self._export_timeout_ms = export_timeout_ms
        self._max_queue_size = max_queue_size
        self._max_export_batch_size = max_export_batch_size

        self._tracer_provider = None
        self._tracer = None
        self._attached = False
        self._stats = {
            "spans_exported": 0,
            "export_errors": 0,
        }

    def attach(self, aigie_client) -> bool:
        """Attach to an Aigie client instance. Returns True if successful."""
        if not _check_otlp_deps():
            logger.warning("OpenTelemetry SDK not installed. Install with: pip install aigie[otlp]")
            return False

        if not _check_otlp_http_exporter():
            logger.warning(
                "OTLP HTTP exporter not installed. "
                "Install with: pip install opentelemetry-exporter-otlp-proto-http"
            )
            return False

        try:
            self._setup_otel()
            self._hook_buffer(aigie_client)
            self._attached = True
            logger.info(f"OTLP export bridge attached — exporting to {self._endpoint}")
            return True
        except Exception as e:
            logger.error(f"Failed to attach OTLP bridge: {e}")
            return False

    def _setup_otel(self) -> None:
        """Initialize OTel TracerProvider with OTLP HTTP exporter."""
        from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
            OTLPSpanExporter,
        )
        from opentelemetry.sdk.resources import Resource
        from opentelemetry.sdk.trace import TracerProvider
        from opentelemetry.sdk.trace.export import BatchSpanProcessor

        resource_attrs = {
            "service.name": self._service_name,
            "telemetry.sdk.name": "aigie",
            "telemetry.sdk.language": "python",
        }
        resource_attrs.update(self._resource_attributes)
        resource = Resource.create(resource_attrs)

        exporter = OTLPSpanExporter(
            endpoint=f"{self._endpoint.rstrip('/')}/v1/traces",
            headers=self._headers,
            timeout=self._export_timeout_ms / 1000,
        )

        self._tracer_provider = TracerProvider(resource=resource)
        self._tracer_provider.add_span_processor(
            BatchSpanProcessor(
                exporter,
                max_queue_size=self._max_queue_size,
                max_export_batch_size=self._max_export_batch_size,
            )
        )
        self._tracer = self._tracer_provider.get_tracer("aigie.otlp_bridge", "0.2.13")

    def _hook_buffer(self, aigie_client) -> None:
        """Hook into the Aigie client's buffer to intercept span events."""
        # Store reference to client for span conversion
        self._aigie = aigie_client

        # Register a post-flush callback if the buffer supports it
        if hasattr(aigie_client, "_buffer") and aigie_client._buffer:
            original_flush = aigie_client._buffer._flush_batch
            bridge = self

            async def _hooked_flush(events):
                # Call original flush first
                result = await original_flush(events)
                # Then export to OTel (non-blocking, best-effort)
                try:
                    bridge._export_events(events)
                except Exception as e:
                    bridge._stats["export_errors"] += 1
                    logger.debug(f"OTLP export error (non-fatal): {e}")
                return result

            aigie_client._buffer._flush_batch = _hooked_flush

    def _export_events(self, events: List[Dict[str, Any]]) -> None:
        """Convert Aigie buffer events to OTel spans and export."""
        if not self._tracer:
            return

        for event in events:
            event_type = event.get("type", "")
            payload = event.get("payload", event)

            # Only export span events
            if "span" not in event_type.lower() and "SPAN" not in event_type:
                continue

            try:
                self._export_single_span(payload)
            except Exception:
                self._stats["export_errors"] += 1

    def _export_single_span(self, payload: Dict[str, Any]) -> None:
        """Convert a single Aigie span payload to an OTel span."""
        from opentelemetry.trace import StatusCode

        span_name = payload.get("name", "aigie.span")
        span_type = payload.get("type", "unknown")

        # Create OTel span
        span = self._tracer.start_span(
            name=f"{span_type}.{span_name}" if span_type != "unknown" else span_name,
        )

        # Set attributes from Aigie span
        attrs = {
            "aigie.span.type": span_type,
            "aigie.trace.id": payload.get("trace_id", ""),
        }

        # Model/provider info
        if payload.get("model"):
            attrs["gen_ai.request.model"] = payload["model"]
        if payload.get("provider"):
            attrs["gen_ai.system"] = payload["provider"]

        # Token usage (following OTel GenAI semantic conventions)
        metadata = payload.get("metadata", {})
        usage = metadata.get("usage", {})
        if usage.get("prompt_tokens"):
            attrs["gen_ai.usage.input_tokens"] = usage["prompt_tokens"]
        if usage.get("completion_tokens"):
            attrs["gen_ai.usage.output_tokens"] = usage["completion_tokens"]

        # Cost
        cost = metadata.get("cost", {})
        if cost.get("total_cost"):
            attrs["aigie.cost.total"] = cost["total_cost"]

        # Duration
        if payload.get("duration_ns"):
            attrs["aigie.duration_ns"] = payload["duration_ns"]

        for k, v in attrs.items():
            if v is not None:
                span.set_attribute(k, v)

        # Status
        error = payload.get("error_message") or payload.get("error")
        if error:
            span.set_status(StatusCode.ERROR, str(error)[:500])
        elif payload.get("status") == "success":
            span.set_status(StatusCode.OK)

        span.end()
        self._stats["spans_exported"] += 1

    def detach(self) -> None:
        """Shut down the OTLP exporter and clean up."""
        if self._tracer_provider:
            try:
                self._tracer_provider.shutdown()
            except Exception as e:
                logger.debug(f"OTLP shutdown error: {e}")
            self._tracer_provider = None
            self._tracer = None
        self._attached = False

    @property
    def is_attached(self) -> bool:
        return self._attached

    def get_stats(self) -> Dict[str, Any]:
        return {
            **self._stats,
            "endpoint": self._endpoint,
            "service_name": self._service_name,
            "attached": self._attached,
        }
