from __future__ import annotations

import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from ocr_pipeline.config import settings
from ocr_pipeline.utils.logging import ProgressTracker, get_logger

logger = get_logger(__name__)


@dataclass
class ProcessedImage:
    path: Path
    image_hash: str
    ocr_result: Any
    preprocessing_time: float
    ocr_time: float
    success: bool
    error: str | None = None


def _process_single_image(args: tuple[Path, dict]) -> ProcessedImage:
    """Worker function for multiprocessing - must be at module level."""
    path, config_dict = args
    from ocr_pipeline.ocr.engine import ImagePreprocessor, OCREngine, compute_image_hash

    preprocessor = ImagePreprocessor()
    ocr_engine = OCREngine()

    try:
        image_hash = compute_image_hash(path)
        prep_start = time.time()
        image = preprocessor.preprocess(path)
        preprocessing_time = time.time() - prep_start

        ocr_start = time.time()
        result = ocr_engine.process(image)
        ocr_time = time.time() - ocr_start

        return ProcessedImage(
            path=path,
            image_hash=image_hash,
            ocr_result=result,
            preprocessing_time=preprocessing_time,
            ocr_time=ocr_time,
            success=True,
        )
    except Exception as e:
        return ProcessedImage(
            path=path,
            image_hash="",
            ocr_result=None,
            preprocessing_time=0,
            ocr_time=0,
            success=False,
            error=str(e),
        )


class ParallelProcessor:
    def __init__(self, workers: int | None = None):
        self.workers = workers or settings.processing.workers
        self.batch_size = settings.processing.batch_size
        self.retry_attempts = settings.processing.retry_attempts
        self.timeout = settings.processing.timeout_per_image

    def find_images(self) -> list[Path]:
        all_images = []
        for base_path_str in settings.input.paths:
            base_path = Path(base_path_str).expanduser()
            if not base_path.exists():
                logger.warning("input_path_not_found", path=str(base_path))
                continue

            for pattern in settings.input.patterns:
                if settings.input.recursive:
                    images = list(base_path.rglob(pattern))
                else:
                    images = list(base_path.glob(pattern))
                all_images.extend(images)

        unique_images = []
        seen = set()
        for img in all_images:
            if img.is_file() and img.suffix.lower() in (".png", ".jpg", ".jpeg", ".tiff", ".bmp"):
                if img not in seen:
                    seen.add(img)
                    unique_images.append(img)

        logger.info("found_images", count=len(unique_images))
        return sorted(unique_images)

    def process(self, image_paths: list[Path] | None = None) -> list[ProcessedImage]:
        if image_paths is None:
            image_paths = self.find_images()

        if not image_paths:
            logger.warning("no_images_to_process")
            return []

        results = []
        failed = []

        with ProgressTracker(len(image_paths), "OCR Processing") as tracker:
            with ProcessPoolExecutor(max_workers=self.workers) as executor:
                future_to_path = {
                    executor.submit(_process_single_image, (path, {})): path for path in image_paths
                }

                for future in as_completed(future_to_path):
                    path = future_to_path[future]
                    try:
                        result = future.result(timeout=self.timeout)
                        if not result.success and self.retry_attempts > 0:
                            for attempt in range(self.retry_attempts):
                                logger.warning("retrying", path=str(path), attempt=attempt + 1)
                                retry_result = _process_single_image((path, {}))
                                if retry_result.success:
                                    result = retry_result
                                    break
                        results.append(result)
                        tracker.update(success=result.success)
                        tracker.log_result(path, result.success, result.error)
                    except Exception as e:
                        logger.error("processing_failed", path=str(path), error=str(e))
                        results.append(
                            ProcessedImage(
                                path=path,
                                image_hash="",
                                ocr_result=None,
                                preprocessing_time=0,
                                ocr_time=0,
                                success=False,
                                error=str(e),
                            )
                        )
                        tracker.update(success=False)
                        tracker.log_result(path, False, str(e))

        successful = [r for r in results if r.success]
        failed = [r for r in results if not r.success]
        logger.info("processing_complete", successful=len(successful), failed=len(failed))
        return results