"""
OCR Worker Process Functions

Contains code that runs in the worker subprocess.
Handles OCR processing with persistent engine instances.

Key features:
- PaddleOCR initialized ONCE at process spawn
- docTR initialized ONCE at process spawn (PyTorch backend)
- Tesseract as fallback/complement engine
- Multi-pass preprocessing (light → medium → tesseract)
- Automatic engine selection based on results
- Memory cleanup after each task
"""

import gc
import io
import os
import time
from dataclasses import dataclass, field
from decimal import Decimal
from typing import Any, Dict, List, Optional, Tuple

import cv2
import numpy as np

# Disable PaddleOCR model source check for faster startup
os.environ['PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK'] = 'True'

# Memory optimization for docTR (prevents memory leak in multiprocessing)
# Source: https://github.com/mindee/doctr/issues/1594
os.environ['DOCTR_MULTIPROCESSING_DISABLE'] = 'TRUE'

# Reduce Intel oneDNN cache to save memory
os.environ['ONEDNN_PRIMITIVE_CACHE_CAPACITY'] = '1'


@dataclass
class OCRResult:
    """Raw OCR result from engine."""
    text: str
    confidence: float
    boxes: List[dict] = field(default_factory=list)
    engine: str = ""


def initialize_paddle_engine():
    """
    Initialize PaddleOCR engine.

    Called once at worker spawn. Returns the engine instance
    that will be reused for all subsequent requests.

    Returns:
        PaddleOCR instance or None if unavailable
    """
    try:
        print(f"[Worker {os.getpid()}] Loading PaddleOCR...", flush=True)
        start_time = time.time()

        from paddleocr import PaddleOCR

        # PaddleOCR 3.x API - optimized for Romanian receipts
        paddle = PaddleOCR(
            lang='en',  # 'en' handles Latin alphabet well for receipts
            det_db_thresh=0.3,
            det_db_box_thresh=0.5,
            det_db_unclip_ratio=1.8,
            rec_batch_num=6,
            use_angle_cls=True,
        )

        elapsed = time.time() - start_time
        print(f"[Worker {os.getpid()}] PaddleOCR loaded in {elapsed:.1f}s", flush=True)
        return paddle

    except Exception as e:
        print(f"[Worker {os.getpid()}] PaddleOCR init failed: {e}", flush=True)
        return None


def initialize_doctr_engine():
    """
    Initialize docTR engine (CPU only).

    Called once at worker spawn. Returns the engine instance
    that will be reused for all subsequent requests.

    Note: DirectML (AMD GPU) has compatibility issues with docTR.
    CUDA (NVIDIA) works but requires separate PyTorch build.
    CPU mode is stable and well-optimized.

    Returns:
        docTR predictor instance or None if unavailable
    """
    try:
        print(f"[Worker {os.getpid()}] Loading docTR (PyTorch backend, CPU)...", flush=True)
        start_time = time.time()

        from doctr.models import ocr_predictor

        # Initialize docTR predictor with pretrained models
        # Uses db_resnet50 for detection and crnn_vgg16_bn for recognition
        doctr = ocr_predictor(
            det_arch='db_resnet50',
            reco_arch='crnn_vgg16_bn',
            pretrained=True,
            assume_straight_pages=True,
            straighten_pages=False,
            preserve_aspect_ratio=True,
        )

        elapsed = time.time() - start_time
        print(f"[Worker {os.getpid()}] docTR loaded in {elapsed:.1f}s", flush=True)
        return doctr

    except Exception as e:
        print(f"[Worker {os.getpid()}] docTR init failed: {e}", flush=True)
        return None


