"""
Batch span processor for efficient span export.

Collects spans and exports them in batches based on:
- Batch size threshold (default: 100 spans)
- Time interval (default: 1000ms)
"""

from __future__ import annotations

import logging
import sys
import threading
import time
from collections import deque
from typing import Callable, Deque, List, Optional, TYPE_CHECKING

from risicare.exporters.base import ExportResult, SpanExporter

logger = logging.getLogger("risicare.exporters.batch")

if TYPE_CHECKING:
    from risicare_core import Span


class BatchSpanProcessor:
    """
    Processor that batches spans before export.

    This processor collects spans in memory and exports them when:
    - The batch reaches the configured size
    - The configured timeout elapses

    Thread-safe for concurrent span submissions.

    Attributes:
        dropped_spans: Count of spans dropped due to queue overflow.
        exported_spans: Count of spans successfully exported.
        failed_exports: Count of failed export attempts.
    """

    def __init__(
        self,
        exporters: List[SpanExporter],
        batch_size: int = 100,
        batch_timeout_ms: int = 1000,
        max_queue_size: int = 10000,
        debug: bool = False,
        on_drop: Optional[Callable[[int], None]] = None,
    ):
        """
        Initialize the batch processor.

        Args:
            exporters: List of exporters to send spans to.
            batch_size: Number of spans to collect before export.
            batch_timeout_ms: Maximum time to wait before flushing.
            max_queue_size: Maximum queue size (warns and drops if exceeded).
            debug: Enable debug logging.
            on_drop: Optional callback invoked with total drop count when spans are dropped.
        """
        self._exporters = list(exporters)
        self._batch_size = max(1, min(batch_size, 10_000))
        self._batch_timeout_ms = max(100, min(batch_timeout_ms, 60_000))
        self._max_queue_size = max(100, min(max_queue_size, 1_000_000))
        self._debug = debug
        self._on_drop = on_drop

        # Unbounded deque — drops handled explicitly in on_span_end()
        self._queue: Deque["Span"] = deque()
        self._lock = threading.Lock()
        self._flush_event = threading.Event()
        self._shutdown_event = threading.Event()
        self._worker_thread: Optional[threading.Thread] = None
        self._started = False

        # Metrics
        self.dropped_spans = 0
        self.exported_spans = 0
        self.failed_exports = 0

    def start(self) -> None:
        """Start the background worker thread."""
        if self._started:
            return

        self._shutdown_event.clear()
        self._worker_thread = threading.Thread(
            target=self._worker_loop,
            name="risicare-batch-processor",
            daemon=True,
        )
        self._worker_thread.start()
        self._started = True

    def shutdown(self, timeout_ms: int = 5000) -> None:
        """
        Shutdown the processor.

        Flushes remaining spans and stops the worker thread.

        Args:
            timeout_ms: Maximum time to wait for flush.
        """
        if not self._started:
            return

        # Signal shutdown
        self._shutdown_event.set()
        self._flush_event.set()

        # Wait for worker to finish
        if self._worker_thread:
            self._worker_thread.join(timeout=timeout_ms / 1000)
            self._worker_thread = None

        # Final flush
        self._export_batch()

        # Shutdown exporters
        for exporter in self._exporters:
            try:
                exporter.shutdown()
            except Exception as e:
                logger.warning("Error shutting down exporter %s: %s", exporter.name, e)

        self._started = False

        if self._debug:
            logger.debug(
                "BatchSpanProcessor shutdown complete. "
                "Exported: %d, Dropped: %d, Failed: %d",
                self.exported_spans, self.dropped_spans, self.failed_exports,
            )

    def on_span_end(self, span: "Span") -> None:
        """
        Called when a span ends.

        Adds the span to the queue for batch export.
        If the queue is full, the span is dropped and a warning is issued.

        Args:
            span: The completed span.
        """
        if not self._started:
            return

        with self._lock:
            # Check queue size before adding
            if len(self._queue) >= self._max_queue_size:
                self.dropped_spans += 1
                # Warn every 1000 drops (not just once)
                if self.dropped_spans == 1 or self.dropped_spans % 1000 == 0:
                    logger.warning(
                        "Span queue full (%d). %d spans dropped so far.",
                        self._max_queue_size, self.dropped_spans,
                    )
                if self._on_drop:
                    self._on_drop(self.dropped_spans)
                return

            self._queue.append(span)

            # Trigger flush if batch size reached
            if len(self._queue) >= self._batch_size:
                self._flush_event.set()

    def flush(self, timeout_ms: int = 5000) -> bool:
        """
        Flush pending spans.

        Args:
            timeout_ms: Maximum time to wait.

        Returns:
            True if flush completed within timeout.
        """
        if not self._started:
            return True

        # Signal flush
        self._flush_event.set()

        # Wait for queue to empty (with timeout)
        start = time.monotonic()
        while True:
            with self._lock:
                if len(self._queue) == 0:
                    return True

            elapsed = (time.monotonic() - start) * 1000
            if elapsed >= timeout_ms:
                return False

            time.sleep(0.01)  # 10ms poll interval

    def add_exporter(self, exporter: SpanExporter) -> None:
        """Add an exporter."""
        self._exporters.append(exporter)

    def _worker_loop(self) -> None:
        """Background worker that exports batches."""
        while not self._shutdown_event.is_set():
            # Wait for flush signal or timeout
            triggered = self._flush_event.wait(
                timeout=self._batch_timeout_ms / 1000
            )

            # Export batch
            self._export_batch()

            # Clear flush event
            if triggered:
                self._flush_event.clear()

    def _export_batch(self) -> None:
        """Export current batch to all exporters."""
        # Collect spans to export
        with self._lock:
            if not self._queue:
                return

            # Take up to batch_size spans
            batch: List["Span"] = []
            while self._queue and len(batch) < self._batch_size:
                batch.append(self._queue.popleft())

        if not batch:
            return

        if self._debug:
            logger.debug("Exporting batch of %d spans", len(batch))

        # Export to all exporters — count batch once, not per-exporter
        batch_exported = False
        for exporter in self._exporters:
            try:
                result = exporter.export(batch)
                if result == ExportResult.SUCCESS:
                    if not batch_exported:
                        self.exported_spans += len(batch)
                        batch_exported = True
                    if self._debug:
                        logger.debug("Export to %s succeeded", exporter.name)
                else:
                    self.failed_exports += 1
                    if self._debug:
                        logger.debug("Export to %s failed: %s", exporter.name, result)
            except Exception as e:
                self.failed_exports += 1
                logger.warning("Export to %s raised exception: %s", exporter.name, e)

    def get_metrics(self) -> dict:
        """Return current processor metrics for monitoring."""
        with self._lock:
            queue_size = len(self._queue)
        return {
            "exported_spans": self.exported_spans,
            "dropped_spans": self.dropped_spans,
            "failed_exports": self.failed_exports,
            "queue_size": queue_size,
            "queue_capacity": self._max_queue_size,
            "queue_utilization": queue_size / self._max_queue_size if self._max_queue_size > 0 else 0,
        }
