"""
Batch processing module for high-volume guardrail checks.

This module provides efficient batch processing capabilities for checking multiple
inputs in parallel with automatic batching, queueing, and GPU optimization.

Example:
    >>> from starvex.batch import BatchProcessor, BatchRequest
    >>> processor = BatchProcessor(rules=[BlockJailbreak(), BlockPII()])
    >>>
    >>> requests = [
    ...     BatchRequest(text="Input 1", metadata={"user_id": "123"}),
    ...     BatchRequest(text="Input 2", metadata={"user_id": "456"}),
    ... ]
    >>>
    >>> results = await processor.process_batch(requests)
    >>> for result in results:
    ...     print(f"Status: {result.status}, Confidence: {result.confidence}")
"""

import asyncio
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
from uuid import uuid4

from .guardrails.guard import StarvexGuard, CheckResult
from .guardrails.rules import GuardRule


@dataclass
class BatchRequest:
    """
    A single request in a batch.

    Attributes:
        text: The text to check against guardrails
        metadata: Optional metadata to attach to this request
        trace_id: Optional trace ID for correlation (auto-generated if not provided)
        check_input: Whether this is an input check (True) or output check (False)
    """
    text: str
    metadata: Optional[Dict[str, Any]] = None
    trace_id: Optional[str] = None
    check_input: bool = True

    def __post_init__(self) -> None:
        if self.trace_id is None:
            self.trace_id = str(uuid4())
        if self.metadata is None:
            self.metadata = {}


@dataclass
class BatchResult:
    """
    Result of a batch processing operation.

    Attributes:
        results: List of CheckResult objects, one per input
        total_ms: Total processing time in milliseconds
        batch_id: Unique identifier for this batch
        throughput: Items processed per second
        avg_latency_ms: Average latency per item in milliseconds
    """
    results: List[CheckResult]
    total_ms: int
    batch_id: str
    throughput: float
    avg_latency_ms: float


@dataclass
class BatchConfig:
    """
    Configuration for batch processing.

    Attributes:
        max_batch_size: Maximum number of items to process in a single batch
        max_queue_size: Maximum number of items to queue before blocking
        max_workers: Maximum number of worker threads for parallel processing
        batch_timeout_ms: Maximum time to wait for batch to fill before processing
        enable_gpu_batching: Whether to batch operations for GPU efficiency
    """
    max_batch_size: int = 100
    max_queue_size: int = 1000
    max_workers: int = 4
    batch_timeout_ms: int = 100
    enable_gpu_batching: bool = True


class BatchProcessor:
    """
    High-performance batch processor for guardrail checks.

    This processor automatically batches requests, uses thread pools for parallelization,
    and optimizes for GPU throughput when available.

    Example:
        >>> from starvex.batch import BatchProcessor, BatchConfig
        >>> from starvex.guardrails.rules import BlockJailbreak, BlockPII
        >>>
        >>> config = BatchConfig(max_batch_size=50, max_workers=8)
        >>> processor = BatchProcessor(
        ...     rules=[BlockJailbreak(), BlockPII()],
        ...     config=config
        ... )
        >>>
        >>> # Process a large batch
        >>> requests = [BatchRequest(text=f"Input {i}") for i in range(500)]
        >>> result = await processor.process_batch(requests)
        >>> print(f"Processed {len(result.results)} items in {result.total_ms}ms")
        >>> print(f"Throughput: {result.throughput:.0f} req/sec")
    """

    def __init__(
        self,
        rules: Optional[List[GuardRule]] = None,
        config: Optional[BatchConfig] = None,
        guard: Optional[StarvexGuard] = None,
    ):
        """
        Initialize the batch processor.

        Args:
            rules: List of guardrail rules to apply (required if guard not provided)
            config: Batch processing configuration
            guard: Pre-configured StarvexGuard instance (optional)
        """
        if guard is None and not rules:
            raise ValueError("Either 'guard' or 'rules' must be provided")

        self.guard = guard or StarvexGuard(rules=rules)
        self.config = config or BatchConfig()
        self.executor = ThreadPoolExecutor(max_workers=self.config.max_workers)
        self._queue: asyncio.Queue[BatchRequest] = asyncio.Queue(
            maxsize=self.config.max_queue_size
        )
        self._processing = False

    async def process_batch(
        self,
        requests: List[BatchRequest],
        chunk_size: Optional[int] = None,
    ) -> BatchResult:
        """
        Process a batch of requests.

        Large batches are automatically chunked for optimal performance.

        Args:
            requests: List of BatchRequest objects to process
            chunk_size: Optional chunk size (defaults to config.max_batch_size)

        Returns:
            BatchResult containing all results and performance metrics
        """
        start_time = time.time()
        batch_id = str(uuid4())
        chunk_size = chunk_size or self.config.max_batch_size

        # Split into chunks if needed
        chunks = [
            requests[i:i + chunk_size]
            for i in range(0, len(requests), chunk_size)
        ]

        # Process chunks in parallel
        all_results: List[CheckResult] = []

        if len(chunks) == 1:
            # Single chunk - process directly
            all_results = await self._process_chunk(chunks[0])
        else:
            # Multiple chunks - process in parallel
            chunk_tasks = [self._process_chunk(chunk) for chunk in chunks]
            chunk_results = await asyncio.gather(*chunk_tasks)
            for results in chunk_results:
                all_results.extend(results)

        # Calculate metrics
        elapsed_ms = int((time.time() - start_time) * 1000)
        throughput = len(requests) / (elapsed_ms / 1000) if elapsed_ms > 0 else 0
        avg_latency = elapsed_ms / len(requests) if requests else 0

        return BatchResult(
            results=all_results,
            total_ms=elapsed_ms,
            batch_id=batch_id,
            throughput=throughput,
            avg_latency_ms=avg_latency,
        )

    async def _process_chunk(self, chunk: List[BatchRequest]) -> List[CheckResult]:
        """
        Process a single chunk of requests.

        Args:
            chunk: List of BatchRequest objects in this chunk

        Returns:
            List of CheckResult objects
        """
        # If GPU batching is enabled and guard supports it, use batch processing
        if self.config.enable_gpu_batching and hasattr(self.guard, 'check_batch'):
            texts = [req.text for req in chunk]
            results = await self.guard.check_batch(
                texts=texts,
                check_input=chunk[0].check_input if chunk else True,
            )

            # Add metadata from original requests
            for i, result in enumerate(results):
                if i < len(chunk):
                    result.trace_id = chunk[i].trace_id
                    if chunk[i].metadata:
                        result.metadata.update(chunk[i].metadata)

            return results

        # Fallback to parallel individual checks
        tasks = [self._process_single(req) for req in chunk]
        return await asyncio.gather(*tasks)

    async def _process_single(self, request: BatchRequest) -> CheckResult:
        """
        Process a single request.

        Args:
            request: BatchRequest to process

        Returns:
            CheckResult for this request
        """
        if request.check_input:
            result = await self.guard.check_input_async(request.text)
        else:
            result = await self.guard.check_output_async(request.text)

        # Add metadata
        result.trace_id = request.trace_id
        if request.metadata:
            result.metadata.update(request.metadata)

        return result

    async def submit(self, request: BatchRequest) -> None:
        """
        Submit a request to the processing queue.

        This method is useful for streaming scenarios where requests arrive
        over time and should be automatically batched.

        Args:
            request: BatchRequest to queue for processing
        """
        await self._queue.put(request)

    async def start_auto_batching(self) -> None:
        """
        Start automatic batch processing from the queue.

        This runs in the background and automatically processes queued requests
        in optimally-sized batches.
        """
        self._processing = True

        while self._processing:
            batch: List[BatchRequest] = []
            timeout = self.config.batch_timeout_ms / 1000

            try:
                # Wait for first item
                request = await asyncio.wait_for(
                    self._queue.get(),
                    timeout=timeout
                )
                batch.append(request)

                # Try to fill batch up to max size
                while len(batch) < self.config.max_batch_size:
                    try:
                        request = self._queue.get_nowait()
                        batch.append(request)
                    except asyncio.QueueEmpty:
                        break

                # Process the batch
                if batch:
                    await self._process_chunk(batch)

            except asyncio.TimeoutError:
                # No requests in timeout period - continue waiting
                continue

    def stop_auto_batching(self) -> None:
        """Stop the automatic batch processing."""
        self._processing = False

    def __enter__(self) -> "BatchProcessor":
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.executor.shutdown(wait=True)

    async def __aenter__(self) -> "BatchProcessor":
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        self.stop_auto_batching()
        self.executor.shutdown(wait=True)