def process_ocr(
    image_bytes: bytes,
    paddle_engine,
    tesseract_engine,
    engine: str = "doctr_plus",
    preprocessing: str = "auto",
    doctr_engine=None
) -> dict:
    """
    Process OCR on image bytes.

    Main entry point for OCR processing in worker process.
    Uses the specified engine for text recognition.

    Args:
        image_bytes: Raw image bytes (JPEG, PNG, or PDF)
        paddle_engine: Pre-initialized PaddleOCR instance (or None)
        tesseract_engine: Pre-initialized TesseractEngine instance (or None)
        engine: Engine selection ('tesseract', 'doctr', 'doctr_plus', 'paddleocr')
        preprocessing: Preprocessing mode ('auto', 'light', 'medium', 'heavy')
        doctr_engine: Pre-initialized docTR instance (or None)

    Returns:
        Dict with extraction results:
        {
            "success": bool,
            "extraction": {...},  # ExtractionResult as dict
            "raw_texts": [...],   # Raw OCR outputs
            "processing_time_ms": int,
            "ocr_engine": str
        }
    """
    import sys
    start_time = time.time()
    print(f"[Worker {os.getpid()}] Processing OCR: engine={engine}, preprocessing={preprocessing}, size={len(image_bytes)} bytes", flush=True)
    sys.stdout.flush()

    try:
        # Decode image from bytes
        print(f"[Worker {os.getpid()}] Decoding image...", flush=True)
        sys.stdout.flush()
        image = _decode_image(image_bytes)
        if image is None:
            return {"success": False, "error": "Failed to decode image"}
        print(f"[Worker {os.getpid()}] Image decoded: shape={image.shape}, dtype={image.dtype}, size={image.nbytes/1024/1024:.1f}MB", flush=True)
        sys.stdout.flush()

        # Import preprocessor
        from backend.modules.data_entry.services.image_preprocessor import ImagePreprocessor
        from backend.modules.data_entry.services.ocr_extractor import ReceiptExtractor

        preprocessor = ImagePreprocessor()
        extractor = ReceiptExtractor()
        print(f"[Worker {os.getpid()}] Preprocessor and extractor initialized", flush=True)
        sys.stdout.flush()

        raw_texts = []
        extraction = None

        # Engine routing (available: tesseract, doctr, doctr_plus, paddleocr)
        print(f"[Worker {os.getpid()}] Routing to engine: {engine}", flush=True)
        sys.stdout.flush()
        if engine == "tesseract":
            extraction, raw_texts = _process_tesseract_only(
                image, tesseract_engine, preprocessor, extractor
            )
        elif engine == "doctr":
            extraction, raw_texts = _process_doctr_only(
                image, doctr_engine, preprocessor, extractor
            )
        elif engine == "doctr_plus":
            extraction, raw_texts = _process_doctr_plus(
                image, doctr_engine, preprocessor, extractor
            )
        elif engine == "paddleocr":
            extraction, raw_texts = _process_paddleocr_only(
                image, paddle_engine, preprocessor, extractor
            )
        else:
            # Default to doctr_plus if unknown engine specified
            print(f"[OCR] Unknown engine '{engine}', defaulting to doctr_plus", flush=True)
            extraction, raw_texts = _process_doctr_plus(
                image, doctr_engine, preprocessor, extractor
            )

        # Calculate processing time
        elapsed_ms = int((time.time() - start_time) * 1000)

        if extraction:
            extraction.processing_time_ms = elapsed_ms

        # Convert extraction to dict for serialization
        result = {
            "success": extraction is not None,
            "extraction": _extraction_to_dict(extraction) if extraction else None,
            "raw_texts": raw_texts,
            "processing_time_ms": elapsed_ms,
            "ocr_engine": extraction.ocr_engine if extraction else "none",
            "pid": os.getpid()
        }

        print(f"[Worker {os.getpid()}] OCR complete in {elapsed_ms}ms", flush=True)
        return result

    except Exception as e:
        elapsed_ms = int((time.time() - start_time) * 1000)
        print(f"[Worker {os.getpid()}] OCR error after {elapsed_ms}ms: {e}", flush=True)
        import traceback
        traceback.print_exc()
        return {
            "success": False,
            "error": str(e),
            "processing_time_ms": elapsed_ms,
            "pid": os.getpid()
        }

    finally:
        # Cleanup memory
        gc.collect()


def _decode_image(image_bytes: bytes) -> Optional[np.ndarray]:
    """Decode image from bytes (JPEG, PNG, or first page of PDF).

    For PDFs, uses 200 DPI which is sufficient for receipt OCR
    and reduces processing time by ~50% vs 300 DPI.
    """
    try:
        # Try as regular image first
        nparr = np.frombuffer(image_bytes, np.uint8)
        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

        if image is not None:
            return image

        # Try as PDF - use 200 DPI for faster processing (sufficient for receipts)
        try:
            import pdf2image
            from PIL import Image

            # 200 DPI is sufficient for receipt text recognition
            # 300 DPI was overkill and slowed down processing
            images = pdf2image.convert_from_bytes(image_bytes, dpi=200)
            if images:
                # Convert first page to numpy array
                pil_img = images[0]
                print(f"[Worker {os.getpid()}] PDF decoded: {pil_img.width}x{pil_img.height} @ 200 DPI", flush=True)
                return np.array(pil_img)
        except Exception as e:
            print(f"[Worker {os.getpid()}] PDF decode error: {e}", flush=True)

        return None

    except Exception as e:
        print(f"[Worker {os.getpid()}] Image decode error: {e}", flush=True)
        return None


