from __future__ import annotations

import atexit
import logging
import queue
import threading
import time
from collections import defaultdict
from typing import Any

import httpx

logger = logging.getLogger("staso")

_MAX_BATCH = 500  # backend hard limit per request


def _spans_to_traces(
    spans: list[dict],
    environment: str = "default",
    workspace_slug: str = "",
) -> list[dict]:
    """Group pre-formed span dicts by trace_id and wrap for the backend."""
    groups: dict[str, list[dict]] = defaultdict(list)
    for span in spans:
        groups[span.get("trace_id", "")].append(span)

    traces: list[dict] = []
    for trace_id, group in groups.items():
        earliest_ts = group[0].get("timestamp", "")
        for s in group:
            ts = s.get("timestamp", "")
            if ts and ts < earliest_ts:
                earliest_ts = ts

        trace_spans: list[dict[str, Any]] = []
        for s in group:
            trace_spans.append({
                "span_id": s.get("span_id"),
                "parent_span_id": s.get("parent_span_id"),
                "name": s.get("name", ""),
                "kind": s.get("kind", "chain"),
                "timestamp": s.get("timestamp", ""),
                "duration_ms": s.get("duration_ms", 0),
                "status": s.get("status", "ok"),
                "model": s.get("model", ""),
                "input_tokens": s.get("input_tokens", 0),
                "output_tokens": s.get("output_tokens", 0),
                "total_tokens": s.get("total_tokens", 0),
                "input": s.get("input", ""),
                "output": s.get("output", ""),
                "attributes": s.get("attributes", "{}"),
                "error_message": s.get("error_message"),
                "session_id": s.get("session_id", ""),
                "agent_id": s.get("agent_id", ""),
                "user_id": s.get("user_id", ""),
            })

        traces.append({
            "trace_id": trace_id,
            "timestamp": earliest_ts,
            "environment": environment,
            "workspace_slug": workspace_slug,
            "spans": trace_spans,
        })

    return traces


class BatchTransport:
    """Background thread that batches span dicts and flushes them to the ingest API."""

    def __init__(
        self,
        base_url: str,
        api_key: str,
        *,
        batch_size: int = 100,
        flush_interval: float = 5.0,
        max_queue_size: int = 10_000,
        environment: str = "default",
        workspace_slug: str = "",
    ) -> None:
        self._url = f"{base_url.rstrip('/')}/v1/traces/ingest"
        self._api_key = api_key
        self._environment = environment
        self._workspace_slug = workspace_slug
        self._batch_size = min(batch_size, _MAX_BATCH)
        self._flush_interval = flush_interval
        self._queue: queue.Queue[dict] = queue.Queue(maxsize=max_queue_size)
        self._shutdown = threading.Event()
        self._client = httpx.Client(
            timeout=10.0,
            headers={"X-API-Key": api_key, "Content-Type": "application/json"},
        )
        self._thread = threading.Thread(
            target=self._worker, daemon=True, name="staso-transport"
        )
        self._thread.start()
        atexit.register(self.shutdown)

    # -- public ---------------------------------------------------------------

    def enqueue(self, span_dict: dict) -> None:
        """Add a span dict to the send queue. Drops silently when the queue is full."""
        if self._shutdown.is_set():
            return
        try:
            self._queue.put_nowait(span_dict)
        except queue.Full:
            logger.warning("staso: event queue full, dropping event")

    def shutdown(self) -> None:
        """Signal the worker to flush remaining spans and stop."""
        if self._shutdown.is_set():
            return
        self._shutdown.set()
        self._thread.join(timeout=5.0)
        try:
            self._client.close()
        except Exception:
            pass

    # -- private --------------------------------------------------------------

    def _worker(self) -> None:
        buf: list[dict] = []
        last_flush = time.monotonic()

        while not self._shutdown.is_set():
            try:
                buf.append(self._queue.get(timeout=0.1))
            except queue.Empty:
                pass

            elapsed = time.monotonic() - last_flush
            if len(buf) >= self._batch_size or (buf and elapsed >= self._flush_interval):
                self._flush(buf)
                buf = []
                last_flush = time.monotonic()

        # Drain remaining spans on shutdown
        while True:
            try:
                buf.append(self._queue.get_nowait())
            except queue.Empty:
                break
        if buf:
            self._flush(buf)

    def _send_with_retry(self, payload: dict, max_attempts: int = 3) -> None:
        """Send payload with exponential backoff. Retry on 5xx/network errors only."""
        delays = [0.5, 1.0, 2.0]
        for attempt in range(max_attempts):
            try:
                resp = self._client.post(self._url, json=payload)
                if resp.status_code < 500:
                    if resp.status_code >= 400:
                        logger.warning(
                            "staso: ingest returned %d – %s",
                            resp.status_code,
                            resp.text[:300],
                        )
                    return  # success or client error (don't retry 4xx)
                # 5xx — retry
                logger.warning(
                    "staso: ingest returned %d, attempt %d/%d",
                    resp.status_code,
                    attempt + 1,
                    max_attempts,
                )
            except Exception as exc:
                logger.warning(
                    "staso: ingest error on attempt %d/%d: %s",
                    attempt + 1,
                    max_attempts,
                    exc,
                )
            if attempt < max_attempts - 1:
                time.sleep(delays[attempt])
        else:
            logger.warning(
                "staso: ingest failed after %d attempts, dropping batch",
                max_attempts,
            )

    def _flush(self, spans: list[dict]) -> None:
        traces = _spans_to_traces(spans, self._environment, self._workspace_slug)
        for i in range(0, len(traces), _MAX_BATCH):
            chunk = traces[i : i + _MAX_BATCH]
            self._send_with_retry({"traces": chunk})
