"""Shared streaming infrastructure for real-time extractors.

Provides common configuration and batching utilities used by SSE,
WebSocket, and file watcher extractors.
"""

from __future__ import annotations

import asyncio
import time
from dataclasses import dataclass, field
from typing import Any


@dataclass(frozen=True, slots=True)
class StreamingConfig:
    """Common configuration for streaming extractors.

    Attributes:
        shutdown_event: External event to signal graceful shutdown.
        max_batches: Stop after yielding this many batches.
        max_records: Stop after processing this many total records.
        idle_timeout: Stop after this many seconds with no new data.
        batch_timeout: Yield a partial batch after this many seconds.
        batch_size: Maximum records per batch before yielding.
    """

    shutdown_event: asyncio.Event | None = None
    max_batches: int | None = None
    max_records: int | None = None
    idle_timeout: float | None = None
    batch_timeout: float = 5.0
    batch_size: int = 1000


class BatchAccumulator:
    """Accumulates records into batches, yielding on size or timeout.

    Args:
        batch_size: Maximum records per batch.
        batch_timeout: Seconds before yielding a partial batch.
    """

    def __init__(self, batch_size: int, batch_timeout: float) -> None:
        self._batch_size = batch_size
        self._batch_timeout = batch_timeout
        self._buffer: list[dict[str, Any]] = []
        self._last_add_time: float = time.monotonic()

    async def add(self, record: dict[str, Any]) -> list[dict[str, Any]] | None:
        """Add a record and return a batch if full.

        Args:
            record: Record to add to the buffer.

        Returns:
            A full batch if the buffer reached batch_size, otherwise None.
        """
        self._buffer.append(record)
        self._last_add_time = time.monotonic()
        if len(self._buffer) >= self._batch_size:
            return self._drain()
        return None

    async def flush(self) -> list[dict[str, Any]] | None:
        """Flush remaining records as a partial batch.

        Returns:
            Remaining records if any, otherwise None.
        """
        if self._buffer:
            return self._drain()
        return None

    def should_flush_timeout(self) -> bool:
        """Check if the batch timeout has elapsed since the last add."""
        if not self._buffer:
            return False
        return (time.monotonic() - self._last_add_time) >= self._batch_timeout

    def _drain(self) -> list[dict[str, Any]]:
        """Drain the buffer and return its contents."""
        batch = self._buffer
        self._buffer = []
        self._last_add_time = time.monotonic()
        return batch