def _process_paddleocr_only(
    image: np.ndarray,
    paddle_engine,
    preprocessor,
    extractor
) -> Tuple[Any, List[str]]:
    """Process using PaddleOCR only (light + medium preprocessing)."""
    raw_texts = []
    extraction = None

    if paddle_engine is None:
        return None, ["PaddleOCR not available"]

    # Step 1: Light preprocessing
    print("[OCR] Step 1: PaddleOCR + Light", flush=True)
    light_img = preprocessor.preprocess_light(image)
    paddle_light = _paddle_recognize(paddle_engine, light_img)

    if paddle_light and paddle_light.text:
        extraction = extractor.extract(paddle_light.text)
        extraction.ocr_engine = "paddle-light"
        raw_texts.append(f"=== PaddleOCR Light (conf: {paddle_light.confidence:.0%}) ===\n{paddle_light.text}")

        if _is_extraction_complete(extraction):
            return extraction, raw_texts

    # Step 2: Medium preprocessing
    print("[OCR] Step 2: PaddleOCR + Medium", flush=True)
    medium_img = preprocessor.preprocess_medium(image)
    paddle_medium = _paddle_recognize(paddle_engine, medium_img)

    if paddle_medium and paddle_medium.text:
        extraction_medium = extractor.extract(paddle_medium.text)
        extraction_medium.ocr_engine = "paddle-medium"
        raw_texts.append(f"=== PaddleOCR Medium (conf: {paddle_medium.confidence:.0%}) ===\n{paddle_medium.text}")

        if extraction:
            extraction = _merge_extractions(extraction, extraction_medium)
            extraction.ocr_engine = "paddle-adaptive"
        else:
            extraction = extraction_medium

    return extraction, raw_texts


def _process_tesseract_only(
    image: np.ndarray,
    tesseract_engine,
    preprocessor,
    extractor
) -> Tuple[Any, List[str]]:
    """Process using Tesseract only with optimized preprocessing."""
    raw_texts = []
    extraction = None

    if tesseract_engine is None:
        return None, ["Tesseract not available"]

    print("[OCR] Tesseract-only mode", flush=True)
    tesseract_img = preprocessor.preprocess_for_tesseract(image)
    tesseract_result = tesseract_engine.recognize(tesseract_img)

    if tesseract_result and tesseract_result.text:
        extraction = extractor.extract(tesseract_result.text)
        extraction.ocr_engine = "tesseract"
        raw_texts.append(f"=== Tesseract (conf: {tesseract_result.confidence:.0%}) ===\n{tesseract_result.text}")

    return extraction, raw_texts


def _process_doctr_only(
    image: np.ndarray,
    doctr_engine,
    preprocessor,
    extractor
) -> Tuple[Any, List[str]]:
    """
    Process using docTR only (light + medium preprocessing).

    docTR uses EXACT same preprocessing as PaddleOCR for consistency.
    """
    raw_texts = []
    extraction = None

    if doctr_engine is None:
        return None, ["docTR not available"]

    # Step 1: Light preprocessing (same as PaddleOCR)
    print("[OCR] Step 1: docTR + Light", flush=True)
    light_img = preprocessor.preprocess_light(image)
    doctr_light = _doctr_recognize(doctr_engine, light_img)

    if doctr_light and doctr_light.text:
        extraction = extractor.extract(doctr_light.text)
        extraction.ocr_engine = "doctr-light"
        raw_texts.append(f"=== docTR Light (conf: {doctr_light.confidence:.0%}) ===\n{doctr_light.text}")

        if _is_extraction_complete(extraction):
            return extraction, raw_texts

    # Step 2: Medium preprocessing (same as PaddleOCR)
    print("[OCR] Step 2: docTR + Medium", flush=True)
    medium_img = preprocessor.preprocess_medium(image)
    doctr_medium = _doctr_recognize(doctr_engine, medium_img)

    if doctr_medium and doctr_medium.text:
        extraction_medium = extractor.extract(doctr_medium.text)
        extraction_medium.ocr_engine = "doctr-medium"
        raw_texts.append(f"=== docTR Medium (conf: {doctr_medium.confidence:.0%}) ===\n{doctr_medium.text}")

        if extraction:
            extraction = _merge_extractions(extraction, extraction_medium)
            extraction.ocr_engine = "doctr-adaptive"
        else:
            extraction = extraction_medium

    return extraction, raw_texts


