from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional, Sequence

from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import Event, ReadableSpan
from opentelemetry.sdk.util.instrumentation import InstrumentationScope
from opentelemetry.trace import (
    Link,
    SpanContext,
    SpanKind,
    Status,
    StatusCode,
    TraceFlags,
    TraceState,
)

from .model import native_from_any_value
from .otel_compat import (
    AggregationTemporality,
    Gauge,
    Histogram,
    HistogramDataPoint,
    Metric,
    MetricsData,
    NumberDataPoint,
    ResourceMetrics,
    ScopeMetrics,
    Sum,
)
from .otlp_endpoints import signal_endpoint

LOGGER = logging.getLogger(__name__)


def _attributes_to_dict(attributes: Optional[List[Dict[str, Any]]]) -> Dict[str, Any]:
    if not attributes:
        return {}
    out: Dict[str, Any] = {}
    for entry in attributes:
        key = entry.get("key")
        if key is None:
            continue
        out[key] = native_from_any_value(entry.get("value", {}))
    return out


def _point_value(point: Dict[str, Any]) -> Optional[int | float]:
    value = point.get("value")
    if isinstance(value, tuple):
        _, value_body = value
        value = value_body.get("value")
    elif isinstance(value, dict):
        value = value.get("value")
    return value


def _scope_from_record(scope: Dict[str, Any]) -> InstrumentationScope:
    return InstrumentationScope(
        name=scope.get("name") or "highlighter.telemetry",
        version=scope.get("version"),
        attributes=_attributes_to_dict(scope.get("attributes")),
    )


def _span_kind_from_str(kind: str) -> SpanKind:
    try:
        return SpanKind[kind.upper()]
    except KeyError:
        return SpanKind.INTERNAL


def _status_from_record(status: Optional[Dict[str, Any]]) -> Status:
    if not status:
        return Status(StatusCode.UNSET)
    code = (status.get("code") or "UNSET").upper()
    description = status.get("message") or status.get("description")
    code_enum = StatusCode.UNSET
    if code == "OK":
        code_enum = StatusCode.OK
    elif code == "ERROR":
        code_enum = StatusCode.ERROR
    return Status(code_enum, description=description)


def _event_from_record(event: Dict[str, Any]) -> Event:
    return Event(
        name=event.get("name") or "",
        attributes=_attributes_to_dict(event.get("attributes")),
        timestamp=event.get("time_unix_nano"),
    )


def _link_from_record(link: Dict[str, Any]) -> Optional[Link]:
    trace_id = link.get("trace_id")
    span_id = link.get("span_id")
    if not trace_id or not span_id:
        return None
    return Link(
        SpanContext(
            trace_id=int.from_bytes(trace_id, "big"),
            span_id=int.from_bytes(span_id, "big"),
            is_remote=False,
            trace_flags=TraceFlags(TraceFlags.SAMPLED),
            trace_state=TraceState.from_header([]),
        ),
        attributes=_attributes_to_dict(link.get("attributes")),
    )


