# TODO: Implement graceful shutdown of threads.

# Copyright © 2025 Dr.-Ing. Paul Wilhelm <paul@wilhelm.dev>
# This file is part of Archive Agent. See LICENSE for details.

import concurrent.futures
from dataclasses import dataclass
from typing import List, Union, Optional, Callable
import io
from logging import Logger

import typer

from PIL import Image

from archive_agent.ai_provider.AiProviderError import AiProviderMaxTokensError
from archive_agent.ai.AiManagerFactory import AiManagerFactory
from archive_agent.data.loader.image import ImageToTextCallback
from archive_agent.util.text_util import splitlines_exact
from archive_agent.core.ProgressManager import ProgressInfo


@dataclass
class VisionRequest:
    """
    Vision processing request containing image data, callback, and formatting logic.
    """
    image_data: Union[bytes, Image.Image]  # Support both PDF bytes and PIL Images
    callback: ImageToTextCallback          # The actual callback to use
    formatter: Callable[[Optional[str]], str]  # Lambda for conditional formatting
    log_header: str                       # Pre-built log message for progress
    image_index: int                      # For logging context
    page_index: int                       # For reassembly into per-page structure


class VisionProcessor:
    """
    Unified vision processor for parallel image-to-text processing.
    Handles both PDF and Binary document vision requests.
    """

    def __init__(
            self,
            ai_factory: AiManagerFactory,
            logger: Logger,
            verbose: bool,
            file_path: str,
            max_workers: int,
            allow_multiline: bool = False,
    ):
        """
        Initialize vision processor.
        :param ai_factory: AI manager factory for creating worker instances.
        :param logger: Logger instance from ai.cli hierarchy.
        :param verbose: Enable verbose output.
        :param file_path: File path for logging context.
        :param max_workers: Max. workers.
        :param allow_multiline: If True, skip the single-line validation so results may
                                contain line breaks (used ONLY by standalone OCR). Default
                                False preserves the original single-line guarantee for all
                                regular ingestion paths.
        """
        self.ai_factory = ai_factory
        self.logger = logger
        self.verbose = verbose
        self.file_path = file_path
        self.max_workers = max_workers
        self.allow_multiline = allow_multiline

    def process_vision_requests_parallel(
            self,
            requests: List[VisionRequest],
            progress_info: ProgressInfo
    ) -> List[str]:
        """
        Process vision requests in parallel with progress tracking.

        THREAD SAFETY: This method uses ThreadPoolExecutor to process requests
        concurrently. Each worker thread gets its own AiManager instance and
        updates progress safely through the progress_manager.

        :param requests: List of VisionRequest objects to process
        :param progress_info: Progress tracking information
        :return: List of formatted result strings in same order as requests
        """
        if not requests:
            return []

        def process_vision_request(request_data: tuple) -> tuple:
            request_index, request = request_data
            try:
                if self.verbose:
                    self.logger.info(request.log_header)

                # Create dedicated AI manager for this vision request
                ai_worker = self.ai_factory.get_ai()
                # Convert bytes to PIL Image if needed
                if isinstance(request.image_data, bytes):
                    with Image.open(io.BytesIO(request.image_data)) as image:
                        vision_result = request.callback(ai_worker, image, progress_info)
                else:
                    # Already PIL Image
                    vision_result = request.callback(ai_worker, request.image_data, progress_info)

                # Validate single-line constraint before formatting (same as original).
                # Skipped when allow_multiline is set (standalone OCR preserves structural line breaks).
                if vision_result is not None and not self.allow_multiline:
                    assert len(splitlines_exact(vision_result)) == 1, f"Text from image must be single line:\n'{vision_result}'"

                # Apply formatter to get final result
                _formatted_result = request.formatter(vision_result)

                return request_index, _formatted_result

            except typer.Exit:
                self.logger.critical(
                    f"VISION SKIPPED: Vision request ({request.image_index + 1}) of {self.file_path} "
                    f"— all retries exhausted"
                )
                _formatted_result = request.formatter(None)
                return request_index, _formatted_result
            except AiProviderMaxTokensError as e:
                self.logger.warning(f"Vision request ({request.image_index + 1}) skipped — max tokens exceeded: {e}")
                _formatted_result = request.formatter(None)
                return request_index, _formatted_result
            except Exception as e:
                self.logger.error(f"Failed to process vision request ({request.image_index + 1}): {e}")
                # Apply formatter to None for error case
                _formatted_result = request.formatter(None)

                return request_index, _formatted_result

        # Use ThreadPoolExecutor for parallel vision processing
        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            # Submit all vision tasks
            future_to_request = {
                executor.submit(process_vision_request, (request_index, request)): (request_index, request)
                for request_index, request in enumerate(requests)
            }

            # Collect results in original order
            results_dict = {}
            for future in concurrent.futures.as_completed(future_to_request):
                request_index, original_request = future_to_request[future]
                try:
                    result_index, formatted_result = future.result()
                    results_dict[result_index] = formatted_result
                except typer.Exit:
                    self.logger.critical(
                        f"VISION SKIPPED: Vision request ({request_index + 1}) of {self.file_path} "
                        f"— all retries exhausted"
                    )
                    formatted_result = original_request.formatter(None)
                    results_dict[request_index] = formatted_result
                except AiProviderMaxTokensError as exc:
                    self.logger.warning(f"Vision request ({request_index + 1}) skipped — max tokens exceeded: {exc}")
                    formatted_result = original_request.formatter(None)
                    results_dict[request_index] = formatted_result
                except Exception as exc:
                    self.logger.error(f"Vision request ({request_index + 1}) generated an exception: {exc}")
                    # Apply formatter to None for exception case
                    formatted_result = original_request.formatter(None)
                    results_dict[request_index] = formatted_result

            # Return results in original order
            return [results_dict[i] for i in range(len(requests))]
