"""
Task functions for autolabel processing in multi-pool architecture.

Provides CPU and GPU task functions that can be called from worker pools.
"""

from __future__ import annotations

import os
import time
import warnings
from pathlib import Path
from typing import Any, Dict, Optional

from shared.config import DEFAULT_DPI

# Global OCR instance (initialized once per GPU worker process)
_ocr_engine: Optional[Any] = None


def init_cpu_worker() -> None:
    """
    Initialize CPU worker process.

    Disables GPU access and suppresses unnecessary warnings.
    """
    os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
    os.environ.setdefault("OMP_NUM_THREADS", "1")
    os.environ.setdefault("MKL_NUM_THREADS", "1")


def init_gpu_worker(gpu_id: int = 0, gpu_mem: int = 4000) -> None:
    """
    Initialize GPU worker process with PaddleOCR.

    Args:
        gpu_id: GPU device ID.
        gpu_mem: Maximum GPU memory in MB.
    """
    global _ocr_engine

    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
    os.environ["GLOG_minloglevel"] = "2"

    # Suppress PaddleX warnings
    warnings.filterwarnings("ignore", message=".*PDX has already been initialized.*")
    warnings.filterwarnings("ignore", message=".*reinitialization.*")

    # Lazy initialization - OCR will be created on first use
    _ocr_engine = None


def _get_ocr_engine():
    """Get or create OCR engine for current GPU worker."""
    global _ocr_engine

    if _ocr_engine is None:
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore")
            from shared.ocr import OCREngine
            _ocr_engine = OCREngine()

    return _ocr_engine


def _save_output_img(output_img, image_path: Path) -> None:
    """Save OCR preprocessed image to replace rendered image."""
    from PIL import Image as PILImage

    if output_img is not None:
        img = PILImage.fromarray(output_img)
        img.save(str(image_path))


def process_text_pdf(task_data: Dict[str, Any]) -> Dict[str, Any]:
    """
    Process a text PDF (CPU task - no OCR needed).

    Args:
        task_data: Dictionary with keys:
            - row_dict: Document fields from CSV
            - pdf_path: Path to PDF file
            - output_dir: Output directory
            - dpi: Rendering DPI
            - min_confidence: Minimum match confidence

    Returns:
        Result dictionary with success status, annotations, and report.
    """
    import shutil
    from training.data.autolabel_report import AutoLabelReport
    from shared.pdf import PDFDocument
    from shared.fields import TRAINING_FIELD_CLASSES as FIELD_CLASSES
    from training.processing.document_processor import process_page, record_unmatched_fields

    row_dict = task_data["row_dict"]
    pdf_path = Path(task_data["pdf_path"])
    output_dir = Path(task_data["output_dir"])
    dpi = task_data.get("dpi", DEFAULT_DPI)
    min_confidence = task_data.get("min_confidence", 0.5)

    start_time = time.time()
    doc_id = row_dict["DocumentId"]

    # Clean up existing temp folder for this document (for re-matching)
    temp_doc_dir = output_dir / "temp" / doc_id
    if temp_doc_dir.exists():
        shutil.rmtree(temp_doc_dir, ignore_errors=True)

    report = AutoLabelReport(document_id=doc_id)
    report.pdf_path = str(pdf_path)
    report.pdf_type = "text"
    # Store metadata fields from CSV (same as single document mode)
    report.split = row_dict.get('split')
    report.customer_number = row_dict.get('customer_number')
    report.supplier_name = row_dict.get('supplier_name')
    report.supplier_organisation_number = row_dict.get('supplier_organisation_number')
    report.supplier_accounts = row_dict.get('supplier_accounts')

    result = {
        "doc_id": doc_id,
        "success": False,
        "pages": [],
        "report": None,
        "stats": {name: 0 for name in FIELD_CLASSES.keys()},
    }

    try:
        with PDFDocument(pdf_path) as pdf_doc:
            page_annotations = []
            matched_fields = set()

            images_dir = output_dir / "temp" / doc_id / "images"
            for page_no, image_path in pdf_doc.render_all_pages(images_dir, dpi=dpi):
                report.total_pages += 1
                img_width, img_height = pdf_doc.get_render_dimensions(page_no, dpi)

                # Text extraction (no OCR)
                tokens = list(pdf_doc.extract_text_tokens(page_no))

                # Get page dimensions for payment line detection
                page = pdf_doc.doc[page_no]
                page_height = page.rect.height
                page_width = page.rect.width

                # Use shared processing logic (same as single document mode)
                matches = {}
                annotations, ann_count = process_page(
                    tokens=tokens,
                    row_dict=row_dict,
                    page_no=page_no,
                    page_height=page_height,
                    page_width=page_width,
                    img_width=img_width,
                    img_height=img_height,
                    dpi=dpi,
                    min_confidence=min_confidence,
                    matches=matches,
                    matched_fields=matched_fields,
                    report=report,
                    result_stats=result["stats"],
                )

                if annotations:
                    page_annotations.append(
                        {
                            "image_path": str(image_path),
                            "page_no": page_no,
                            "count": ann_count,
                        }
                    )
                    report.annotations_generated += ann_count

            # Record unmatched fields using shared logic
            record_unmatched_fields(row_dict, matched_fields, report)

            if page_annotations:
                result["pages"] = page_annotations
                result["success"] = True
                report.success = True
            else:
                report.errors.append("No annotations generated")

    except Exception as e:
        report.errors.append(str(e))

    report.processing_time_ms = (time.time() - start_time) * 1000
    result["report"] = report.to_dict()

    return result