def _process_doctr_plus(
    image: np.ndarray,
    doctr_engine,
    preprocessor,
    extractor
) -> Tuple[Any, List[str]]:
    """
    docTR Plus - Optimized 2-tier sequential processing with early exit.

    Architecture:
    - Tier 1: Light preprocessing (~4-5s)
      → Early exit if confidence >= 0.75 AND all fields valid AND cross-validations pass
    - Tier 2: Medium preprocessing (only if Tier 1 insufficient, ~4-5s additional)
      → Merge with Tier 1 results
      → Mark for review if still problems

    Performance:
    - Fast path (80% receipts): ~4-5s (Tier 1 only)
    - Slow path (20% receipts): ~8-9s (Tier 1 + Tier 2)
    - Average: ~5-6s

    Returns:
        Tuple of (extraction_result, raw_texts_list)
        extraction_result.needs_review = True if validation issues remain
    """
    raw_texts = []
    extraction = None

    if doctr_engine is None:
        return None, ["docTR not available"]

    # ========== TIER 1: Light Preprocessing ==========
    print("[docTR+] TIER 1: Light preprocessing", flush=True)
    import time
    tier1_start = time.time()

    light_img = preprocessor.preprocess_light(image)
    doctr_light = _doctr_recognize(doctr_engine, light_img)

    tier1_time = time.time() - tier1_start
    print(f"[docTR+] TIER 1 completed in {tier1_time:.1f}s", flush=True)

    if doctr_light and doctr_light.text:
        extraction = extractor.extract(doctr_light.text)
        extraction.ocr_engine = "doctr-plus-light"
        raw_texts.append(f"=== docTR+ Tier1/Light (conf: {doctr_light.confidence:.0%}) ===\n{doctr_light.text}")

        # Early Exit Check: confidence >= 0.75 + cross-validations
        if _is_extraction_valid_for_early_exit(extraction, min_confidence=0.75):
            print(f"[docTR+] EARLY EXIT - Tier 1 sufficient (conf: {extraction.overall_confidence:.0%})", flush=True)
            extraction.ocr_engine = "doctr-plus"
            return extraction, raw_texts

        print(f"[docTR+] Tier 1 incomplete or validation failed, proceeding to Tier 2...", flush=True)

    # ========== TIER 2: Medium Preprocessing (only if needed) ==========
    print("[docTR+] TIER 2: Medium preprocessing", flush=True)
    tier2_start = time.time()

    medium_img = preprocessor.preprocess_medium(image)
    doctr_medium = _doctr_recognize(doctr_engine, medium_img)

    tier2_time = time.time() - tier2_start
    print(f"[docTR+] TIER 2 completed in {tier2_time:.1f}s", flush=True)

    if doctr_medium and doctr_medium.text:
        extraction_medium = extractor.extract(doctr_medium.text)
        extraction_medium.ocr_engine = "doctr-plus-medium"
        raw_texts.append(f"=== docTR+ Tier2/Medium (conf: {doctr_medium.confidence:.0%}) ===\n{doctr_medium.text}")

        if extraction:
            # Merge Tier 1 + Tier 2 results
            extraction = _merge_extractions(extraction, extraction_medium)
        else:
            extraction = extraction_medium

    # ========== FINAL VALIDATION ==========
    if extraction:
        extraction.ocr_engine = "doctr-plus"

        # Mark for review if validation still fails after both tiers
        passes_validation, penalty, errors = _quick_cross_validate(extraction)

        if not passes_validation or extraction.overall_confidence < 0.75:
            # Mark for human review using existing fields
            extraction.needs_manual_review = True

            if extraction.overall_confidence < 0.75:
                extraction.validation_warnings.append(f"Low confidence: {extraction.overall_confidence:.0%}")

            if not extraction.amount:
                extraction.validation_errors.append("TOTAL not detected")
            if not extraction.cui:
                extraction.validation_warnings.append("CUI not detected")
            if not extraction.tva_total and not extraction.tva_entries:
                extraction.validation_warnings.append("TVA not detected")
            if not extraction.receipt_date:
                extraction.validation_warnings.append("Date not detected")

            # Add cross-validation errors
            extraction.validation_errors.extend(errors)

            print(f"[docTR+] Marked for review: {extraction.validation_errors + extraction.validation_warnings}", flush=True)
        else:
            extraction.needs_manual_review = False

    total_time = tier1_time + (tier2_time if 'tier2_time' in dir() else 0)
    print(f"[docTR+] Total processing time: {total_time:.1f}s", flush=True)

    return extraction, raw_texts


# =============================================================================
# VALIDATION HELPERS (used by doctr_plus for early exit decisions)
# =============================================================================

