"""Batch transport for sending trace data to AgentLens API."""

import json
import logging
import threading
import time
from typing import Any, Dict, List, Optional

import httpx

from agentlens.models import TraceData

logger = logging.getLogger("agentlens")


class BatchTransport:
    """Thread-safe batch transport that buffers and sends traces to API."""

    def __init__(
        self,
        api_key: str,
        endpoint: str,
        max_batch_size: int = 10,
        flush_interval: float = 5.0,
    ) -> None:
        self._api_key = api_key
        self._endpoint = endpoint.rstrip("/")
        self._max_batch_size = max_batch_size
        self._flush_interval = flush_interval
        self._buffer: List[Dict[str, Any]] = []
        self._lock = threading.Lock()
        self._shutdown_event = threading.Event()
        self._flush_thread = threading.Thread(
            target=self._flush_loop, daemon=True, name="agentlens-flush"
        )
        self._flush_thread.start()
        self._client = httpx.Client(timeout=30.0)

    def add(self, trace: TraceData) -> None:
        """Add a completed trace to send buffer."""
        trace_dict = trace.to_dict()
        with self._lock:
            self._buffer.append(trace_dict)
            should_flush = len(self._buffer) >= self._max_batch_size
        if should_flush:
            self._do_flush()

    def _flush_loop(self) -> None:
        """Background loop that periodically flushes buffer."""
        while not self._shutdown_event.is_set():
            self._shutdown_event.wait(timeout=self._flush_interval)
            if not self._shutdown_event.is_set():
                self._do_flush()

    def _do_flush(self) -> None:
        """Flush all buffered traces to the API."""
        with self._lock:
            if not self._buffer:
                return
            batch = self._buffer.copy()
            self._buffer.clear()

        try:
            response = self._client.post(
                f"{self._endpoint}/api/traces",
                json={"traces": batch},
                headers={
                    "Authorization": f"Bearer {self._api_key}",
                    "Content-Type": "application/json",
                },
            )
            if response.status_code >= 400:
                logger.warning(
                    "AgentLens: Failed to send traces (HTTP %d): %s",
                    response.status_code,
                    response.text[:200],
                )
        except Exception as e:
            logger.warning("AgentLens: Failed to send traces: %s", e)
            # Put traces back in buffer for retry (optional, up to you)
            # For now, drop on failure to avoid unbounded growth

    def flush(self) -> None:
        """Manually flush all buffered traces."""
        self._do_flush()

    def shutdown(self) -> None:
        """Shutdown transport: flush remaining traces and stop background thread."""
        self._shutdown_event.set()
        self._flush_thread.join(timeout=5.0)
        self._do_flush()  # Final flush
        self._client.close()
