#!/usr/bin/env python3
"""Unified OCR processing task for DocuElevate.

This task replaces the single-provider ``process_with_azure_document_intelligence``
task with a multi-engine OCR pipeline that:

1. Runs every OCR provider listed in ``OCR_PROVIDERS`` (default: ``azure``).
2. Merges/cross-checks the results using the configured AI model when more
   than one provider is active (see ``OCR_MERGE_STRATEGY``).
3. Writes the best searchable PDF back to the working directory.
4. Optionally compares the OCR output against the original embedded text
   (passed as *original_text*) using a head-to-head AI review and keeps the
   higher-quality text for downstream processing.
5. Hands off to the page-rotation and metadata-extraction pipeline exactly as
   the legacy Azure task did.
"""

import logging
import os

from app.celery_app import celery
from app.config import settings
from app.database import SessionLocal
from app.models import FileRecord
from app.tasks.retry_config import OcrTaskWithRetry
from app.tasks.rotate_pdf_pages import rotate_pdf_pages
from app.utils import log_task_progress
from app.utils.ocr_provider import OCRResult, embed_text_layer, get_ocr_providers, merge_ocr_results
from app.utils.text_quality import TextSource, check_text_quality, compare_text_quality

logger = logging.getLogger(__name__)


@celery.task(base=OcrTaskWithRetry, bind=True)
def process_with_ocr(
    self,
    filename: str,
    file_id: int | None = None,
    original_text: str | None = None,
    language: str | None = None,
):
    """Run the configured OCR providers on *filename* and continue the pipeline.

    When multiple OCR providers are configured the results are merged using the
    AI model (or a simpler strategy controlled by ``OCR_MERGE_STRATEGY``).

    If *original_text* is provided (the original embedded text that failed the
    quality check), the OCR result is compared against it using a head-to-head
    AI review.  The higher-quality text is passed to downstream tasks.

    Args:
        filename: Base name of the file inside ``<workdir>/tmp/``.
        file_id: Optional database record ID passed through to downstream tasks.
        original_text: Optional original embedded text for head-to-head comparison.
        language: Optional Tesseract-style language code(s) (e.g. ``"eng+deu"``)
            to override the global OCR language settings for this specific run.
            Pass ``None`` or ``"auto"`` to use the global settings.  This
            enables per-pipeline language configuration.
    """
    task_id = self.request.id
    log_task_progress(
        task_id,
        "process_with_ocr",
        "in_progress",
        f"Starting OCR for {filename}",
        file_id=file_id,
    )

    try:
        tmp_file_path = os.path.join(settings.workdir, "tmp", filename)
        if not os.path.exists(tmp_file_path):
            raise FileNotFoundError(f"Local file not found: {tmp_file_path}")

        providers = get_ocr_providers(language=language)
        provider_names = [p.name for p in providers]
        logger.info(f"[{task_id}] Running {len(providers)} OCR provider(s): {provider_names}")

        log_task_progress(
            task_id,
            "run_ocr_providers",
            "in_progress",
            f"Running OCR providers: {', '.join(provider_names)}",
            file_id=file_id,
        )

        results = []
        errors = []
        for provider in providers:
            pname = provider.__class__.__name__
            try:
                logger.info(f"[{task_id}] Running {pname} on {filename}")
                result: OCRResult = provider.process(tmp_file_path)
                results.append(result)
                logger.info(f"[{task_id}] {pname} extracted {len(result.text)} chars")
            except Exception as exc:
                logger.error(f"[{task_id}] {pname} failed for {filename}: {exc}")
                errors.append(f"{pname}: {exc}")

        if not results:
            error_summary = "; ".join(errors)
            log_task_progress(
                task_id,
                "run_ocr_providers",
                "failure",
                "All OCR providers failed",
                file_id=file_id,
                detail=error_summary,
            )
            raise RuntimeError(f"All OCR providers failed for {filename}: {error_summary}")

        if errors:
            logger.warning(f"[{task_id}] Some OCR providers failed: {'; '.join(errors)}")

        log_task_progress(
            task_id,
            "run_ocr_providers",
            "success",
            f"{len(results)} of {len(providers)} OCR provider(s) succeeded",
            file_id=file_id,
        )

        # Merge results (no-op when only one provider succeeded)
        extracted_text, searchable_pdf_path, rotation_data = merge_ocr_results(results, filename)
        logger.info(
            f"[{task_id}] Merged OCR text: {len(extracted_text)} chars, "
            f"pdf={'yes' if searchable_pdf_path else 'no'}, "
            f"rotations={len(rotation_data)}"
        )

        # If no provider produced a searchable PDF, post-process the original
        # PDF with ocrmypdf to embed an invisible text layer so the output is
        # selectable/searchable in PDF viewers.
        if searchable_pdf_path is None:
            # Use the per-call language override; fall back to global setting
            embed_lang = (
                language
                if language and language != "auto"
                else (getattr(settings, "tesseract_language", None) or "eng")
            )
            log_task_progress(
                task_id,
                "embed_text_layer",
                "in_progress",
                "Embedding searchable text layer into PDF",
                file_id=file_id,
            )
            embedded = embed_text_layer(tmp_file_path, tmp_file_path, language=embed_lang)
            if embedded:
                searchable_pdf_path = tmp_file_path
                log_task_progress(
                    task_id,
                    "embed_text_layer",
                    "success",
                    "Searchable text layer embedded via ocrmypdf",
                    file_id=file_id,
                )
            else:
                log_task_progress(
                    task_id,
                    "embed_text_layer",
                    "skipped",
                    "ocrmypdf unavailable – PDF will not have a searchable text layer",
                    file_id=file_id,
                )

        # ----------------------------------------------------------------
        # Head-to-head comparison with original embedded text (if provided)
        # ----------------------------------------------------------------
        final_text = extracted_text
        if original_text and original_text.strip() and extracted_text.strip():
            logger.info(f"[{task_id}] Original embedded text provided; running head-to-head quality comparison")
            log_task_progress(
                task_id,
                "compare_ocr_quality",
                "in_progress",
                "Comparing OCR result against original embedded text",
                file_id=file_id,
            )
            try:
                comparison = compare_text_quality(original_text, extracted_text)
                comparison_detail = (
                    f"Original score: {comparison.original_score}/100, "
                    f"OCR score: {comparison.ocr_score}/100, "
                    f"Preferred: {comparison.preferred}\n"
                    f"AI explanation: {comparison.explanation}"
                )
                logger.info(f"[{task_id}] OCR comparison – {comparison_detail}")

                if comparison.preferred == "original":
                    # Original text is actually better – use it instead of OCR.
                    final_text = original_text
                    logger.info(
                        f"[{task_id}] Original embedded text selected "
                        f"(original={comparison.original_score} > ocr={comparison.ocr_score})"
                    )
                    log_task_progress(
                        task_id,
                        "compare_ocr_quality",
                        "success",
                        f"Original text preferred (original={comparison.original_score}/100 vs "
                        f"ocr={comparison.ocr_score}/100)",
                        file_id=file_id,
                        detail=comparison_detail,
                    )
                else:
                    logger.info(
                        f"[{task_id}] OCR text selected "
                        f"(preferred={comparison.preferred!r}, "
                        f"ocr={comparison.ocr_score}, original={comparison.original_score})"
                    )
                    log_task_progress(
                        task_id,
                        "compare_ocr_quality",
                        "success",
                        f"OCR text preferred (ocr={comparison.ocr_score}/100 vs "
                        f"original={comparison.original_score}/100)",
                        file_id=file_id,
                        detail=comparison_detail,
                    )
            except Exception as cmp_exc:
                logger.warning(f"[{task_id}] Head-to-head comparison failed ({cmp_exc}); keeping OCR text")
                log_task_progress(
                    task_id,
                    "compare_ocr_quality",
                    "skipped",
                    f"Comparison failed ({cmp_exc}); keeping OCR output",
                    file_id=file_id,
                )
        elif original_text is not None:
            # original_text was provided but one side is empty – pick whichever has content.
            if not extracted_text.strip() and original_text.strip():
                final_text = original_text
                logger.info(f"[{task_id}] OCR returned empty text; falling back to original embedded text")
                log_task_progress(
                    task_id,
                    "compare_ocr_quality",
                    "success",
                    "OCR empty – using original embedded text",
                    file_id=file_id,
                )
            else:
                log_task_progress(
                    task_id,
                    "compare_ocr_quality",
                    "skipped",
                    "No original text to compare; using OCR output",
                    file_id=file_id,
                )

        log_task_progress(
            task_id,
            "process_with_ocr",
            "success",
            f"OCR complete for {filename}",
            file_id=file_id,
            detail=f"Extracted {len(extracted_text)} chars using {len(results)} provider(s); "
            f"final text length: {len(final_text)} chars",
        )

        # Score the final embedded text — the text that will land in ocr_text.
        # We always call check_text_quality() on final_text because:
        #   - merge_ocr_results() may have AI-merged output from several engines
        #   - compare_text_quality() scores are relative (not the same scale)
        #   - The original may have been preferred, reversing the OCR output
        # OCR-produced (or AI-merged) text is treated as TextSource.OCR_PREVIOUS
        # so the quality AI call is always made.
        if file_id is not None:
            try:
                quality_result = check_text_quality(final_text, TextSource.OCR_PREVIOUS)
                logger.info(
                    f"[{task_id}] Final text quality: score={quality_result.quality_score}/100, "
                    f"good={quality_result.is_good_quality}, feedback={quality_result.feedback!r}"
                )
                with SessionLocal() as _db:
                    _rec = _db.query(FileRecord).filter_by(id=file_id).first()
                    if _rec:
                        _rec.ocr_quality_score = quality_result.quality_score
                        _db.commit()
                logger.info(f"[{task_id}] Saved ocr_quality_score={quality_result.quality_score} for file_id={file_id}")
            except Exception as _score_exc:
                logger.warning(f"[{task_id}] Could not persist ocr_quality_score: {_score_exc}")

        # Continue pipeline: rotate pages (if needed), then extract metadata
        rotate_pdf_pages.delay(filename, final_text, rotation_data, file_id)

        return {
            "file": filename,
            "searchable_pdf": searchable_pdf_path or tmp_file_path,
            "cleaned_text": final_text,
            "providers_used": [r.provider for r in results],
        }

    except Exception as exc:
        logger.error(f"[{task_id}] OCR failed for {filename}: {exc}")
        log_task_progress(
            task_id,
            "process_with_ocr",
            "failure",
            f"OCR failed for {filename}",
            file_id=file_id,
            detail=str(exc),
        )
        raise