def _quick_cross_validate(extraction) -> tuple[bool, float, list[str]]:
    """
    Quick cross-validation for OCR results.

    Checks critical field correlations to detect obvious OCR errors.
    Used by doctr_plus to decide whether to proceed to Tier 2 or exit early.

    Returns:
        Tuple of (passes_validation, confidence_penalty, error_messages)
    """
    try:
        from backend.modules.data_entry.services.ocr.validation import OCRValidationEngine

        if extraction is None:
            return False, 1.0, ["No extraction result"]

        # Convert extraction to dict for validation
        # Build TVA entries dict for TVAEntriesSumRule (expects {code: amount})
        tva_entries_dict = {}
        if extraction.tva_entries:
            for entry in extraction.tva_entries:
                if isinstance(entry, dict):
                    code = entry.get('code', 'A')
                    amount = entry.get('amount', 0)
                    try:
                        tva_entries_dict[code] = float(amount)
                    except (TypeError, ValueError):
                        pass

        validation_data = {
            "amount": float(extraction.amount) if extraction.amount else None,
            "tva": float(extraction.tva_total) if extraction.tva_total else None,
            "tva_entries": tva_entries_dict,  # For TVAEntriesSumRule: {code: amount}
            "cui": extraction.cui,  # For CUI checksum validation
        }

        # Also pass raw tva_entries for TVABasedTotalRule (for rate detection)
        if extraction.tva_entries:
            validation_data['tva_entries_raw'] = extraction.tva_entries

        # Add payment methods if available (for TOTAL vs CARD+CASH validation)
        if extraction.payment_methods:
            try:
                card_amount = sum(
                    float(p.get('amount', 0) if isinstance(p, dict) else 0)
                    for p in extraction.payment_methods
                    if isinstance(p, dict) and p.get('method') == 'CARD'
                )
                cash_amount = sum(
                    float(p.get('amount', 0) if isinstance(p, dict) else 0)
                    for p in extraction.payment_methods
                    if isinstance(p, dict) and p.get('method') == 'NUMERAR'
                )
                validation_data['card_amount'] = card_amount
                validation_data['cash_amount'] = cash_amount
            except Exception as e:
                print(f"[Worker {os.getpid()}] Payment method validation error: {e}", flush=True)

        # Run quick validation
        validator = OCRValidationEngine()
        return validator.quick_validate_for_hybrid(validation_data)

    except Exception as e:
        # Never crash the process on validation errors
        print(f"[Worker {os.getpid()}] Cross-validation error: {e}", flush=True)
        import traceback
        traceback.print_exc()
        # Return "passes" to allow processing to continue
        return True, 0.0, [f"Validation skipped due to error: {str(e)}"]


def _is_extraction_valid_for_early_exit(extraction, min_confidence: float = 0.85) -> bool:
    """
    Check if extraction is valid for early exit in doctr_plus.

    Combines confidence check with cross-validation to prevent
    early exit on OCR errors (e.g., wrong TOTAL but correct TVA).

    Returns:
        True only if:
        1. Overall confidence >= min_confidence
        2. Critical fields are present (AMOUNT, DATE, CUI)
        3. Cross-validation passes (TOTAL matches TVA calculation, or no TVA)
    """
    try:
        # First check basic completeness (relaxed for early exit)
        if not _is_extraction_complete(extraction, min_confidence, for_early_exit=True):
            return False

        # Then run cross-validation
        passes_validation, penalty, errors = _quick_cross_validate(extraction)

        if not passes_validation:
            print(f"[Early Exit] BLOCKED: cross-validation failed: {errors}", flush=True)
            return False

        print(f"[Early Exit] OK: conf={extraction.overall_confidence:.0%}, validation passed", flush=True)
        return True

    except Exception as e:
        # Never crash on validation - just continue to next engine
        print(f"[Worker {os.getpid()}] Early exit check error: {e}", flush=True)
        return False  # Continue to next engine on error

def _paddle_recognize(paddle_engine, image: np.ndarray) -> Optional[OCRResult]:
    """Run PaddleOCR recognition on image."""
    try:
        # Ensure 3-channel image
        if len(image.shape) == 2:
            image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)

        result = paddle_engine.predict(image, use_textline_orientation=True)

        if not result or len(result) == 0:
            return OCRResult(text="", confidence=0.0, boxes=[], engine="paddleocr")

        ocr_result = result[0]
        rec_texts = ocr_result.get('rec_texts', [])
        rec_scores = ocr_result.get('rec_scores', [])
        dt_polys = ocr_result.get('dt_polys', [])

        if not rec_texts:
            return OCRResult(text="", confidence=0.0, boxes=[], engine="paddleocr")

        boxes = []
        for i, text in enumerate(rec_texts):
            conf = rec_scores[i] if i < len(rec_scores) else 0.0
            box = dt_polys[i].tolist() if i < len(dt_polys) else []
            boxes.append({'text': text, 'confidence': float(conf), 'box': box})

        avg_conf = sum(rec_scores) / len(rec_scores) if rec_scores else 0.0
        text_result = '\n'.join(rec_texts)

        return OCRResult(
            text=text_result,
            confidence=float(avg_conf),
            boxes=boxes,
            engine="paddleocr"
        )

    except Exception as e:
        print(f"[Worker] PaddleOCR error: {e}", flush=True)
        return None


