"""Observability primitives for Logos ADK."""
from __future__ import annotations

import asyncio
import base64
import gzip
import json
import re
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import Any, Awaitable, Callable

import httpx

from logos_adk.version import __version__

@dataclass
class Span:
    """A traced operation."""

    name: str
    trace_id: str
    span_id: str
    parent_span_id: str | None
    started_at: datetime
    ended_at: datetime
    duration_ms: float
    agent_name: str | None = None
    session_id: str | None = None
    status: str = "ok"
    attributes: dict[str, Any] = field(default_factory=dict)


@dataclass
class Metric:
    """A recorded metric."""

    name: str
    value: float
    unit: str | None = None
    attributes: dict[str, Any] = field(default_factory=dict)


class Exporter(ABC):
    """Base exporter interface."""

    @abstractmethod
    async def on_span(self, span: Span) -> None:
        ...

    @abstractmethod
    async def on_metric(self, metric: Metric) -> None:
        ...

    @abstractmethod
    async def flush(self) -> None:
        ...


class ConsoleExporter(Exporter):
    """Simple stdout exporter for spans and metrics."""

    async def on_span(self, span: Span) -> None:
        print(
            f"[span] {span.name} status={span.status} "
            f"duration_ms={span.duration_ms:.2f} trace_id={span.trace_id}"
        )

    async def on_metric(self, metric: Metric) -> None:
        unit = f" {metric.unit}" if metric.unit else ""
        print(f"[metric] {metric.name}={metric.value}{unit}")

    async def flush(self) -> None:
        return


class InMemoryExporter(Exporter):
    """Exporter that retains spans and metrics in memory."""

    def __init__(self) -> None:
        self.spans: list[Span] = []
        self.metrics: list[Metric] = []
        self.flush_count = 0

    async def on_span(self, span: Span) -> None:
        self.spans.append(span)

    async def on_metric(self, metric: Metric) -> None:
        self.metrics.append(metric)

    async def flush(self) -> None:
        self.flush_count += 1