class OtelExporter:
    def __init__(
        self,
        endpoint: str,
        headers: Optional[Dict[str, str]],
        protocol: str,
        resource_attributes: Dict[str, Any],
        scope: Dict[str, Any],
    ) -> None:
        self._resource = Resource(attributes=resource_attributes)
        self._scope = _scope_from_record(scope)
        self._span_exporter = None
        self._metric_exporter = None
        self._log_exporter = None

        if protocol == "grpc":
            self._setup_grpc_exporters(endpoint, headers)
        else:
            self._setup_http_exporters(endpoint, headers)

    def _setup_http_exporters(self, endpoint: str, headers: Optional[Dict[str, str]]) -> None:
        from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
        from opentelemetry.exporter.otlp.proto.http.metric_exporter import (
            OTLPMetricExporter,
        )
        from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
            OTLPSpanExporter,
        )

        trace_endpoint = signal_endpoint(endpoint, "traces")
        metric_endpoint = signal_endpoint(endpoint, "metrics")
        log_endpoint = signal_endpoint(endpoint, "logs")

        base_headers = dict(headers or {})
        for key in list(base_headers.keys()):
            if key.lower() == "stream-name":
                base_headers.pop(key)
                break

        trace_headers = dict(base_headers)
        trace_headers.setdefault("stream-name", "highlighter_traces")
        log_headers = dict(base_headers)
        log_headers.setdefault("stream-name", "highlighter_logs")
        metric_headers = dict(base_headers)

        self._span_exporter = OTLPSpanExporter(endpoint=trace_endpoint, headers=trace_headers)
        self._metric_exporter = OTLPMetricExporter(endpoint=metric_endpoint, headers=metric_headers)
        self._log_exporter = OTLPLogExporter(endpoint=log_endpoint, headers=log_headers)

    def _setup_grpc_exporters(self, endpoint: str, headers: Optional[Dict[str, str]]) -> None:
        from opentelemetry.exporter.otlp.proto.grpc._log_exporter import OTLPLogExporter
        from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import (
            OTLPMetricExporter,
        )
        from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
            OTLPSpanExporter,
        )

        self._span_exporter = OTLPSpanExporter(endpoint=endpoint, headers=headers)
        self._metric_exporter = OTLPMetricExporter(endpoint=endpoint, headers=headers)
        self._log_exporter = OTLPLogExporter(endpoint=endpoint, headers=headers)

    def export_spans(self, spans: Sequence[Dict[str, Any]]) -> None:
        if not spans or not self._span_exporter:
            return
        readable_spans: List[ReadableSpan] = []
        for span in spans:
            trace_id = span.get("trace_id")
            span_id = span.get("span_id")
            if not trace_id or not span_id:
                continue
            parent_span_id = span.get("parent_span_id")
            parent_context = None
            if parent_span_id:
                parent_context = SpanContext(
                    trace_id=int.from_bytes(trace_id, "big"),
                    span_id=int.from_bytes(parent_span_id, "big"),
                    is_remote=False,
                    trace_flags=TraceFlags(TraceFlags.SAMPLED),
                    trace_state=TraceState.from_header([]),
                )
            context = SpanContext(
                trace_id=int.from_bytes(trace_id, "big"),
                span_id=int.from_bytes(span_id, "big"),
                is_remote=False,
                trace_flags=TraceFlags(TraceFlags.SAMPLED),
                trace_state=TraceState.from_header([]),
            )
            readable_spans.append(
                ReadableSpan(
                    name=span.get("name") or "",
                    context=context,
                    parent=parent_context,
                    resource=self._resource,
                    attributes=_attributes_to_dict(span.get("attributes")),
                    events=[_event_from_record(event) for event in span.get("events", [])],
                    links=[
                        link for link in (_link_from_record(link) for link in span.get("links", [])) if link
                    ],
                    kind=_span_kind_from_str(span.get("kind") or "INTERNAL"),
                    status=_status_from_record(span.get("status")),
                    start_time=span.get("start_time_unix_nano"),
                    end_time=span.get("end_time_unix_nano"),
                    instrumentation_scope=self._scope,
                )
            )
        if readable_spans:
            self._span_exporter.export(readable_spans)

    def export_metrics(self, metrics: Sequence[Dict[str, Any]]) -> None:
        if not metrics or not self._metric_exporter:
            return
        otel_metrics: List[Metric] = []
        for metric in metrics:
            data = metric.get("data")
            if not data:
                continue
            if isinstance(data, tuple):
                data_type, data_body = data
            else:
                data_type, data_body = data.get("type"), data
            metric_name = metric.get("name") or ""
            if metric_name and not metric_name.startswith("hl_"):
                metric_name = f"hl_{metric_name}"
            points = []
            if data_type == "Gauge":
                for point in data_body.get("data_points", []):
                    value = _point_value(point)
                    if value is None:
                        continue
                    points.append(
                        NumberDataPoint(
                            attributes=_attributes_to_dict(point.get("attributes")),
                            start_time_unix_nano=point.get("start_time_unix_nano") or 0,
                            time_unix_nano=point.get("time_unix_nano") or 0,
                            value=value,
                        )
                    )
                if not points:
                    continue
                otel_metrics.append(
                    Metric(
                        name=metric_name,
                        description=metric.get("description"),
                        unit=metric.get("unit"),
                        data=Gauge(points),
                    )
                )
                continue

            if data_type == "Sum":
                for point in data_body.get("data_points", []):
                    value = _point_value(point)
                    if value is None:
                        continue
                    points.append(
                        NumberDataPoint(
                            attributes=_attributes_to_dict(point.get("attributes")),
                            start_time_unix_nano=point.get("start_time_unix_nano") or 0,
                            time_unix_nano=point.get("time_unix_nano") or 0,
                            value=value,
                        )
                    )
                if not points:
                    continue
                temporality = AggregationTemporality[
                    (data_body.get("aggregation_temporality") or "CUMULATIVE").upper()
                ]
                is_monotonic = bool(data_body.get("is_monotonic"))
                otel_metrics.append(
                    Metric(
                        name=metric_name,
                        description=metric.get("description"),
                        unit=metric.get("unit"),
                        data=Sum(points, temporality, is_monotonic),
                    )
                )
                continue

            if data_type == "Histogram":
                for point in data_body.get("data_points", []):
                    count = point.get("count")
                    sum_value = point.get("sum")
                    bucket_counts = point.get("bucket_counts")
                    explicit_bounds = point.get("explicit_bounds")
                    if count is None or sum_value is None or bucket_counts is None or explicit_bounds is None:
                        continue
                    points.append(
                        HistogramDataPoint(
                            attributes=_attributes_to_dict(point.get("attributes")),
                            start_time_unix_nano=point.get("start_time_unix_nano") or 0,
                            time_unix_nano=point.get("time_unix_nano") or 0,
                            count=int(count),
                            sum=float(sum_value),
                            bucket_counts=[int(v) for v in bucket_counts],
                            explicit_bounds=[float(v) for v in explicit_bounds],
                            min=float(point.get("min", sum_value)),
                            max=float(point.get("max", sum_value)),
                        )
                    )
                if not points:
                    continue
                temporality = AggregationTemporality[
                    (data_body.get("aggregation_temporality") or "DELTA").upper()
                ]
                otel_metrics.append(
                    Metric(
                        name=metric_name,
                        description=metric.get("description"),
                        unit=metric.get("unit"),
                        data=Histogram(points, temporality),
                    )
                )
                continue
        if not otel_metrics:
            return
        metrics_data = MetricsData(
            [
                ResourceMetrics(
                    resource=self._resource,
                    scope_metrics=[ScopeMetrics(scope=self._scope, metrics=otel_metrics, schema_url="")],
                    schema_url="",
                )
            ]
        )
        self._metric_exporter.export(metrics_data)

    def export_logs(self, logs: Sequence[Dict[str, Any]]) -> None:
        if not logs or not self._log_exporter:
            return
        from opentelemetry._logs import SeverityNumber
        from opentelemetry.sdk._logs import LogData, LogRecord

        log_data: List[LogData] = []
        for record in logs:
            severity_number = record.get("severity_number")
            severity_enum = None
            if isinstance(severity_number, int):
                try:
                    severity_enum = SeverityNumber(severity_number)
                except ValueError:
                    severity_enum = None
            log_data.append(
                LogData(
                    log_record=LogRecord(
                        timestamp=record.get("time_unix_nano"),
                        observed_timestamp=record.get("observed_time_unix_nano"),
                        trace_id=(
                            int.from_bytes(record.get("trace_id"), "big") if record.get("trace_id") else None
                        ),
                        span_id=(
                            int.from_bytes(record.get("span_id"), "big") if record.get("span_id") else None
                        ),
                        trace_flags=TraceFlags(TraceFlags.SAMPLED),
                        severity_text=record.get("severity_text"),
                        severity_number=severity_enum,
                        body=native_from_any_value(record.get("body", {})),
                        resource=self._resource,
                        attributes=_attributes_to_dict(record.get("attributes")),
                    ),
                    instrumentation_scope=self._scope,
                )
            )
        if log_data:
            self._log_exporter.export(log_data)

    def shutdown(self) -> None:
        for exporter in (self._span_exporter, self._metric_exporter, self._log_exporter):
            if exporter:
                try:
                    exporter.shutdown()
                except Exception:
                    LOGGER.exception("Failed to shutdown OTLP exporter; continuing shutdown")