def _doctr_recognize(doctr_engine, image: np.ndarray) -> Optional[OCRResult]:
    """
    Run docTR recognition on image.

    docTR requires RGB images, handles conversion automatically.
    Uses same preprocessing as PaddleOCR for consistent results.
    """
    try:
        # docTR requires RGB images
        if len(image.shape) == 2:
            # Convert grayscale to RGB
            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
        elif image.shape[2] == 3:
            # Convert BGR (OpenCV) to RGB
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        elif image.shape[2] == 4:
            # Convert RGBA to RGB
            image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)

        # docTR expects a list of numpy arrays (pages)
        result = doctr_engine([image])

        if not result or not result.pages:
            return OCRResult(text="", confidence=0.0, boxes=[], engine="doctr")

        # Extract text from all pages
        all_texts = []
        all_confidences = []
        boxes = []

        for page in result.pages:
            for block in page.blocks:
                for line in block.lines:
                    line_text = ' '.join(word.value for word in line.words)
                    line_confidence = sum(w.confidence for w in line.words) / len(line.words) if line.words else 0.0
                    all_texts.append(line_text)
                    all_confidences.append(line_confidence)

                    # Store word-level boxes
                    for word in line.words:
                        boxes.append({
                            'text': word.value,
                            'confidence': float(word.confidence),
                            'box': word.geometry  # (xmin, ymin), (xmax, ymax)
                        })

        text_result = '\n'.join(all_texts)
        avg_conf = sum(all_confidences) / len(all_confidences) if all_confidences else 0.0

        return OCRResult(
            text=text_result,
            confidence=float(avg_conf),
            boxes=boxes,
            engine="doctr"
        )

    except Exception as e:
        print(f"[Worker] docTR error: {e}", flush=True)
        return None


def _is_extraction_complete(ext, min_confidence: float = 0.85, for_early_exit: bool = False) -> bool:
    """
    Check if extraction has required fields.

    Args:
        ext: Extraction result
        min_confidence: Minimum overall confidence
        for_early_exit: If True, use relaxed criteria (AMOUNT + DATE + CUI required)
                       If False, require all fields (strict mode for final validation)

    Returns:
        True if extraction meets completeness criteria
    """
    # Check confidence first
    if ext.overall_confidence < min_confidence:
        if for_early_exit:
            print(f"[Early Exit] BLOCKED: confidence {ext.overall_confidence:.0%} < {min_confidence:.0%}", flush=True)
        return False

    has_date = bool(ext.receipt_date)
    has_amount = bool(ext.amount)
    has_tva = bool(ext.tva_total) or bool(ext.tva_entries)
    has_cui = bool(ext.cui)

    if for_early_exit:
        # Relaxed criteria for early exit:
        # - AMOUNT is required (core field)
        # - DATE is required (needed for accounting)
        # - CUI is required (needed for supplier identification)
        # - TVA is NOT required (some receipts have 0% TVA)
        # - receipt_number is NOT required (often missing)
        required_ok = all([has_amount, has_date, has_cui])

        if not required_ok:
            missing = []
            if not has_amount: missing.append("AMOUNT")
            if not has_date: missing.append("DATE")
            if not has_cui: missing.append("CUI")
            print(f"[Early Exit] BLOCKED: missing required fields: {', '.join(missing)}", flush=True)

        return required_ok
    else:
        # Strict criteria for final validation (all fields required)
        has_number = bool(ext.receipt_number)
        return all([has_number, has_date, has_amount, has_tva, has_cui])