class StreamingBatchProcessor:
    """
    Streaming batch processor with automatic flushing.

    This processor accumulates requests and automatically flushes them when:
    - The batch reaches max_batch_size
    - The flush_interval_ms timeout expires
    - flush() is called manually

    Example:
        >>> async with StreamingBatchProcessor(rules=[BlockPII()]) as processor:
        ...     for user_input in user_inputs:
        ...         result = await processor.add(user_input)
        ...         if result:  # Result available immediately
        ...             print(f"Status: {result.status}")
        ...
        ...     # Flush any remaining items
        ...     final_results = await processor.flush()
    """

    def __init__(
        self,
        rules: Optional[List[GuardRule]] = None,
        max_batch_size: int = 32,
        flush_interval_ms: int = 100,
        guard: Optional[StarvexGuard] = None,
    ):
        """
        Initialize the streaming batch processor.

        Args:
            rules: List of guardrail rules to apply
            max_batch_size: Maximum batch size before auto-flush
            flush_interval_ms: Maximum time before auto-flush
            guard: Pre-configured StarvexGuard instance (optional)
        """
        self.processor = BatchProcessor(
            rules=rules,
            guard=guard,
            config=BatchConfig(max_batch_size=max_batch_size),
        )
        self.pending: List[BatchRequest] = []
        self.flush_interval_ms = flush_interval_ms
        self.last_flush = time.time()

    async def add(self, text: str, metadata: Optional[Dict[str, Any]] = None) -> Optional[CheckResult]:
        """
        Add a text to the batch.

        Args:
            text: Text to check
            metadata: Optional metadata

        Returns:
            CheckResult if batch was flushed, None otherwise
        """
        request = BatchRequest(text=text, metadata=metadata)
        self.pending.append(request)

        # Auto-flush if batch is full or timeout expired
        should_flush = (
            len(self.pending) >= self.processor.config.max_batch_size or
            (time.time() - self.last_flush) * 1000 >= self.flush_interval_ms
        )

        if should_flush:
            results = await self.flush()
            return results[-1] if results else None

        return None

    async def flush(self) -> List[CheckResult]:
        """
        Flush all pending requests.

        Returns:
            List of CheckResult objects for all pending requests
        """
        if not self.pending:
            return []

        result = await self.processor.process_batch(self.pending)
        self.pending = []
        self.last_flush = time.time()

        return result.results

    async def __aenter__(self) -> "StreamingBatchProcessor":
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        # Flush any remaining items
        await self.flush()


__all__ = [
    "BatchRequest",
    "BatchResult",
    "BatchConfig",
    "BatchProcessor",
    "StreamingBatchProcessor",
]
