src/extractor.py

import os
import cv2
import numpy as np
import easyocr
import warnings
import ssl
import re
from passporteye import read_mrz
from pdf2image import convert_from_path
from PIL import Image
import string as st

from src.utils import (
    clean_string,
    clean_mrz_line,
    parse_date,
    get_country_name,
    get_sex,
    setup_logger,
    clean_name_field
)

from src.fallback_mrz import FallbackMRZ
from config.settings import USE_GPU, OCR_LANGUAGES, TEMP_DIR

warnings.filterwarnings("ignore")
logger = setup_logger(__name__)


# Fix SSL issue (Mac EasyOCR model download fix)
try:
    _create_unverified_https_context = ssl._create_unverified_context
except AttributeError:
    pass
else:
    ssl._create_default_https_context = _create_unverified_https_context


class PassportExtractor:

    def __init__(self, use_gpu=USE_GPU, languages=None):
        self.languages = languages if languages else OCR_LANGUAGES

        base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        model_dir = os.path.join(base_dir, "data", "models")
        os.makedirs(model_dir, exist_ok=True)

        logger.debug(f"Initializing EasyOCR Reader (GPU={use_gpu})...")
        
        # Create custom user network directory to avoid permission issues
        user_network_dir = os.path.join(base_dir, "data", "easyocr_user_network")
        os.makedirs(user_network_dir, exist_ok=True)
        
        self.reader = easyocr.Reader(
            self.languages,
            gpu=use_gpu,
            model_storage_directory=model_dir,
            user_network_directory=user_network_dir
        )
        logger.debug("EasyOCR initialized.")

    # ---------------------------------------------------
    # VISUAL GIVEN NAME EXTRACTION (PRIMARY SOURCE)
    # ---------------------------------------------------
    def extract_given_names_from_visual(self, img_path):
        try:
            results = self.reader.readtext(img_path, detail=0)
            lines = [r.strip() for r in results if r.strip()]

            for i, line in enumerate(lines):
                upper_line = line.upper()

                if "GIVEN" in upper_line and "NAME" in upper_line:

                    if ":" in line:
                        candidate = line.split(":")[1].strip()
                    else:
                        if i + 1 < len(lines):
                            candidate = lines[i + 1].strip()
                        else:
                            return ""

                    # Keep only letters and spaces, but preserve spaces between names
                    candidate = re.sub(r'[^A-Za-z\s]', '', candidate)
                    
                    # Clean up extra spaces but keep single spaces between names
                    candidate = re.sub(r'\s+', ' ', candidate).strip()

                    # Remove trailing single letter only if it's clearly an OCR artifact (not part of a name)
                    # This is more conservative - only removes single letters that are likely OCR errors
                    candidate = re.sub(r'([A-Z]{2,})[K]$', r'\1', candidate)  # K is common OCR error for <

                    return candidate.strip()

            return ""

        except Exception as e:
            logger.error(f"Given Names extraction failed: {e}")
            return ""

    # ---------------------------------------------------
    # MRZ EXTRACTION
    # ---------------------------------------------------
    def extract_mrz_from_roi(self, img_path):
        try:
            mrz = read_mrz(img_path, save_roi=True)

            if not mrz:
                return None, None, None

            roi = mrz.aux["roi"]

            if roi.dtype != np.uint8:
                roi = (roi * 255).astype(np.uint8)

            img_resized = cv2.resize(roi, (1110, 140))
            allow = st.ascii_uppercase + st.digits + "<"

            code = self.reader.readtext(
                img_resized,
                detail=0,
                allowlist=allow,
                batch_size=1,  # Single image processing for speed
                workers=0,     # Use main thread for stability
                decoder='greedy'  # Faster decoding
            )

            if len(code) < 2:
                return None, None, mrz

            line1 = clean_mrz_line(code[0])
            line2 = clean_mrz_line(code[1])

            return line1, line2, mrz

        except Exception as e:
            logger.error(f"MRZ extraction failed: {e}")
            return None, None, None

    # ---------------------------------------------------
    # MAIN DATA FUNCTION
    # ---------------------------------------------------
    def get_data(self, img_path, airline="iraqi"):

        if not os.path.exists(img_path):
            logger.error(f"File not found: {img_path}")
            return None

        line1, line2, mrz = self.extract_mrz_from_roi(img_path)

        if line1 and line2:
            mrz = FallbackMRZ(line1, line2)

        if mrz is None:
            logger.warning("MRZ not detected.")
            # Return data with placeholder values instead of None
            return {
                "surname": "•••",
                "name": "•••",
                "country": "•••",
                "nationality": "•••",
                "passport_number": "•••",
                "sex": "•••",
                "date_of_birth": "•••",
                "expiration_date": "•••",
                "personal_number": "•••",
                "mrz_full_string": "",
                "valid_score": 0,
                "mrz_found": False
            }

        surname = clean_name_field(getattr(mrz, "surname", ""))

        # 🔥 ALWAYS prefer visual name
        visual_name = self.extract_given_names_from_visual(img_path)

        if visual_name:
            name = visual_name
        else:
            name = clean_name_field(
                getattr(mrz, "names", getattr(mrz, "name", ""))
            )

            # Final defensive cleanup - only remove trailing K which is a common OCR artifact
            name = re.sub(r'([A-Z]{2,})[K]$', r'\1', name)

        data = {
            "surname": surname,
            "name": name,
            "country": get_country_name(getattr(mrz, "country", "")),
            "nationality": get_country_name(getattr(mrz, "nationality", "")),
            "passport_number": clean_string(getattr(mrz, "number", "")),
            "sex": get_sex(getattr(mrz, "sex", "")),
            "date_of_birth": parse_date(getattr(mrz, "date_of_birth", ""), airline=airline),
            "expiration_date": parse_date(getattr(mrz, "expiration_date", ""), airline=airline),
            "mrz_full_string": (line1 or "") + (line2 or ""),
            "valid_score": getattr(mrz, "valid_score", 0),
            "mrz_found": True,
        }

        return data

    # ---------------------------------------------------
    # PDF PROCESSING
    # ---------------------------------------------------
    def process_pdf(self, pdf_path, progress_callback=None, airline="iraqi"):
        """
        Memory-safe PDF processing for Streamlit free tier with fallback support.
        Converts PDF pages to images and extracts passport data from each page.
        
        Args:
            pdf_path (str): Path to the PDF file.
            progress_callback (function, optional): Progress callback function.
            airline (str): Airline format for date formatting ("iraqi", "default", "fly dubai", "fly baghdad").
            
        Returns:
            list: List of dictionaries with extracted passport data per page.
        """
        # Ensure temp directory exists before processing
        os.makedirs(TEMP_DIR, exist_ok=True)
        
        # Check file size for Streamlit free tier (max 10 MB)
        try:
            file_size = os.path.getsize(pdf_path)
            if file_size > 10 * 1024 * 1024:  # 10 MB limit
                logger.error(f"PDF too large for free tier: {file_size / (1024*1024):.1f} MB")
                return []
        except Exception as e:
            logger.error(f"Could not check file size: {e}")
            return []
        
        results = []
        
        try:
            # Try pdf2image first (primary method)
            from pdf2image import pdfinfo_from_path, convert_from_path
            
            info = pdfinfo_from_path(pdf_path)
            total_pages = info["Pages"]
            
            logger.debug(f"Processing PDF with {total_pages} pages (size: {file_size / (1024*1024):.1f} MB) using pdf2image")
            
            for page in range(1, total_pages + 1):
                try:
                    # Update progress if callback provided
                    if progress_callback:
                        progress_callback(page / total_pages)
                    
                    # Convert ONE page at a time, optimize for speed
                    images = convert_from_path(
                        pdf_path,
                        dpi=150,  # Further reduced DPI for faster processing
                        first_page=page,
                        last_page=page,
                        thread_count=1,  # Single thread for stability
                        use_pdftocairo=True  # Faster backend
                    )
                    
                    image = images[0]
                    
                    # Save temporary image to TEMP_DIR with lower quality for speed
                    temp_image_path = os.path.join(TEMP_DIR, f"temp_page_{page}.jpg")
                    image.save(temp_image_path, "JPEG", quality=70, optimize=True)
                    
                    # Extract passport data with airline-specific formatting
                    result = self.get_data(temp_image_path, airline=airline)
                    
                    if result:
                        result["page_number"] = page
                        results.append(result)
                        logger.debug(f"Successfully extracted data from page {page}")
                    
                    # Delete temp image immediately to free memory
                    if os.path.exists(temp_image_path):
                        os.remove(temp_image_path)
                        
                except Exception as e:
                    logger.error(f"Error on page {page} with pdf2image: {e}")
                    continue
            
            logger.debug(f"PDF processing finished with pdf2image. Valid pages: {len(results)}")
            
        except Exception as pdf2image_error:
            logger.warning(f"pdf2image failed: {pdf2image_error}. Trying fallback with PyMuPDF...")
            
            # Fallback to PyMuPDF (fitz)
            try:
                import fitz
                
                doc = fitz.open(pdf_path)
                total_pages = len(doc)
                
                logger.debug(f"Processing PDF with {total_pages} pages using PyMuPDF fallback")
                
                for i in range(total_pages):
                    try:
                        # Update progress if callback provided
                        if progress_callback:
                            progress_callback((i + 1) / total_pages)
                        
                        page = doc.load_page(i)
                        pix = page.get_pixmap(dpi=200)
                        
                        # Convert to PIL Image
                        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                        
                        # Save temporary image to TEMP_DIR
                        temp_image_path = os.path.join(TEMP_DIR, f"temp_page_{i+1}.png")
                        img.save(temp_image_path, "PNG")
                        
                        # Extract passport data with airline-specific formatting
                        result = self.get_data(temp_image_path, airline=airline)
                        
                        if result:
                            result["page_number"] = i + 1
                            results.append(result)
                            logger.debug(f"Successfully extracted data from page {i+1}")
                        
                        # Cleanup
                        if os.path.exists(temp_image_path):
                            os.remove(temp_image_path)
                            
                    except Exception as e:
                        logger.error(f"Error on page {i+1} with PyMuPDF: {e}")
                        continue
                
                doc.close()
                logger.debug(f"PDF processing finished with PyMuPDF fallback. Valid pages: {len(results)}")
                
            except Exception as fitz_error:
                logger.error(f"PyMuPDF fallback also failed: {fitz_error}")
                return []
        
        return results