class OpenTelemetryExporter(Exporter):
    """OTLP-like HTTP exporter for spans and metrics."""

    def __init__(
        self,
        endpoint: str,
        *,
        timeout: float = 5.0,
        transport: httpx.AsyncBaseTransport | None = None,
        batch_size: int = 100,
        max_retries: int = 3,
        backoff_seconds: float = 0.25,
        headers: dict[str, str] | None = None,
        bearer_token: str | None = None,
        compression: str = "none",
        sleeper: Callable[[float], Awaitable[None]] | None = None,
    ) -> None:
        if timeout <= 0:
            raise ValueError("OpenTelemetryExporter timeout must be positive")
        if batch_size <= 0:
            raise ValueError("OpenTelemetryExporter batch_size must be positive")
        if max_retries < 0:
            raise ValueError("OpenTelemetryExporter max_retries must be non-negative")
        if backoff_seconds < 0:
            raise ValueError("OpenTelemetryExporter backoff_seconds must be non-negative")
        if compression not in {"none", "gzip"}:
            raise ValueError("OpenTelemetryExporter compression must be 'none' or 'gzip'")

        self.endpoint = endpoint
        self.timeout = timeout
        self._transport = transport
        self.batch_size = batch_size
        self.max_retries = max_retries
        self.backoff_seconds = backoff_seconds
        self.headers = dict(headers or {})
        self.bearer_token = bearer_token
        self.compression = compression
        self._sleeper = sleeper or asyncio.sleep
        self._client: httpx.AsyncClient | None = None
        self.spans: list[Span] = []
        self.metrics: list[Metric] = []

    async def on_span(self, span: Span) -> None:
        self.spans.append(span)

    async def on_metric(self, metric: Metric) -> None:
        self.metrics.append(metric)

    def _derive_signal_endpoint(self, signal: str) -> str:
        if signal == "traces" and self.endpoint.endswith("/v1/metrics"):
            return f"{self.endpoint[:-len('/v1/metrics')]}/v1/traces"
        if signal == "metrics" and self.endpoint.endswith("/v1/traces"):
            return f"{self.endpoint[:-len('/v1/traces')]}/v1/metrics"
        return self.endpoint

    def _normalize_hex_id(self, value: str, length: int) -> str:
        hex_only = re.sub(r"[^0-9a-fA-F]", "", value).lower()
        if len(hex_only) >= length:
            return hex_only[-length:]
        return hex_only.rjust(length, "0")

    def _to_unix_nano(self, value: datetime) -> str:
        timestamp = value if value.tzinfo is not None else value.replace(tzinfo=UTC)
        return str(int(timestamp.timestamp() * 1_000_000_000))

    def _attribute_value(self, value: Any) -> dict[str, Any]:
        if isinstance(value, bool):
            return {"boolValue": value}
        if isinstance(value, int) and not isinstance(value, bool):
            return {"intValue": str(value)}
        if isinstance(value, float):
            return {"doubleValue": value}
        return {"stringValue": str(value)}

    def _attributes_to_otel(self, attributes: dict[str, Any]) -> list[dict[str, Any]]:
        return [
            {"key": key, "value": self._attribute_value(value)}
            for key, value in sorted(attributes.items())
        ]

    def _base_resource_attributes(self) -> list[dict[str, Any]]:
        return self._attributes_to_otel({"service.name": "logos_adk", "telemetry.sdk.language": "python"})

    def _serialize_span(self, span: Span) -> dict[str, Any]:
        attributes = dict(span.attributes)
        if span.agent_name:
            attributes["agent.name"] = span.agent_name
        if span.session_id:
            attributes["session.id"] = span.session_id
        return {
            "traceId": self._normalize_hex_id(span.trace_id, 32),
            "spanId": self._normalize_hex_id(span.span_id, 16),
            "parentSpanId": self._normalize_hex_id(span.parent_span_id, 16) if span.parent_span_id else None,
            "name": span.name,
            "startTimeUnixNano": self._to_unix_nano(span.started_at),
            "endTimeUnixNano": self._to_unix_nano(span.ended_at),
            "attributes": self._attributes_to_otel(attributes),
            "status": {
                "code": 2 if span.status == "error" else 1 if span.status == "ok" else 0,
                "message": span.status,
            },
        }

    def _serialize_metric(self, metric: Metric) -> dict[str, Any]:
        return {
            "name": metric.name,
            "unit": metric.unit or "",
            "gauge": {
                "dataPoints": [
                    {
                        "asDouble": float(metric.value),
                        "timeUnixNano": self._to_unix_nano(datetime.now(UTC)),
                        "attributes": self._attributes_to_otel(metric.attributes),
                    }
                ]
            },
        }

    def _build_traces_payload(self, spans: list[Span]) -> dict[str, Any]:
        return {
            "resourceSpans": [
                {
                    "resource": {"attributes": self._base_resource_attributes()},
                    "scopeSpans": [
                        {
                            "scope": {"name": "logos_adk.observe", "version": __version__},
                            "spans": [self._serialize_span(span) for span in spans],
                        }
                    ],
                }
            ]
        }

    def _build_metrics_payload(self, metrics: list[Metric]) -> dict[str, Any]:
        return {
            "resourceMetrics": [
                {
                    "resource": {"attributes": self._base_resource_attributes()},
                    "scopeMetrics": [
                        {
                            "scope": {"name": "logos_adk.observe", "version": __version__},
                            "metrics": [self._serialize_metric(metric) for metric in metrics],
                        }
                    ],
                }
            ]
        }

    def _request_headers(self) -> dict[str, str]:
        headers = {"Content-Type": "application/json", **self.headers}
        if self.bearer_token and "Authorization" not in headers:
            headers["Authorization"] = f"Bearer {self.bearer_token}"
        if self.compression == "gzip":
            headers["Content-Encoding"] = "gzip"
        return headers

    def _encode_payload(self, payload: dict[str, Any]) -> bytes:
        encoded = json.dumps(payload, separators=(",", ":")).encode()
        if self.compression == "gzip":
            return gzip.compress(encoded)
        return encoded

    async def _post_payload(self, endpoint: str, payload: dict[str, Any]) -> None:
        body = self._encode_payload(payload)
        headers = self._request_headers()
        for attempt in range(self.max_retries + 1):
            try:
                if self._client is None:
                    self._client = httpx.AsyncClient(
                        timeout=self.timeout,
                        transport=self._transport,
                    )
                response = await self._client.post(
                    endpoint,
                    content=body,
                    headers=headers,
                )
                response.raise_for_status()
                return
            except httpx.HTTPError:
                if attempt >= self.max_retries:
                    raise
                delay = self.backoff_seconds * (2**attempt)
                if delay > 0:
                    await self._sleeper(delay)

    async def _flush_batches(
        self,
        items: list[Span] | list[Metric],
        *,
        endpoint: str,
        build_payload: Callable[[Any], dict[str, Any]],
    ) -> None:
        while items:
            chunk = items[: self.batch_size]
            await self._post_payload(endpoint, build_payload(chunk))
            del items[: len(chunk)]

    async def flush(self) -> None:
        if not self.spans and not self.metrics:
            return

        if self.spans:
            await self._flush_batches(
                self.spans,
                endpoint=self._derive_signal_endpoint("traces"),
                build_payload=self._build_traces_payload,
            )
        if self.metrics:
            await self._flush_batches(
                self.metrics,
                endpoint=self._derive_signal_endpoint("metrics"),
                build_payload=self._build_metrics_payload,
            )