def _merge_extractions(primary, secondary):
    """Merge two extractions, picking best fields from each.

    Primary should be the higher-quality engine (e.g., docTR).
    Secondary is the fallback engine (e.g., Tesseract).

    Priority logic:
    - AMOUNT: TVA validation wins over confidence. If both valid or both invalid,
      uses confidence (or TVA diff for invalid cases).
    - DATE/CUI: Validation-based, then confidence, then primary wins ties.
    - OTHER FIELDS: Primary wins when both have values.
    """
    from backend.modules.data_entry.services.ocr_extractor import ExtractionResult

    result = ExtractionResult()

    # Helper: Check if amount matches TVA calculation
    def amount_passes_tva_validation(amount, tva_total, tva_entries):
        if not amount or not tva_total:
            return False, 0.0
        try:
            tva_rate = 0.21  # Default Romanian TVA
            if tva_entries:
                for entry in tva_entries:
                    if isinstance(entry, dict) and entry.get('percent'):
                        tva_rate = float(entry['percent']) / 100.0
                        break
            # Expected TOTAL = TVA / rate * (1 + rate)
            expected = float(tva_total) * (1 + tva_rate) / tva_rate
            actual = float(amount)
            diff_percent = abs(actual - expected) / expected if expected > 0 else 1.0
            return diff_percent < 0.03, diff_percent  # 3% tolerance
        except:
            return False, 1.0

    # Amount - prefer TVA-validated value over confidence
    if primary.amount and secondary.amount:
        # Get TVA from the one with entries, or use any available
        tva_total = primary.tva_total or secondary.tva_total
        tva_entries = primary.tva_entries or secondary.tva_entries

        primary_valid, primary_diff = amount_passes_tva_validation(
            primary.amount, tva_total, tva_entries
        )
        secondary_valid, secondary_diff = amount_passes_tva_validation(
            secondary.amount, tva_total, tva_entries
        )

        print(f"[Merge] Amount comparison: primary={primary.amount} (valid={primary_valid}, diff={primary_diff:.1%}), "
              f"secondary={secondary.amount} (valid={secondary_valid}, diff={secondary_diff:.1%})", flush=True)

        if secondary_valid and not primary_valid:
            # Secondary passes validation, primary doesn't - use secondary!
            print(f"[Merge] Using secondary amount {secondary.amount} (passes TVA validation)", flush=True)
            result.amount = secondary.amount
            result.confidence_amount = secondary.confidence_amount
        elif primary_valid and not secondary_valid:
            # Primary passes validation
            result.amount = primary.amount
            result.confidence_amount = primary.confidence_amount
        elif primary_valid and secondary_valid:
            # Both valid - use higher confidence
            if primary.confidence_amount >= secondary.confidence_amount:
                result.amount = primary.amount
                result.confidence_amount = primary.confidence_amount
            else:
                result.amount = secondary.amount
                result.confidence_amount = secondary.confidence_amount
        else:
            # Neither valid - use the one closer to TVA calculation
            if secondary_diff < primary_diff:
                print(f"[Merge] Neither valid, using secondary {secondary.amount} (closer to TVA)", flush=True)
                result.amount = secondary.amount
                result.confidence_amount = secondary.confidence_amount
            else:
                result.amount = primary.amount
                result.confidence_amount = primary.confidence_amount
    elif primary.amount:
        result.amount = primary.amount
        result.confidence_amount = primary.confidence_amount
    elif secondary.amount:
        result.amount = secondary.amount
        result.confidence_amount = secondary.confidence_amount

    # Date - prefer higher confidence
    if primary.receipt_date and secondary.receipt_date:
        if primary.confidence_date >= secondary.confidence_date:
            result.receipt_date = primary.receipt_date
            result.confidence_date = primary.confidence_date
        else:
            result.receipt_date = secondary.receipt_date
            result.confidence_date = secondary.confidence_date
    elif primary.receipt_date:
        result.receipt_date = primary.receipt_date
        result.confidence_date = primary.confidence_date
    elif secondary.receipt_date:
        result.receipt_date = secondary.receipt_date
        result.confidence_date = secondary.confidence_date

    # CUI - prefer valid format and version with RO prefix
    # Use CUIChecksumRule static methods (single source of truth)
    from backend.modules.data_entry.services.ocr.validation import CUIChecksumRule

    def is_valid_cui(cui):
        if not cui:
            return False
        digits = CUIChecksumRule.extract_digits(cui)
        return len(digits) >= 6 and len(digits) <= 10

    if primary.cui and secondary.cui:
        if is_valid_cui(primary.cui) and not is_valid_cui(secondary.cui):
            result.cui = primary.cui
        elif is_valid_cui(secondary.cui) and not is_valid_cui(primary.cui):
            result.cui = secondary.cui
        else:
            # Both valid - prefer the one with RO prefix if digits match
            primary_digits = CUIChecksumRule.extract_digits(primary.cui)
            secondary_digits = CUIChecksumRule.extract_digits(secondary.cui)
            if primary_digits == secondary_digits:
                if CUIChecksumRule.has_ro_prefix(secondary.cui) and not CUIChecksumRule.has_ro_prefix(primary.cui):
                    result.cui = secondary.cui  # Prefer version with RO
                    print(f"[CUI Complement] Preferring secondary with RO: {secondary.cui}", flush=True)
                else:
                    result.cui = primary.cui
            else:
                result.cui = primary.cui
    elif primary.cui:
        result.cui = primary.cui
    elif secondary.cui:
        result.cui = secondary.cui

    # TVA entries - ALWAYS prefer primary (docTR) when both have entries
    if primary.tva_entries and secondary.tva_entries:
        # Always use primary (docTR) - higher quality OCR
        result.tva_entries = primary.tva_entries
        result.tva_total = primary.tva_total
        result.confidence_tva = getattr(primary, 'confidence_tva', 0.0)
    elif primary.tva_entries:
        result.tva_entries = primary.tva_entries
        result.tva_total = primary.tva_total
        result.confidence_tva = getattr(primary, 'confidence_tva', 0.0)
    elif secondary.tva_entries:
        result.tva_entries = secondary.tva_entries
        result.tva_total = secondary.tva_total
        result.confidence_tva = getattr(secondary, 'confidence_tva', 0.0)

    # Other fields - prefer primary
    result.receipt_number = primary.receipt_number or secondary.receipt_number
    result.receipt_series = primary.receipt_series or secondary.receipt_series
    result.receipt_type = primary.receipt_type or secondary.receipt_type
    result.partner_name = primary.partner_name or secondary.partner_name
    result.address = primary.address or secondary.address
    result.items_count = primary.items_count or secondary.items_count
    # Payment methods with confidence
    if primary.payment_methods:
        result.payment_methods = primary.payment_methods
        result.confidence_payment = getattr(primary, 'confidence_payment', 0.0)
    elif secondary.payment_methods:
        result.payment_methods = secondary.payment_methods
        result.confidence_payment = getattr(secondary, 'confidence_payment', 0.0)
    result.suggested_payment_mode = getattr(primary, 'suggested_payment_mode', None) or getattr(secondary, 'suggested_payment_mode', None)

    # Client fields
    result.client_name = primary.client_name or secondary.client_name
    result.client_cui = primary.client_cui or secondary.client_cui
    result.client_address = primary.client_address or secondary.client_address

    # Confidence fields - preserve from primary or pick best
    if primary.confidence_vendor >= secondary.confidence_vendor:
        result.confidence_vendor = primary.confidence_vendor
    else:
        result.confidence_vendor = secondary.confidence_vendor

    if hasattr(primary, 'confidence_client') and hasattr(secondary, 'confidence_client'):
        if primary.confidence_client >= secondary.confidence_client:
            result.confidence_client = primary.confidence_client
        else:
            result.confidence_client = secondary.confidence_client

    # Raw text - combine both for debugging/display
    raw_texts = []
    if primary.raw_text:
        raw_texts.append(primary.raw_text)
    if secondary.raw_text and secondary.raw_text != primary.raw_text:
        raw_texts.append(secondary.raw_text)
    result.raw_text = '\n---\n'.join(raw_texts) if raw_texts else ''

    # Note: overall_confidence is a computed @property on ExtractionResult
    # It automatically calculates from confidence_amount, confidence_date, confidence_vendor

    return result


