"""Bounded application parallelism for the v3 pipeline.

Provides async worker pools, specialist micro-batching, adjudicator
semaphore with queue backpressure, and load-shedding rules that never
drop safety-critical documents silently.
"""

from __future__ import annotations

import asyncio
import enum
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Callable, Coroutine


class LoadSheddingAction(str, enum.Enum):
    """Actions when load shedding is triggered."""

    QUEUE = "queue"  # Re-queue for later processing
    REJECT = "reject"  # Reject with error (non-safety-critical only)
    DEGRADE = "degrade"  # Process with reduced quality (skip optional stages)


class DocumentPriority(str, enum.Enum):
    """Document priority classes for load shedding decisions."""

    SAFETY_CRITICAL = "safety_critical"  # Never silently dropped
    HIGH = "high"
    NORMAL = "normal"
    LOW = "low"


@dataclass
class WorkerPoolConfig:
    """Configuration for an async worker pool."""

    max_workers: int = 4
    batch_size: int = 8
    batch_timeout_ms: int = 100
    queue_max_depth: int = 500
    shed_threshold: float = 0.8  # Start shedding at 80% capacity


@dataclass
class WorkerStats:
    """Runtime statistics for a worker pool."""

    active_workers: int = 0
    queued_items: int = 0
    processed_total: int = 0
    shed_total: int = 0
    errors_total: int = 0
    avg_latency_ms: float = 0.0
    last_activity: datetime | None = None


class AsyncWorkerPool:
    """Configurable async worker pool with bounded concurrency.

    Replaces the single sequential extraction loop with concurrent
    processing while respecting resource limits.
    """

    def __init__(self, config: WorkerPoolConfig | None = None) -> None:
        self.config = config or WorkerPoolConfig()
        self._semaphore = asyncio.Semaphore(self.config.max_workers)
        self._stats = WorkerStats()
        self._running = False
        self._tasks: set[asyncio.Task[Any]] = set()

    @property
    def stats(self) -> WorkerStats:
        return self._stats

    @property
    def is_running(self) -> bool:
        return self._running

    @property
    def available_slots(self) -> int:
        """Number of available worker slots."""
        return max(0, self.config.max_workers - self._stats.active_workers)

    def should_shed_load(self) -> bool:
        """Whether load shedding should be active."""
        if self.config.queue_max_depth <= 0:
            return False
        ratio = self._stats.queued_items / self.config.queue_max_depth
        return ratio >= self.config.shed_threshold

    async def submit(
        self,
        coro_fn: Callable[..., Coroutine[Any, Any, Any]],
        *args: Any,
        document_id: str = "",
        priority: DocumentPriority = DocumentPriority.NORMAL,
    ) -> LoadSheddingAction | None:
        """Submit work to the pool.

        Returns None on successful submission, or a LoadSheddingAction
        if load shedding was applied. Safety-critical documents are
        never silently rejected.
        """
        if self.should_shed_load():
            if priority == DocumentPriority.SAFETY_CRITICAL:
                # Safety-critical: always queue, never shed
                pass
            elif priority == DocumentPriority.LOW:
                self._stats.shed_total += 1
                return LoadSheddingAction.REJECT
            else:
                self._stats.shed_total += 1
                return LoadSheddingAction.QUEUE

        self._stats.queued_items += 1
        task = asyncio.create_task(self._run_with_semaphore(coro_fn, *args))
        self._tasks.add(task)
        task.add_done_callback(self._tasks.discard)
        return None

    async def _run_with_semaphore(
        self,
        coro_fn: Callable[..., Coroutine[Any, Any, Any]],
        *args: Any,
    ) -> Any:
        """Execute work bounded by the semaphore."""
        async with self._semaphore:
            self._stats.active_workers += 1
            self._stats.queued_items = max(0, self._stats.queued_items - 1)
            start = datetime.now(timezone.utc)
            try:
                result = await coro_fn(*args)
                self._stats.processed_total += 1
                return result
            except Exception:
                self._stats.errors_total += 1
                raise
            finally:
                self._stats.active_workers -= 1
                elapsed = (
                    datetime.now(timezone.utc) - start
                ).total_seconds() * 1000
                # Rolling average
                n = self._stats.processed_total + self._stats.errors_total
                if n > 0:
                    self._stats.avg_latency_ms = (
                        self._stats.avg_latency_ms * (n - 1) + elapsed
                    ) / n
                self._stats.last_activity = datetime.now(timezone.utc)

    async def start(self) -> None:
        """Mark the pool as running."""
        self._running = True

    async def shutdown(self, timeout: float = 30.0) -> None:
        """Wait for all active tasks to complete."""
        self._running = False
        if self._tasks:
            await asyncio.wait(self._tasks, timeout=timeout)


class AdjudicatorSemaphore:
    """GPU-safe concurrency control for the 9B adjudicator.

    Limits concurrent adjudication requests to match vLLM's max-num-seqs
    setting. Provides queue-depth monitoring and backpressure signaling.
    """

    def __init__(
        self,
        max_concurrent: int = 8,
        max_queued: int = 32,
    ) -> None:
        self.max_concurrent = max_concurrent
        self.max_queued = max_queued
        self._semaphore = asyncio.Semaphore(max_concurrent)
        self._queued = 0
        self._active = 0
        self._total_processed = 0

    @property
    def active_count(self) -> int:
        return self._active

    @property
    def queued_count(self) -> int:
        return self._queued

    @property
    def is_backpressured(self) -> bool:
        """Whether the adjudicator queue is full."""
        return self._queued >= self.max_queued

    async def acquire(self) -> bool:
        """Acquire adjudicator access.

        Returns False if backpressure prevents queuing.
        """
        if self._queued >= self.max_queued:
            return False
        self._queued += 1
        await self._semaphore.acquire()
        self._queued -= 1
        self._active += 1
        return True

    def release(self) -> None:
        """Release adjudicator slot."""
        self._active -= 1
        self._total_processed += 1
        self._semaphore.release()

    @property
    def utilization(self) -> float:
        """Current GPU utilization fraction."""
        return self._active / self.max_concurrent if self.max_concurrent > 0 else 0.0


@dataclass
class MicroBatcher:
    """Specialist micro-batching with configurable latency limits.

    Accumulates items until batch_size is reached or timeout expires,
    then processes the batch together for efficiency.
    """

    batch_size: int = 16
    timeout_ms: int = 50
    _buffer: list[Any] = field(default_factory=list)
    _batch_count: int = 0

    def add(self, item: Any) -> list[Any] | None:
        """Add an item. Returns a full batch if ready, else None."""
        self._buffer.append(item)
        if len(self._buffer) >= self.batch_size:
            return self.flush()
        return None

    def flush(self) -> list[Any]:
        """Force-flush the current buffer as a batch."""
        batch = self._buffer[:]
        self._buffer.clear()
        if batch:
            self._batch_count += 1
        return batch

    @property
    def pending_count(self) -> int:
        return len(self._buffer)

    @property
    def total_batches(self) -> int:
        return self._batch_count

    @property
    def is_empty(self) -> bool:
        return len(self._buffer) == 0