class LangfuseExporter(OpenTelemetryExporter):
    """Langfuse-compatible exporter using the public ingestion endpoint.

    This exporter reuses the OTLP-like batching logic and targets a Langfuse
    ingestion endpoint. The default path is inferred as ``/api/public/ingestion``.
    """

    def __init__(
        self,
        *,
        base_url: str,
        public_key: str,
        secret_key: str,
        ingestion_path: str = "/api/public/ingestion",
        timeout: float = 5.0,
        transport: httpx.AsyncBaseTransport | None = None,
        batch_size: int = 100,
        max_retries: int = 3,
        backoff_seconds: float = 0.25,
        headers: dict[str, str] | None = None,
        compression: str = "none",
        sleeper: Callable[[float], Awaitable[None]] | None = None,
    ) -> None:
        token = base64.b64encode(f"{public_key}:{secret_key}".encode()).decode()
        merged_headers = {"Authorization": f"Basic {token}", **(headers or {})}
        endpoint = f"{base_url.rstrip('/')}{ingestion_path}"
        super().__init__(
            endpoint=endpoint,
            timeout=timeout,
            transport=transport,
            batch_size=batch_size,
            max_retries=max_retries,
            backoff_seconds=backoff_seconds,
            headers=merged_headers,
            compression=compression,
            sleeper=sleeper,
        )
        self.base_url = base_url.rstrip("/")
        self.public_key = public_key
        self.secret_key = secret_key
        self.ingestion_path = ingestion_path

from logos_adk.observe.native import (
    AlertEvent,
    AlertRule,
    LogEntry,
    NativeObservabilityExporter,
    NativeObservabilityStore,
    PostgreSQLObservabilityStore,
    RetentionPolicy,
    get_native_exporter,
    get_observability_store,
    reset_observability_store,
)


__all__ = [
    "ConsoleExporter",
    "Exporter",
    "InMemoryExporter",
    "Metric",
    "AlertEvent",
    "AlertRule",
    "LogEntry",
    "NativeObservabilityExporter",
    "NativeObservabilityStore",
    "PostgreSQLObservabilityStore",
    "RetentionPolicy",
    "LangfuseExporter",
    "OpenTelemetryExporter",
    "Span",
    "get_native_exporter",
    "get_observability_store",
    "reset_observability_store",
]