def _complement_extraction(primary, secondary):
    """Complement primary with missing fields from secondary (never overrides)."""
    # Only fill missing fields
    if not primary.amount and secondary.amount:
        primary.amount = secondary.amount
        primary.confidence_amount = secondary.confidence_amount

    if not primary.receipt_date and secondary.receipt_date:
        primary.receipt_date = secondary.receipt_date
        primary.confidence_date = secondary.confidence_date

    if not primary.partner_name and secondary.partner_name:
        primary.partner_name = secondary.partner_name
        primary.confidence_vendor = secondary.confidence_vendor

    if not primary.cui and secondary.cui:
        primary.cui = secondary.cui

    if not primary.tva_entries and secondary.tva_entries:
        primary.tva_entries = secondary.tva_entries
        primary.tva_total = secondary.tva_total
        primary.confidence_tva = getattr(secondary, 'confidence_tva', 0.0)

    if not getattr(primary, 'payment_methods', None) and getattr(secondary, 'payment_methods', None):
        primary.payment_methods = secondary.payment_methods
        primary.confidence_payment = getattr(secondary, 'confidence_payment', 0.0)

    if not primary.receipt_number and secondary.receipt_number:
        primary.receipt_number = secondary.receipt_number

    if not primary.address and secondary.address:
        primary.address = secondary.address

    if not primary.client_name and secondary.client_name:
        primary.client_name = secondary.client_name
        primary.client_cui = secondary.client_cui
        primary.client_address = secondary.client_address

    return primary


def _extraction_to_dict(extraction) -> dict:
    """Convert ExtractionResult to serializable dict."""
    if extraction is None:
        return None

    def safe_decimal(val):
        if val is None:
            return None
        return float(val) if isinstance(val, Decimal) else val

    def safe_date(val):
        if val is None:
            return None
        return val.isoformat() if hasattr(val, 'isoformat') else str(val)

    return {
        "receipt_type": extraction.receipt_type,
        "receipt_number": extraction.receipt_number,
        "receipt_series": extraction.receipt_series,
        "receipt_date": safe_date(extraction.receipt_date),
        "amount": safe_decimal(extraction.amount),
        "partner_name": extraction.partner_name,
        "cui": extraction.cui,
        "description": extraction.description,
        "tva_entries": extraction.tva_entries,
        "tva_total": safe_decimal(extraction.tva_total),
        "address": extraction.address,
        "items_count": extraction.items_count,
        "payment_methods": extraction.payment_methods,
        "suggested_payment_mode": getattr(extraction, 'suggested_payment_mode', None),
        # Client data
        "client_name": extraction.client_name,
        "client_cui": extraction.client_cui,
        "client_address": extraction.client_address,
        # Confidence scores
        "confidence_amount": extraction.confidence_amount,
        "confidence_date": extraction.confidence_date,
        "confidence_vendor": extraction.confidence_vendor,
        "confidence_client": getattr(extraction, 'confidence_client', 0.0),
        "confidence_tva": getattr(extraction, 'confidence_tva', 0.0),
        "confidence_payment": getattr(extraction, 'confidence_payment', 0.0),
        "overall_confidence": extraction.overall_confidence,
        "raw_text": extraction.raw_text,
        "ocr_engine": extraction.ocr_engine,
        "processing_time_ms": extraction.processing_time_ms,
        # Validation (if present)
        "needs_manual_review": getattr(extraction, 'needs_manual_review', None),
        "validation_warnings": getattr(extraction, 'validation_warnings', []),
        "validation_errors": getattr(extraction, 'validation_errors', []),
        "inter_ocr_ratios": getattr(extraction, 'inter_ocr_ratios', {}),
    }
