"""
Dual Pool Coordinator for managing CPU and GPU worker pools.

Coordinates task distribution between CPU and GPU pools, handles result
collection using as_completed(), and provides callbacks for progress tracking.
"""

from __future__ import annotations

import logging
import time
from concurrent.futures import Future, TimeoutError, as_completed
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional

from training.processing.cpu_pool import CPUWorkerPool
from training.processing.gpu_pool import GPUWorkerPool
from training.processing.task_dispatcher import Task, TaskDispatcher, TaskType
from training.processing.worker_pool import TaskResult

logger = logging.getLogger(__name__)


@dataclass
class BatchStats:
    """Statistics for a batch processing run."""

    total: int = 0
    cpu_submitted: int = 0
    gpu_submitted: int = 0
    successful: int = 0
    failed: int = 0
    cpu_time: float = 0.0
    gpu_time: float = 0.0
    errors: List[str] = field(default_factory=list)

    @property
    def success_rate(self) -> float:
        """Calculate success rate as percentage."""
        if self.total == 0:
            return 0.0
        return (self.successful / self.total) * 100


class DualPoolCoordinator:
    """
    Coordinates CPU and GPU worker pools for parallel document processing.

    Uses separate ProcessPoolExecutor instances for CPU and GPU tasks,
    with as_completed() for efficient result collection across both pools.

    Key features:
    - Automatic task classification (CPU vs GPU)
    - Parallel submission to both pools
    - Unified result collection with timeouts
    - Progress callbacks for UI integration
    - Proper resource cleanup

    Example:
        with DualPoolCoordinator(cpu_workers=4, gpu_workers=1) as coord:
            results = coord.process_batch(
                documents=docs,
                cpu_task_fn=process_text_pdf,
                gpu_task_fn=process_scanned_pdf,
                on_result=lambda r: save_to_db(r),
            )
    """

    def __init__(
        self,
        cpu_workers: int = 4,
        gpu_workers: int = 1,
        gpu_id: int = 0,
        task_timeout: float = 300.0,
    ) -> None:
        """
        Initialize the dual pool coordinator.

        Args:
            cpu_workers: Number of CPU worker processes.
            gpu_workers: Number of GPU worker processes (usually 1).
            gpu_id: GPU device ID to use.
            task_timeout: Timeout in seconds for individual tasks.
        """
        self.cpu_workers = cpu_workers
        self.gpu_workers = gpu_workers
        self.gpu_id = gpu_id
        self.task_timeout = task_timeout

        self._cpu_pool: Optional[CPUWorkerPool] = None
        self._gpu_pool: Optional[GPUWorkerPool] = None
        self._dispatcher = TaskDispatcher()
        self._started = False

    def start(self) -> None:
        """Start both worker pools."""
        if self._started:
            raise RuntimeError("Coordinator already started")

        logger.info(
            f"Starting DualPoolCoordinator: "
            f"{self.cpu_workers} CPU workers, {self.gpu_workers} GPU workers"
        )

        self._cpu_pool = CPUWorkerPool(max_workers=self.cpu_workers)
        self._gpu_pool = GPUWorkerPool(
            max_workers=self.gpu_workers,
            gpu_id=self.gpu_id,
        )

        self._cpu_pool.start()
        self._gpu_pool.start()
        self._started = True

    def shutdown(self, wait: bool = True) -> None:
        """Shutdown both worker pools."""
        logger.info("Shutting down DualPoolCoordinator")

        if self._cpu_pool is not None:
            self._cpu_pool.shutdown(wait=wait)
            self._cpu_pool = None

        if self._gpu_pool is not None:
            self._gpu_pool.shutdown(wait=wait)
            self._gpu_pool = None

        self._started = False

    def __enter__(self) -> "DualPoolCoordinator":
        """Context manager entry."""
        self.start()
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """Context manager exit."""
        self.shutdown(wait=True)

    def process_batch(
        self,
        documents: List[dict],
        cpu_task_fn: Callable[[dict], Any],
        gpu_task_fn: Callable[[dict], Any],
        on_result: Optional[Callable[[TaskResult], None]] = None,
        on_error: Optional[Callable[[str, Exception], None]] = None,
        on_progress: Optional[Callable[[int, int], None]] = None,
        id_field: str = "id",
    ) -> List[TaskResult]:
        """
        Process a batch of documents using both CPU and GPU pools.

        Documents are automatically classified and routed to the appropriate
        pool. Results are collected as they complete.

        Args:
            documents: List of document info dicts to process.
            cpu_task_fn: Function to process text PDFs (called in CPU pool).
            gpu_task_fn: Function to process scanned PDFs (called in GPU pool).
            on_result: Callback for each successful result.
            on_error: Callback for each failed task (task_id, exception).
            on_progress: Callback for progress updates (completed, total).
            id_field: Field name to use as task ID.

        Returns:
            List of TaskResult objects for all tasks.

        Raises:
            RuntimeError: If coordinator is not started.
        """
        if not self._started:
            raise RuntimeError("Coordinator not started. Use context manager or call start().")

        if not documents:
            return []

        stats = BatchStats(total=len(documents))

        # Create and partition tasks
        tasks = self._dispatcher.create_tasks(documents, id_field=id_field)
        cpu_tasks, gpu_tasks = self._dispatcher.partition_tasks(tasks)

        # Submit tasks to pools
        futures_map: Dict[Future, Task] = {}

        # Submit CPU tasks
        cpu_start = time.time()
        for task in cpu_tasks:
            future = self._cpu_pool.submit(cpu_task_fn, task.data)
            futures_map[future] = task
            stats.cpu_submitted += 1

        # Submit GPU tasks
        gpu_start = time.time()
        for task in gpu_tasks:
            future = self._gpu_pool.submit(gpu_task_fn, task.data)
            futures_map[future] = task
            stats.gpu_submitted += 1

        logger.info(
            f"Submitted {stats.cpu_submitted} CPU tasks, {stats.gpu_submitted} GPU tasks"
        )

        # Collect results as they complete
        results: List[TaskResult] = []
        completed = 0

        for future in as_completed(futures_map.keys(), timeout=self.task_timeout * len(documents)):
            task = futures_map[future]
            pool_type = "CPU" if task.task_type == TaskType.CPU else "GPU"
            start_time = time.time()

            try:
                data = future.result(timeout=self.task_timeout)
                processing_time = time.time() - start_time

                result = TaskResult(
                    task_id=task.id,
                    success=True,
                    data=data,
                    pool_type=pool_type,
                    processing_time=processing_time,
                )
                stats.successful += 1

                if pool_type == "CPU":
                    stats.cpu_time += processing_time
                else:
                    stats.gpu_time += processing_time

                if on_result is not None:
                    try:
                        on_result(result)
                    except Exception as e:
                        logger.warning(f"on_result callback failed: {e}")

            except TimeoutError:
                error_msg = f"Task timed out after {self.task_timeout}s"
                logger.error(f"[{pool_type}] Task {task.id}: {error_msg}")

                result = TaskResult(
                    task_id=task.id,
                    success=False,
                    data=None,
                    error=error_msg,
                    pool_type=pool_type,
                )
                stats.failed += 1
                stats.errors.append(f"{task.id}: {error_msg}")

                if on_error is not None:
                    try:
                        on_error(task.id, TimeoutError(error_msg))
                    except Exception as e:
                        logger.warning(f"on_error callback failed: {e}")

            except Exception as e:
                error_msg = str(e)
                logger.error(f"[{pool_type}] Task {task.id} failed: {error_msg}")

                result = TaskResult(
                    task_id=task.id,
                    success=False,
                    data=None,
                    error=error_msg,
                    pool_type=pool_type,
                )
                stats.failed += 1
                stats.errors.append(f"{task.id}: {error_msg}")

                if on_error is not None:
                    try:
                        on_error(task.id, e)
                    except Exception as callback_error:
                        logger.warning(f"on_error callback failed: {callback_error}")

            results.append(result)
            completed += 1

            if on_progress is not None:
                try:
                    on_progress(completed, stats.total)
                except Exception as e:
                    logger.warning(f"on_progress callback failed: {e}")

        # Log final stats
        logger.info(
            f"Batch complete: {stats.successful}/{stats.total} successful "
            f"({stats.success_rate:.1f}%), {stats.failed} failed"
        )
        if stats.cpu_submitted > 0:
            logger.info(f"CPU: {stats.cpu_submitted} tasks, {stats.cpu_time:.2f}s total")
        if stats.gpu_submitted > 0:
            logger.info(f"GPU: {stats.gpu_submitted} tasks, {stats.gpu_time:.2f}s total")

        return results

    def process_single(
        self,
        document: dict,
        cpu_task_fn: Callable[[dict], Any],
        gpu_task_fn: Callable[[dict], Any],
        id_field: str = "id",
    ) -> TaskResult:
        """
        Process a single document.

        Convenience method for processing one document at a time.

        Args:
            document: Document info dict.
            cpu_task_fn: Function for text PDF processing.
            gpu_task_fn: Function for scanned PDF processing.
            id_field: Field name for task ID.

        Returns:
            TaskResult for the document.
        """
        results = self.process_batch(
            documents=[document],
            cpu_task_fn=cpu_task_fn,
            gpu_task_fn=gpu_task_fn,
            id_field=id_field,
        )
        return results[0] if results else TaskResult(
            task_id=str(document.get(id_field, "unknown")),
            success=False,
            data=None,
            error="No result returned",
        )

    @property
    def is_running(self) -> bool:
        """Check if both pools are running."""
        return (
            self._started
            and self._cpu_pool is not None
            and self._gpu_pool is not None
            and self._cpu_pool.is_running
            and self._gpu_pool.is_running
        )