def process_scanned_pdf(task_data: Dict[str, Any]) -> Dict[str, Any]:
    """
    Process a scanned PDF (GPU task - requires OCR).

    Args:
        task_data: Dictionary with keys:
            - row_dict: Document fields from CSV
            - pdf_path: Path to PDF file
            - output_dir: Output directory
            - dpi: Rendering DPI
            - min_confidence: Minimum match confidence

    Returns:
        Result dictionary with success status, annotations, and report.
    """
    import shutil
    from training.data.autolabel_report import AutoLabelReport
    from shared.pdf import PDFDocument
    from shared.fields import TRAINING_FIELD_CLASSES as FIELD_CLASSES
    from training.processing.document_processor import process_page, record_unmatched_fields

    row_dict = task_data["row_dict"]
    pdf_path = Path(task_data["pdf_path"])
    output_dir = Path(task_data["output_dir"])
    dpi = task_data.get("dpi", DEFAULT_DPI)
    min_confidence = task_data.get("min_confidence", 0.5)

    start_time = time.time()
    doc_id = row_dict["DocumentId"]

    # Clean up existing temp folder for this document (for re-matching)
    temp_doc_dir = output_dir / "temp" / doc_id
    if temp_doc_dir.exists():
        shutil.rmtree(temp_doc_dir, ignore_errors=True)

    report = AutoLabelReport(document_id=doc_id)
    report.pdf_path = str(pdf_path)
    report.pdf_type = "scanned"
    # Store metadata fields from CSV (same as single document mode)
    report.split = row_dict.get('split')
    report.customer_number = row_dict.get('customer_number')
    report.supplier_name = row_dict.get('supplier_name')
    report.supplier_organisation_number = row_dict.get('supplier_organisation_number')
    report.supplier_accounts = row_dict.get('supplier_accounts')

    result = {
        "doc_id": doc_id,
        "success": False,
        "pages": [],
        "report": None,
        "stats": {name: 0 for name in FIELD_CLASSES.keys()},
    }

    try:
        # Get OCR engine from worker cache
        ocr_engine = _get_ocr_engine()

        with PDFDocument(pdf_path) as pdf_doc:
            page_annotations = []
            matched_fields = set()

            images_dir = output_dir / "temp" / doc_id / "images"
            for page_no, image_path in pdf_doc.render_all_pages(images_dir, dpi=dpi):
                report.total_pages += 1
                img_width, img_height = pdf_doc.get_render_dimensions(page_no, dpi)

                # Get page dimensions for payment line detection
                page = pdf_doc.doc[page_no]
                page_height = page.rect.height
                page_width = page.rect.width

                # OCR extraction
                ocr_result = ocr_engine.extract_with_image(
                    str(image_path),
                    page_no,
                    scale_to_pdf_points=72 / dpi,
                )
                tokens = ocr_result.tokens

                # Save preprocessed image
                _save_output_img(ocr_result.output_img, image_path)

                # Update dimensions to match OCR output
                if ocr_result.output_img is not None:
                    img_height, img_width = ocr_result.output_img.shape[:2]

                # Use shared processing logic (same as single document mode)
                matches = {}
                annotations, ann_count = process_page(
                    tokens=tokens,
                    row_dict=row_dict,
                    page_no=page_no,
                    page_height=page_height,
                    page_width=page_width,
                    img_width=img_width,
                    img_height=img_height,
                    dpi=dpi,
                    min_confidence=min_confidence,
                    matches=matches,
                    matched_fields=matched_fields,
                    report=report,
                    result_stats=result["stats"],
                )

                if annotations:
                    page_annotations.append(
                        {
                            "image_path": str(image_path),
                            "page_no": page_no,
                            "count": ann_count,
                        }
                    )
                    report.annotations_generated += ann_count

            # Record unmatched fields using shared logic
            record_unmatched_fields(row_dict, matched_fields, report)

            if page_annotations:
                result["pages"] = page_annotations
                result["success"] = True
                report.success = True
            else:
                report.errors.append("No annotations generated")

    except Exception as e:
        report.errors.append(str(e))

    report.processing_time_ms = (time.time() - start_time) * 1000
    result["report"] = report.to_dict()

    return result
