"""
Task Dispatcher for classifying and routing tasks to appropriate worker pools.

Determines whether a document should be processed by CPU (text PDF) or
GPU (scanned PDF requiring OCR) workers.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from enum import Enum, auto
from pathlib import Path
from typing import Any, List, Tuple

logger = logging.getLogger(__name__)


class TaskType(Enum):
    """Task type classification."""

    CPU = auto()  # Text PDF - no OCR needed
    GPU = auto()  # Scanned PDF - requires OCR


@dataclass
class Task:
    """
    Represents a processing task.

    Attributes:
        id: Unique task identifier.
        task_type: Whether task needs CPU or GPU processing.
        data: Task payload (document info, paths, etc.).
    """

    id: str
    task_type: TaskType
    data: Any


class TaskDispatcher:
    """
    Classifies and partitions tasks for CPU and GPU worker pools.

    Uses PDF characteristics to determine if OCR is needed:
    - Text PDFs with extractable text -> CPU
    - Scanned PDFs / image-based PDFs -> GPU (OCR)

    Example:
        dispatcher = TaskDispatcher()
        tasks = [Task(id="1", task_type=dispatcher.classify(doc), data=doc) for doc in docs]
        cpu_tasks, gpu_tasks = dispatcher.partition_tasks(tasks)
    """

    def __init__(
        self,
        text_char_threshold: int = 100,
        ocr_ratio_threshold: float = 0.3,
    ) -> None:
        """
        Initialize the task dispatcher.

        Args:
            text_char_threshold: Minimum characters to consider as text PDF.
            ocr_ratio_threshold: If text/expected ratio below this, use OCR.
        """
        self.text_char_threshold = text_char_threshold
        self.ocr_ratio_threshold = ocr_ratio_threshold

    def classify_by_pdf_info(
        self,
        has_text: bool,
        text_length: int,
        page_count: int = 1,
    ) -> TaskType:
        """
        Classify task based on PDF text extraction info.

        Args:
            has_text: Whether PDF has extractable text layer.
            text_length: Number of characters extracted.
            page_count: Number of pages in PDF.

        Returns:
            TaskType.CPU for text PDFs, TaskType.GPU for scanned PDFs.
        """
        if not has_text:
            return TaskType.GPU

        # Check if text density is reasonable
        avg_chars_per_page = text_length / max(page_count, 1)
        if avg_chars_per_page < self.text_char_threshold:
            return TaskType.GPU

        return TaskType.CPU

    def classify_document(self, doc_info: dict) -> TaskType:
        """
        Classify a document based on its metadata.

        Args:
            doc_info: Document information dict with keys like:
                     - 'is_scanned': bool (if known)
                     - 'text_length': int
                     - 'page_count': int
                     - 'pdf_path': str

        Returns:
            TaskType for the document.
        """
        # If explicitly marked as scanned
        if doc_info.get("is_scanned", False):
            return TaskType.GPU

        # If we have text extraction info
        text_length = doc_info.get("text_length", 0)
        page_count = doc_info.get("page_count", 1)
        has_text = doc_info.get("has_text", text_length > 0)

        return self.classify_by_pdf_info(
            has_text=has_text,
            text_length=text_length,
            page_count=page_count,
        )

    def partition_tasks(
        self,
        tasks: List[Task],
    ) -> Tuple[List[Task], List[Task]]:
        """
        Partition tasks into CPU and GPU groups.

        Args:
            tasks: List of Task objects with task_type set.

        Returns:
            Tuple of (cpu_tasks, gpu_tasks).
        """
        cpu_tasks = [t for t in tasks if t.task_type == TaskType.CPU]
        gpu_tasks = [t for t in tasks if t.task_type == TaskType.GPU]

        logger.info(
            f"Task partition: {len(cpu_tasks)} CPU tasks, {len(gpu_tasks)} GPU tasks"
        )

        return cpu_tasks, gpu_tasks

    def create_tasks(
        self,
        documents: List[dict],
        id_field: str = "id",
    ) -> List[Task]:
        """
        Create Task objects from document dicts.

        Args:
            documents: List of document info dicts.
            id_field: Field name to use as task ID.

        Returns:
            List of Task objects with types classified.
        """
        tasks = []
        for doc in documents:
            task_id = str(doc.get(id_field, id(doc)))
            task_type = self.classify_document(doc)
            tasks.append(Task(id=task_id, task_type=task_type, data=doc))

        cpu_count = sum(1 for t in tasks if t.task_type == TaskType.CPU)
        gpu_count = len(tasks) - cpu_count
        logger.debug(f"Created {len(tasks)} tasks: {cpu_count} CPU, {gpu_count} GPU")

        return tasks
