"""
Galtea Span Processor for OpenTelemetry.

This processor collects finished spans and exports them to the Galtea API
as trace records. It batches spans by inference_result_id and sends them
when the trace context is flushed or closed.
"""

import json
import logging
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor
from opentelemetry.trace import StatusCode

from galtea.domain.models.trace import TraceBase
from galtea.infrastructure.telemetry.attributes import (
    GALTEA_ATTR_INFERENCE_RESULT_ID,
    GALTEA_ATTR_TRACE_DESCRIPTION,
    GALTEA_ATTR_TRACE_ERROR,
    GALTEA_ATTR_TRACE_INPUT,
    GALTEA_ATTR_TRACE_METADATA,
    GALTEA_ATTR_TRACE_OUTPUT,
    GALTEA_ATTR_TRACE_TYPE,
)


class GalteaSpanProcessor(SpanProcessor):
    """OpenTelemetry span processor that exports spans to Galtea API.

    This processor collects finished spans and converts them to Galtea TraceBase
    objects. Spans are batched by inference_result_id and can be exported either:
    - Immediately on span end (real-time mode)
    - On explicit flush (batch mode, default)

    The processor uses the galtea.* namespace attributes to extract trace data.
    """

    def __init__(
        self,
        http_client: Any,
        batch_mode: bool = True,
        max_batch_size: int = 100,
    ):
        """Initialize the Galtea span processor.

        Args:
                http_client: Galtea HTTP client for API requests.
                batch_mode: If True, batch spans until flush. If False, export immediately.
                max_batch_size: Maximum spans to batch before auto-flush.
        """
        self._client = http_client
        self._batch_mode = batch_mode
        self._max_batch_size = max_batch_size
        self._pending_spans: Dict[str, List[ReadableSpan]] = {}
        self._logger = logging.getLogger(__name__)

    def on_start(self, span: ReadableSpan, parent_context: Optional[Any] = None) -> None:
        """Called when a span is started. No-op for this processor."""
        pass

    def on_end(self, span: ReadableSpan) -> None:
        """Called when a span is ended.

        Collects the span for later export or exports immediately based on mode.

        Args:
                span: The finished span.
        """
        inference_result_id = self._get_inference_result_id(span)
        if not inference_result_id:
            self._logger.debug(f"Span {span.name} has no inference_result_id, skipping")
            return

        if self._batch_mode:
            if inference_result_id not in self._pending_spans:
                self._pending_spans[inference_result_id] = []
            self._pending_spans[inference_result_id].append(span)

            if len(self._pending_spans[inference_result_id]) >= self._max_batch_size:
                self._export_batch(inference_result_id)
        else:
            self._export_span(span, inference_result_id)

    def shutdown(self) -> None:
        """Shutdown the processor, exporting any pending spans."""
        self.force_flush()

    def force_flush(self, timeout_millis: Optional[int] = None) -> bool:
        """Force flush all pending spans to Galtea API.

        Args:
                timeout_millis: Optional timeout (not currently used).

        Returns:
                True if flush was successful.
        """
        for inference_result_id in list(self._pending_spans.keys()):
            self._export_batch(inference_result_id)
        return True

    def flush_inference(self, inference_result_id: str) -> None:
        """Flush spans for a specific inference result.

        Args:
                inference_result_id: The inference result ID to flush.
        """
        if inference_result_id in self._pending_spans:
            self._export_batch(inference_result_id)

    def _export_batch(self, inference_result_id: str) -> None:
        """Export a batch of spans for an inference result.

        Args:
                inference_result_id: The inference result ID.
        """
        spans = self._pending_spans.pop(inference_result_id, [])
        if not spans:
            return

        try:
            traces = [self._span_to_trace(span, inference_result_id) for span in spans]
            traces = [t for t in traces if t is not None]

            if traces:
                self._client.post(
                    "traces/batch",
                    json={"traces": [t.model_dump(by_alias=True, exclude_none=True) for t in traces]},
                )
                self._logger.debug(f"Exported {len(traces)} traces for inference {inference_result_id}")
        except Exception as e:
            self._logger.error(f"Failed to export traces: {e}")

    def _export_span(self, span: ReadableSpan, inference_result_id: str) -> None:
        """Export a single span immediately.

        Args:
                span: The span to export.
                inference_result_id: The inference result ID.
        """
        try:
            trace = self._span_to_trace(span, inference_result_id)
            if trace:
                self._client.post(
                    "traces",
                    json=trace.model_dump(by_alias=True, exclude_none=True),
                )
                self._logger.debug(f"Exported trace {trace.name} for inference {inference_result_id}")
        except Exception as e:
            self._logger.error(f"Failed to export trace: {e}")

    def _span_to_trace(self, span: ReadableSpan, inference_result_id: str) -> Optional[TraceBase]:
        """Convert an OpenTelemetry span to a Galtea TraceBase.

        Args:
                span: The OpenTelemetry span.
                inference_result_id: The inference result ID.

        Returns:
                TraceBase object or None if conversion fails.
        """
        try:
            attributes = dict(span.attributes) if span.attributes else {}

            # Extract timing
            start_time = None
            end_time = None
            latency_ms = None

            if span.start_time:
                start_time = datetime.fromtimestamp(span.start_time / 1e9, tz=timezone.utc).isoformat()

            if span.end_time:
                end_time = datetime.fromtimestamp(span.end_time / 1e9, tz=timezone.utc).isoformat()

            if span.start_time and span.end_time:
                latency_ms = (span.end_time - span.start_time) / 1e6

            # Extract Galtea-specific attributes
            trace_type = attributes.get(GALTEA_ATTR_TRACE_TYPE)
            description = attributes.get(GALTEA_ATTR_TRACE_DESCRIPTION)
            input_data = self._parse_json_attribute(attributes.get(GALTEA_ATTR_TRACE_INPUT))
            output_data = self._parse_json_attribute(attributes.get(GALTEA_ATTR_TRACE_OUTPUT))
            error = attributes.get(GALTEA_ATTR_TRACE_ERROR)
            metadata = self._parse_json_attribute(attributes.get(GALTEA_ATTR_TRACE_METADATA))

            # Extract parent trace ID from OpenTelemetry span parent context
            parent_trace_id = None
            if span.parent and span.parent.span_id:
                # Convert parent span ID to Galtea trace ID format
                parent_trace_id = f"trace_{span.parent.span_id:016x}"

            # Handle error from span status
            if span.status and span.status.status_code == StatusCode.ERROR:
                error = error or span.status.description or "Unknown error"

            # Generate trace ID from span ID
            trace_id = f"trace_{span.context.span_id:016x}"

            return TraceBase(
                id=trace_id,
                inference_result_id=inference_result_id,
                name=span.name,
                description=description,
                type=trace_type,
                input_data=input_data,
                output_data=output_data,
                error=error,
                latency_ms=latency_ms,
                metadata=metadata,
                parent_trace_id=parent_trace_id,
                start_time=start_time,
                end_time=end_time,
            )
        except Exception as e:
            self._logger.error(f"Failed to convert span to trace: {e}")
            return None

    def _get_inference_result_id(self, span: ReadableSpan) -> Optional[str]:
        """Extract inference_result_id from span attributes.

        Args:
                span: The span.

        Returns:
                The inference result ID or None.
        """
        if span.attributes:
            return span.attributes.get(GALTEA_ATTR_INFERENCE_RESULT_ID)
        return None

    def _parse_json_attribute(self, value: Any) -> Optional[Dict[str, Any]]:
        """Parse a JSON string attribute to a dictionary.

        Args:
                value: The attribute value (string or dict).

        Returns:
                Parsed dictionary or None.
        """
        if value is None:
            return None
        if isinstance(value, dict):
            return value
        if isinstance(value, str):
            try:
                return json.loads(value)
            except json.JSONDecodeError:
                return {"value": value}
        return {"value": str(value)}
