app.py

# import os
# os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # Suppress TensorFlow INFO and WARNING messages
# os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
import gradio as gr
import base64
import requests
import json
import re
import os
import uuid
from datetime import datetime
import tempfile  # ✅ Add this
import shutil 
import time # For potential sleeps if needed, or timing

# Attempt to import deepface and handle import error gracefully
try:
    import fitz  # PyMuPDF
    PYMUPDF_AVAILABLE = True
except ImportError:
    PYMUPDF_AVAILABLE = False
    print("Warning: PyMuPDF not found. PDF processing will be disabled.")

try:
    import docx
    from PIL import Image, ImageDraw, ImageFont
    DOCX_AVAILABLE = True
except ImportError:
    DOCX_AVAILABLE = False
    print("Warning: python-docx or Pillow not found. DOCX processing will be disabled.")
try:
    from deepface import DeepFace
    # from deepface.commons import functions as deepface_functions
    DEEPFACE_AVAILABLE = True
    print(f"Got DeepFace")
except ImportError:
    DEEPFACE_AVAILABLE = False
    print("Warning: deepface library not found. Facial recognition features will be disabled.")
    # Mock DeepFace object if not available to prevent NameErrors, though functions won't work
    class DeepFaceMock:
        def represent(self, *args, **kwargs): return []
        def verify(self, *args, **kwargs): return {'verified': False, 'distance': float('inf')}
        def detectFace(self, *args, **kwargs): raise NotImplementedError("DeepFace not installed")
    DeepFace = DeepFaceMock()


# --- Configuration ---
OPENROUTER_API_KEY = "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"
IMAGE_MODEL = "opengvlab/internvl3-14b:free"
OPENROUTER_API_URL = "https://openrouter.ai/api/v1/chat/completions"

# Facial Recognition Configuration
FACE_DETECTOR_BACKEND = 'retinaface' # common and effective
FACE_RECOGNITION_MODEL_NAME = 'VGG-Face' # good balance
# Threshold for deepface.verify (model-specific, VGG-Face with cosine is often around 0.40 for verification)
# Lower threshold means stricter match for verify. For similarity search, we might use raw distance.
# DeepFace.verify uses model-specific thresholds internally. Let's rely on its 'verified' flag.
FACE_SIMILARITY_THRESHOLD = 0.60 # For cosine distance, lower is more similar. For similarity, higher is better.
                                 # Deepface verify returns 'distance'. For cosine, lower distance = more similar.
                                 # Let's use a distance threshold. For VGG-Face with cosine, this might be < 0.4 for a match.
                                 # We will use deepface.verify which handles this internally.

# --- Global State ---
processed_files_data = []
person_profiles = {}

# --- Helper Functions ---

def render_text_to_image(text, output_path):
    """Renders a string of text onto a new image file."""
    if not DOCX_AVAILABLE:
        raise ImportError("Pillow or python-docx is not installed.")

    try:
        # Use a built-in font if available, otherwise this might fail on minimal OS
        font = ImageFont.truetype("DejaVuSans.ttf", 15)
    except IOError:
        print("Default font not found, using basic PIL font.")
        font = ImageFont.load_default()

    padding = 20
    image_width = 800
    
    # Simple text wrapping
    lines = []
    for paragraph in text.split('\n'):
        words = paragraph.split()
        line = ""
        for word in words:
            # Use getbbox for more accurate width calculation if available (Pillow >= 9.2.0)
            if hasattr(font, 'getbbox'):
                box = font.getbbox(line + word)
                line_width = box[2] - box[0]
            else: # Fallback for older Pillow
                line_width = font.getsize(line + word)[0]

            if line_width <= image_width - 2 * padding:
                line += word + " "
            else:
                lines.append(line.strip())
                line = word + " "
        lines.append(line.strip())
    
    # Calculate image height
    _, top, _, bottom = font.getbbox("A")
    line_height = bottom - top + 5 # Add some line spacing
    image_height = len(lines) * line_height + 2 * padding
    
    img = Image.new('RGB', (image_width, int(image_height)), color='white')
    draw = ImageDraw.Draw(img)
    
    y = padding
    for line in lines:
        draw.text((padding, y), line, font=font, fill='black')
        y += line_height
        
    img.save(output_path, format='PNG')


def convert_file_to_images(original_filepath, temp_output_dir):
    """
    Converts an uploaded file (PDF, DOCX) into one or more images.
    If the file is already an image, it returns its own path.
    Returns a list of dictionaries, each with 'path' and 'page' keys.
    """
    filename_lower = original_filepath.lower()
    output_paths = []

    if filename_lower.endswith('.pdf'):
        if not PYMUPDF_AVAILABLE:
            raise RuntimeError("PDF processing is disabled (PyMuPDF not installed).")
        doc = fitz.open(original_filepath)
        for i, page in enumerate(doc):
            pix = page.get_pixmap(dpi=200) # Render page to image
            output_filepath = os.path.join(temp_output_dir, f"{os.path.basename(original_filepath)}_page_{i+1}.png")
            pix.save(output_filepath)
            output_paths.append({"path": output_filepath, "page": i + 1})
        doc.close()

    elif filename_lower.endswith('.docx'):
        if not DOCX_AVAILABLE:
            raise RuntimeError("DOCX processing is disabled (python-docx or Pillow not installed).")
        doc = docx.Document(original_filepath)
        full_text = "\n".join([para.text for para in doc.paragraphs])
        if not full_text.strip():
             full_text = "--- Document is empty or contains only images/tables ---"
        output_filepath = os.path.join(temp_output_dir, f"{os.path.basename(original_filepath)}.png")
        render_text_to_image(full_text, output_filepath)
        output_paths.append({"path": output_filepath, "page": 1})

    elif filename_lower.endswith(('.png', '.jpg', '.jpeg', '.webp', '.bmp', '.tiff')):
        # File is already an image, just return its path
        output_paths.append({"path": original_filepath, "page": 1})

    else:
        raise TypeError(f"Unsupported file type: {os.path.basename(original_filepath)}")

    return output_paths


def extract_json_from_text(text):
    if not text:
        return {"error": "Empty text provided for JSON extraction."}
    match_block = re.search(r"```json\s*(\{.*?\})\s*```", text, re.DOTALL | re.IGNORECASE)
    if match_block:
        json_str = match_block.group(1)
    else:
        text_stripped = text.strip()
        if text_stripped.startswith("`") and text_stripped.endswith("`"):
            json_str = text_stripped[1:-1]
        else:
            json_str = text_stripped
    try:
        return json.loads(json_str)
    except json.JSONDecodeError as e:
        try:
            first_brace = json_str.find('{')
            last_brace = json_str.rfind('}')
            if first_brace != -1 and last_brace != -1 and last_brace > first_brace:
                potential_json_str = json_str[first_brace : last_brace+1]
                return json.loads(potential_json_str)
            else:
                return {"error": f"Invalid JSON structure (no outer braces found): {str(e)}", "original_text": text}
        except json.JSONDecodeError as e2:
             return {"error": f"Invalid JSON structure after attempting substring: {str(e2)}", "original_text": text}

def get_ocr_prompt():
    # Enhanced prompt
    return f"""You are an advanced OCR and information extraction AI.
Your task is to meticulously analyze this image and extract all relevant information.

Output Format Instructions:
Provide your response as a SINGLE, VALID JSON OBJECT. Do not include any explanatory text before or after the JSON.
The JSON object should have the following top-level keys:
- "document_type_detected": (string) Your best guess of the specific document type (e.g., "Passport Front", "Passport Back", "National ID Card", "Photo of a person", "Hotel Reservation", "Bank Statement").
- "extracted_fields": (object) A key-value map of all extracted information. Be comprehensive.
    - For ALL document types, if a primary person is the subject, try to include: "Primary Person Name", "Full Name".
    - List other names found under specific keys like "Guest Name", "Account Holder Name", "Mother's Name", "Spouse's Name".
    - Extract critical identifiers like "Passport Number", "Document Number", "ID Number", "Account Number", "Reservation Number" FROM ANY PART OF THE DOCUMENT where they appear. Use consistent key names for these if possible.
    - For passports/IDs: "Surname", "Given Names", "Nationality", "Date of Birth", "Sex", "Place of Birth", "Date of Issue", "Date of Expiry".
    - For photos: "Description" (e.g., "Portrait of John Doe", "User's profile photo"), "People Present" (array of names if discernible).
- "mrz_data": (object or null) If a Machine Readable Zone (MRZ) is present.
- "full_text_ocr": (string) Concatenation of all text found on the document.

Extraction Guidelines:
1. Extract "Passport Number" or "Document Number" even from back sides or less prominent areas.
2. Identify and list all prominent names. If one person is clearly the main subject, label their name as "Primary Person Name" or "Full Name".
3. For dates, aim for YYYY-MM-DD.

Ensure the entire output strictly adheres to the JSON format.
"""

def call_openrouter_ocr(image_filepath):
    # (User's existing function - kept mostly as is, ensure YOUR_SPACE is updated if needed)
    if not OPENROUTER_API_KEY:
        return {"error": "OpenRouter API Key not configured."}
    try:
        with open(image_filepath, "rb") as f:
            encoded_image = base64.b64encode(f.read()).decode("utf-8")
        mime_type = "image/jpeg"
        if image_filepath.lower().endswith(".png"): mime_type = "image/png"
        elif image_filepath.lower().endswith(".webp"): mime_type = "image/webp"
        data_url = f"data:{mime_type};base64,{encoded_image}"
        prompt_text = get_ocr_prompt()
        payload = {
            "model": IMAGE_MODEL,
            "messages": [{"role": "user", "content": [{"type": "text", "text": prompt_text}, {"type": "image_url", "image_url": {"url": data_url}}]}],
            "max_tokens": 3500, "temperature": 0.1,
        }
        headers = {
            "Authorization": f"Bearer {OPENROUTER_API_KEY}", "Content-Type": "application/json",
            "HTTP-Referer": os.environ.get("GRADIO_ROOT_PATH", "http://localhost:7860"), # Better placeholder
            "X-Title": "Gradio Document Processor"
        }
        response = requests.post(OPENROUTER_API_URL, headers=headers, json=payload, timeout=180)
        response.raise_for_status()
        result = response.json()
        if "choices" in result and result["choices"]:
            raw_content = result["choices"][0]["message"]["content"]
            return extract_json_from_text(raw_content)
        else:
            return {"error": "No 'choices' in API response from OpenRouter.", "details": result}
    except requests.exceptions.Timeout: return {"error": "API request timed out."}
    except requests.exceptions.RequestException as e:
        error_message = f"API Request Error: {str(e)}"
        if hasattr(e, 'response') and e.response is not None: error_message += f" Status: {e.response.status_code}, Response: {e.response.text}"
        return {"error": error_message}
    except Exception as e: return {"error": f"An unexpected error occurred during OCR: {str(e)}"}

def get_facial_embeddings_with_deepface(image_filepath):
    if not DEEPFACE_AVAILABLE:
        return {"error": "DeepFace library not installed.", "embeddings": []}
    try:
        # Use represent to get embeddings. It can find multiple faces.
        # Setting align=True, detector_backend for robustness.
        # enforce_detection=False will return empty list if no face, rather than error.
        embedding_objs = DeepFace.represent(
            img_path=image_filepath,
            model_name=FACE_RECOGNITION_MODEL_NAME,
            detector_backend=FACE_DETECTOR_BACKEND,
            enforce_detection=False, # Don't raise error if no face
            align=True
        )
        # DeepFace.represent returns a list of dictionaries, each with an 'embedding' key
        embeddings = [obj['embedding'] for obj in embedding_objs if 'embedding' in obj]
        if not embeddings:
            return {"message": "No face detected or embedding failed.", "embeddings": []}
        return {"embeddings": embeddings, "count": len(embeddings)}
    except Exception as e:
        # Catch errors from DeepFace if enforce_detection was True or other issues
        # Like "Face detector ... could not find anyıs face"
        if "could not find any face" in str(e).lower():
             return {"message": "No face detected.", "embeddings": []}
        return {"error": f"Facial embedding extraction failed: {str(e)}", "embeddings": []}


def extract_entities_from_ocr(ocr_json):
    if not ocr_json or not isinstance(ocr_json, dict) or "extracted_fields" not in ocr_json or not isinstance(ocr_json.get("extracted_fields"), dict):
        doc_type_from_ocr = "Unknown"
        if isinstance(ocr_json, dict):
            doc_type_from_ocr = ocr_json.get("document_type_detected", "Unknown (error in OCR)")
        return {"name": None, "dob": None, "main_id": None, "doc_type": doc_type_from_ocr, "all_names_roles": []}

    fields = ocr_json["extracted_fields"]
    doc_type = ocr_json.get("document_type_detected", "Unknown")

    # Expanded and prioritized name keys
    # Order matters: more specific or primary names first
    name_keys = [
        "primary person name", "full name", "name", "account holder name", "guest name", 
        "cardholder name", "policy holder name", "applicant name", "beneficiary name", 
        "student name", "employee name", "sender name", "receiver name", 
        "patient name", "traveler name", "customer name", "member name", "user name"
    ]
    dob_keys = ["date of birth", "dob"]
    # Expanded ID keys (passport, national ID, etc.)
    id_keys = ["passport number", "document number", "id number", "personal no", "member id", "customer id", "account number", "reservation number"]

    extracted_name = None
    all_names_roles = [] # To store all names found with their original JSON key

    for key in name_keys:
        for field_key, value in fields.items():
            if key == field_key.lower():
                if value and isinstance(value, str) and value.strip():
                    if not extracted_name: # Take the first one found as primary for now
                        extracted_name = value.strip()
                    all_names_roles.append({"name_text": value.strip(), "source_key": field_key})
    # If "People Present" exists (e.g., for photos), add them
    if "people present" in (k.lower() for k in fields.keys()):
        people = fields.get([k for k in fields if k.lower() == "people present"][0])
        if isinstance(people, list):
            for person_name in people:
                if isinstance(person_name, str) and person_name.strip():
                     all_names_roles.append({"name_text": person_name.strip(), "source_key": "People Present"})
                     if not extracted_name: extracted_name = person_name.strip() # Prioritize if no other name found

    extracted_dob = None
    for key in dob_keys:
        for field_key, value in fields.items():
            if key == field_key.lower() and value and isinstance(value, str):
                extracted_dob = value.strip()
                break
        if extracted_dob: break
    
    extracted_main_id = None
    for key in id_keys:
        for field_key, value in fields.items():
            if key == field_key.lower() and value and isinstance(value, str):
                extracted_main_id = value.replace(" ", "").upper().strip() # Normalize
                break
        if extracted_main_id: break
        
    return {
        "name": extracted_name,
        "dob": extracted_dob,
        "main_id": extracted_main_id, # This will be used as the primary linking ID
        "doc_type": doc_type,
        "all_names_roles": list({tuple(d.items()): d for d in all_names_roles}.values()) # Deduplicate
    }

def normalize_name(name):
    if not name: return ""
    return "".join(filter(str.isalnum, name)).lower()

def are_faces_similar(emb1_list, emb2_gallery_list):
    if not DEEPFACE_AVAILABLE or not emb1_list or not emb2_gallery_list:
        return False
    # Compare each embedding from emb1_list against each in emb2_gallery_list
    for emb1 in emb1_list:
        for emb2 in emb2_gallery_list:
            try:
                # DeepFace.verify expects embeddings directly if not paths
                # It uses built-in thresholds per model.
                result = DeepFace.verify(
                    img1_path=emb1,  # Pass embedding directly
                    img2_path=emb2,  # Pass embedding directly
                    model_name=FACE_RECOGNITION_MODEL_NAME,
                    detector_backend=FACE_DETECTOR_BACKEND, # Though not used for verify with embeddings
                    distance_metric='cosine' # Or 'euclidean', 'euclidean_l2'
                )
                if result.get("verified", False):
                    # print(f"Face match found: distance {result.get('distance')}")
                    return True
            except Exception as e:
                print(f"DeepFace verify error: {e}") # e.g. if embeddings are not in expected format
    return False

def get_person_id_and_update_profiles(doc_id, entities, facial_embeddings, current_persons_data, linking_method_log):
    main_id = entities.get("main_id") # Passport No, Document No, Account No etc.
    name = entities.get("name")
    dob = entities.get("dob")

    # Tier 1: Match by Main ID (Passport, National ID, etc.)
    if main_id:
        for p_key, p_data in current_persons_data.items():
            if main_id in p_data.get("ids", set()):
                p_data["doc_ids"].add(doc_id)
                if name and normalize_name(name) not in p_data["names"]: p_data["names"].add(normalize_name(name))
                if dob and dob not in p_data["dobs"]: p_data["dobs"].add(dob)
                if facial_embeddings: p_data["face_gallery"].extend(facial_embeddings) # Add new faces
                linking_method_log.append(f"Linked by Main ID ({main_id}) to {p_key}")
                return p_key
        # New person based on this main_id
        new_person_key = f"person_id_{main_id}"
        current_persons_data[new_person_key] = {
            "display_name": name or f"Person (ID: {main_id})",
            "names": {normalize_name(name)} if name else set(),
            "dobs": {dob} if dob else set(),
            "ids": {main_id},
            "face_gallery": list(facial_embeddings or []), # Initialize gallery
            "doc_ids": {doc_id}
        }
        linking_method_log.append(f"New person by Main ID ({main_id}): {new_person_key}")
        return new_person_key

    # Tier 2: Match by Facial Recognition
    if facial_embeddings:
        for p_key, p_data in current_persons_data.items():
            if are_faces_similar(facial_embeddings, p_data.get("face_gallery", [])):
                p_data["doc_ids"].add(doc_id)
                if name and normalize_name(name) not in p_data["names"]: p_data["names"].add(normalize_name(name))
                if dob and dob not in p_data["dobs"]: p_data["dobs"].add(dob)
                p_data["face_gallery"].extend(facial_embeddings) # Freshen gallery
                linking_method_log.append(f"Linked by Facial Match to {p_key}")
                return p_key
        # If no facial match to existing, but we have a face and name/dob, it will be used for new profile below

    # Tier 3: Match by Normalized Name + DOB
    if name and dob:
        norm_name = normalize_name(name)
        for p_key, p_data in current_persons_data.items():
            if norm_name in p_data.get("names", set()) and dob in p_data.get("dobs", set()):
                p_data["doc_ids"].add(doc_id)
                if facial_embeddings: p_data["face_gallery"].extend(facial_embeddings)
                linking_method_log.append(f"Linked by Name+DOB to {p_key}")
                return p_key
        # New person based on name and DOB
        new_person_key = f"person_{norm_name}_{dob}_{str(uuid.uuid4())[:4]}"
        current_persons_data[new_person_key] = {
            "display_name": name, "names": {norm_name}, "dobs": {dob}, "ids": set(),
            "face_gallery": list(facial_embeddings or []), "doc_ids": {doc_id}
        }
        linking_method_log.append(f"New person by Name+DOB: {new_person_key}")
        return new_person_key
    
    # Tier 4: Match by Normalized Name only (creates a more tentative profile)
    if name:
        norm_name = normalize_name(name)
        # Check if any existing profile primarily matches this name AND has no stronger identifiers yet (e.g. no DOB, no ID, no face)
        # This logic could be refined to prevent overly aggressive merging or splitting.
        # For now, we'll create a new profile if not matched above.
        new_person_key = f"person_name_{norm_name}_{str(uuid.uuid4())[:4]}"
        current_persons_data[new_person_key] = {
            "display_name": name, "names": {norm_name}, "dobs": set(), "ids": set(),
            "face_gallery": list(facial_embeddings or []), "doc_ids": {doc_id}
        }
        linking_method_log.append(f"New person by Name only: {new_person_key}")
        return new_person_key

    # Tier 5: Unclassifiable by PII, but might have a face
    generic_person_key = f"unidentified_person_{str(uuid.uuid4())[:6]}"
    current_persons_data[generic_person_key] = {
        "display_name": f"Unknown Person ({doc_id[:6]})",
        "names": set(), "dobs": set(), "ids": set(),
        "face_gallery": list(facial_embeddings or []), "doc_ids": {doc_id}
    }
    linking_method_log.append(f"New Unidentified Person: {generic_person_key}")
    return generic_person_key


def format_dataframe_data(current_files_data):
    df_rows = []
    for f_data in current_files_data:
        entities = f_data.get("entities") or {}
        face_info = f_data.get("face_analysis_result", {}) or {}
        face_detected_status = "Y" if face_info.get("count", 0) > 0 else "N"
        if "error" in face_info : face_detected_status = "Error"
        elif "message" in face_info and "No face detected" in face_info["message"]: face_detected_status = "N"
        
        df_rows.append([
            f_data.get("doc_id", "N/A")[:8],
            f_data.get("filename", "N/A"),
            f_data.get("status", "N/A"),
            entities.get("doc_type", "N/A"),
            face_detected_status,
            entities.get("name", "N/A"),
            entities.get("dob", "N/A"),
            entities.get("main_id", "N/A"), # Changed from passport_no to main_id
            f_data.get("assigned_person_key", "N/A"),
            f_data.get("linking_method", "N/A")
        ])
    return df_rows

def format_persons_markdown(current_persons_data, current_files_data):
    if not current_persons_data: return "No persons identified yet."
    md_parts = ["## Classified Persons & Documents\n"]
    for p_key, p_data in sorted(current_persons_data.items()): # Sort for consistent display
        display_name = p_data.get('display_name', p_key)
        md_parts.append(f"### Person: {display_name} (Profile Key: {p_key})")
        if p_data.get("dobs"): md_parts.append(f"* Known DOB(s): {', '.join(p_data['dobs'])}")
        if p_data.get("ids"): md_parts.append(f"* Known ID(s): {', '.join(p_data['ids'])}")
        if p_data.get("face_gallery") and len(p_data.get("face_gallery")) > 0:
            md_parts.append(f"* Facial Signatures Stored: {len(p_data.get('face_gallery'))}")
        md_parts.append("* Documents:")
        doc_ids_for_person = sorted(list(p_data.get("doc_ids", set()))) # Sort for consistency
        if doc_ids_for_person:
            for doc_id in doc_ids_for_person:
                doc_detail = next((f for f in current_files_data if f["doc_id"] == doc_id), None)
                if doc_detail:
                    filename = doc_detail.get("filename", "Unknown File")
                    doc_entities = doc_detail.get("entities") or {}
                    doc_type = doc_entities.get("doc_type", "Unknown Type")
                    linking_method = doc_detail.get("linking_method", "")
                    md_parts.append(f"  - {filename} (`{doc_type}`) {linking_method}")
                else: md_parts.append(f"  - Document ID: {doc_id[:8]} (details error)")
        else: md_parts.append("  - No documents currently assigned.")
        md_parts.append("\n---\n")
    return "\n".join(md_parts)

def process_uploaded_files_old(files_list, progress=gr.Progress(track_tqdm=True)):
    global processed_files_data, person_profiles
    processed_files_data = []
    person_profiles = {}
    if not OPENROUTER_API_KEY:
        # Expected number of output components: df_data, persons_md, ocr_json_output, status_textbox
        yield ([["N/A", "ERROR", "API Key Missing", "N/A","N/A", "N/A", "N/A", "N/A","N/A", "N/A"]], "API Key Missing.", "{}", "Error: API Key not set.")
        return
    if not files_list:
        yield ([], "No files uploaded.", "{}", "Upload files to begin.")
        return

    # Initialize file data structures
    for i, file_obj_path in enumerate(files_list): # gr.Files with type="filepath" returns list of path strings
        doc_uid = str(uuid.uuid4())
        processed_files_data.append({
            "doc_id": doc_uid,
            "filename": os.path.basename(file_obj_path),
            "filepath": file_obj_path,
            "status": "Queued", "ocr_json": None, "entities": None,
            "face_analysis_result": None, "facial_embeddings": None,
            "assigned_person_key": None, "linking_method": ""
        })
    
    df_data = format_dataframe_data(processed_files_data)
    persons_md = format_persons_markdown(person_profiles, processed_files_data)
    yield (df_data, persons_md, "{}", f"Initialized {len(files_list)} files.")

    for i, file_data_item in enumerate(progress.tqdm(processed_files_data, desc="Processing Documents")):
        current_doc_id = file_data_item["doc_id"]
        current_filename = file_data_item["filename"]
        linking_method_log_for_doc = [] # To store how this doc was linked

        if not file_data_item["filepath"] or not os.path.exists(file_data_item["filepath"]):
            file_data_item["status"] = "Error: Invalid file"
            linking_method_log_for_doc.append("File path error.")
            file_data_item["linking_method"] = " ".join(linking_method_log_for_doc)
            df_data = format_dataframe_data(processed_files_data)
            persons_md = format_persons_markdown(person_profiles, processed_files_data)
            yield(df_data, persons_md, "{}", f"({i+1}/{len(processed_files_data)}) Error for {current_filename}")
            continue

        # 1. OCR
        file_data_item["status"] = "OCR..."
        df_data = format_dataframe_data(processed_files_data); yield (df_data, persons_md, file_data_item.get("ocr_json_str","{}"), f"OCR: {current_filename}")
        ocr_result = call_openrouter_ocr(file_data_item["filepath"])
        file_data_item["ocr_json"] = ocr_result
        if "error" in ocr_result:
            file_data_item["status"] = f"OCR Err: {str(ocr_result['error'])[:30]}.."
            linking_method_log_for_doc.append("OCR Failed.")
            file_data_item["linking_method"] = " ".join(linking_method_log_for_doc)
            df_data = format_dataframe_data(processed_files_data); yield (df_data, persons_md, json.dumps(ocr_result, indent=2), f"OCR Err: {current_filename}")
            continue
        file_data_item["status"] = "OCR OK. Entities..."
        df_data = format_dataframe_data(processed_files_data); yield (df_data, persons_md, json.dumps(ocr_result, indent=2), f"Entities: {current_filename}")
        
        # 2. Entity Extraction
        entities = extract_entities_from_ocr(ocr_result)
        file_data_item["entities"] = entities
        file_data_item["status"] = "Entities OK. Face..."
        df_data = format_dataframe_data(processed_files_data); yield (df_data, persons_md, json.dumps(ocr_result, indent=2), f"Face Detect: {current_filename}")

        # 3. Facial Feature Extraction
        doc_type_lower = (entities.get("doc_type") or "").lower()
        # Attempt face detection on photos, passports, IDs.
        if DEEPFACE_AVAILABLE and ("photo" in doc_type_lower or "passport" in doc_type_lower or "id card" in doc_type_lower or "selfie" in doc_type_lower):
            face_result = get_facial_embeddings_with_deepface(file_data_item["filepath"])
            file_data_item["face_analysis_result"] = face_result
            if "embeddings" in face_result and face_result["embeddings"]:
                file_data_item["facial_embeddings"] = face_result["embeddings"]
                file_data_item["status"] = f"Face OK ({face_result.get('count',0)}). Classify..."
                linking_method_log_for_doc.append(f"{face_result.get('count',0)} face(s).")
            elif "error" in face_result:
                file_data_item["status"] = f"Face Err: {face_result['error'][:20]}.."
                linking_method_log_for_doc.append("Face Ext. Error.")
            else: # No error, but no embeddings (e.g. no face detected)
                file_data_item["status"] = "No Face. Classify..."
                linking_method_log_for_doc.append("No face det.")
        else:
            file_data_item["status"] = "No Face Ext. Classify..."
            linking_method_log_for_doc.append("Face Ext. Skipped.")
        df_data = format_dataframe_data(processed_files_data); yield (df_data, persons_md, json.dumps(ocr_result, indent=2), f"Classifying: {current_filename}")
        
        # 4. Person Classification
        person_key = get_person_id_and_update_profiles(current_doc_id, entities, file_data_item.get("facial_embeddings"), person_profiles, linking_method_log_for_doc)
        file_data_item["assigned_person_key"] = person_key
        file_data_item["status"] = "Classified"
        file_data_item["linking_method"] = " ".join(linking_method_log_for_doc)

        df_data = format_dataframe_data(processed_files_data)
        persons_md = format_persons_markdown(person_profiles, processed_files_data)
        yield (df_data, persons_md, json.dumps(ocr_result, indent=2), f"Done: {current_filename} -> {person_key}")

    final_df_data = format_dataframe_data(processed_files_data)
    final_persons_md = format_persons_markdown(person_profiles, processed_files_data)
    yield (final_df_data, final_persons_md, "{}", f"All {len(processed_files_data)} documents processed.")

def process_uploaded_files(files_list, progress=gr.Progress(track_tqdm=True)):
    global processed_files_data, person_profiles
    processed_files_data = []
    person_profiles = {}
    temp_dir = tempfile.mkdtemp()  # Create a temporary directory for converted images

    empty_df_row = [["N/A"] * 11]  # Match number of headers
    if not OPENROUTER_API_KEY:
        yield (empty_df_row, "API Key Missing.", "{}", "Error: API Key not set.")
        shutil.rmtree(temp_dir)
        return
    if not files_list:
        yield ([], "No files uploaded.", "{}", "Upload files to begin.")
        shutil.rmtree(temp_dir)
        return

    # --- Stage 1: Pre-process files into a job queue of images ---
    job_queue = []
    for original_file_obj in progress.tqdm(files_list, desc="Pre-processing Files"):
        try:
            image_page_list = convert_file_to_images(original_file_obj.name, temp_dir)
            total_pages = len(image_page_list)
            for item in image_page_list:
                job_queue.append({
                    "original_filename": os.path.basename(original_file_obj.name),
                    "page_number": item["page"],
                    "total_pages": total_pages,
                    "image_path": item["path"]
                })
        except Exception as e:
            job_queue.append({"original_filename": os.path.basename(original_file_obj.name), "error": str(e)})

    for job in job_queue:
        if "error" in job:
            processed_files_data.append({
                "doc_id": str(uuid.uuid4()),
                "original_filename": job["original_filename"],
                "page_number": 1,
                "status": f"Error: {job['error']}"
            })
        else:
            processed_files_data.append({
                "doc_id": str(uuid.uuid4()),
                "original_filename": job["original_filename"],
                "page_number": job["page_number"],
                "total_pages": job["total_pages"],
                "filepath": job["image_path"],
                "status": "Queued",
                "ocr_json": None,
                "entities": None,
                "face_analysis_result": None,
                "facial_embeddings": None,
                "assigned_person_key": None,
                "linking_method": ""
            })

    initial_df_data = format_dataframe_data(processed_files_data)
    initial_persons_md = format_persons_markdown(person_profiles, processed_files_data)
    yield (initial_df_data, initial_persons_md, "{}", f"Pre-processing complete. Analyzing {len(processed_files_data)} pages.")

    # --- Stage 2: Analyze each page ---
    current_ocr_json_display = "{}"
    for i, file_data_item in enumerate(progress.tqdm(processed_files_data, desc="Analyzing Pages")):
        if file_data_item["status"].startswith("Error"):
            continue

        current_filename = f"{file_data_item['original_filename']} (p.{file_data_item['page_number']})"
        linking_method_log_for_doc = []

        # 1. OCR
        file_data_item["status"] = "OCR..."
        persons_md = format_persons_markdown(person_profiles, processed_files_data)
        df_data = format_dataframe_data(processed_files_data)
        yield (df_data, persons_md, current_ocr_json_display, f"OCR: {current_filename}")

        ocr_result = call_openrouter_ocr(file_data_item["filepath"])
        file_data_item["ocr_json"] = ocr_result
        current_ocr_json_display = json.dumps(ocr_result, indent=2)

        if "error" in ocr_result:
            file_data_item["status"] = f"OCR Err: {str(ocr_result['error'])[:30]}.."
            linking_method_log_for_doc.append("OCR Failed.")
            file_data_item["linking_method"] = " ".join(linking_method_log_for_doc)
            persons_md = format_persons_markdown(person_profiles, processed_files_data)
            df_data = format_dataframe_data(processed_files_data)
            yield (df_data, persons_md, current_ocr_json_display, f"OCR Err: {current_filename}")
            continue

        # 2. Entity Extraction
        file_data_item["status"] = "OCR OK. Entities..."
        persons_md = format_persons_markdown(person_profiles, processed_files_data)
        df_data = format_dataframe_data(processed_files_data)
        yield (df_data, persons_md, current_ocr_json_display, f"Entities: {current_filename}")
        entities = extract_entities_from_ocr(ocr_result)
        file_data_item["entities"] = entities

        # 3. Facial Feature Extraction
        file_data_item["status"] = "Entities OK. Face..."
        persons_md = format_persons_markdown(person_profiles, processed_files_data)
        df_data = format_dataframe_data(processed_files_data)
        yield (df_data, persons_md, current_ocr_json_display, f"Face Detect: {current_filename}")
        doc_type_lower = (entities.get("doc_type") or "").lower()

        if DEEPFACE_AVAILABLE and (
            "photo" in doc_type_lower or
            "passport" in doc_type_lower or
            "id" in doc_type_lower or
            "selfie" in doc_type_lower or
            not doc_type_lower
        ):
            face_result = get_facial_embeddings_with_deepface(file_data_item["filepath"])
            file_data_item["face_analysis_result"] = face_result
            if "embeddings" in face_result and face_result["embeddings"]:
                file_data_item["facial_embeddings"] = face_result["embeddings"]
                linking_method_log_for_doc.append(f"{face_result.get('count', 0)} face(s).")
            elif "error" in face_result:
                linking_method_log_for_doc.append("Face Ext. Error.")
            else:
                linking_method_log_for_doc.append("No face det.")
        else:
            linking_method_log_for_doc.append("Face Ext. Skipped.")

        file_data_item["status"] = "Face Done. Classify..."
        persons_md = format_persons_markdown(person_profiles, processed_files_data)
        df_data = format_dataframe_data(processed_files_data)
        yield (df_data, persons_md, current_ocr_json_display, f"Classifying: {current_filename}")

        # 4. Person Classification
        person_key = get_person_id_and_update_profiles(
            file_data_item["doc_id"],
            entities,
            file_data_item.get("facial_embeddings"),
            person_profiles,
            linking_method_log_for_doc
        )
        file_data_item["assigned_person_key"] = person_key
        file_data_item["status"] = "Classified"
        file_data_item["linking_method"] = " ".join(linking_method_log_for_doc)

        persons_md = format_persons_markdown(person_profiles, processed_files_data)
        df_data = format_dataframe_data(processed_files_data)
        yield (df_data, persons_md, current_ocr_json_display, f"Done: {current_filename} -> {person_key}")

    # Final Result
    final_df_data = format_dataframe_data(processed_files_data)
    final_persons_md = format_persons_markdown(person_profiles, processed_files_data)
    yield (final_df_data, final_persons_md, "{}", f"All {len(processed_files_data)} pages analyzed.")

    # Cleanup
    try:
        shutil.rmtree(temp_dir)
        print(f"Cleaned up temporary directory: {temp_dir}")
    except Exception as e:
        print(f"Error cleaning up temporary directory {temp_dir}: {e}")

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📄 Intelligent Document Processor & Classifier v2 (with Face ID)")
    gr.Markdown(
        "**Upload multiple documents. The system will OCR, extract entities & faces, and classify documents by person.**\n"
        "Ensure `OPENROUTER_API_KEY` is set as a Secret. Facial recognition uses `deepface` ('VGG-Face' model, 'retinaface' detector)."
    )
    if not OPENROUTER_API_KEY: gr.Markdown("<h3 style='color:red;'>⚠️ ERROR: `OPENROUTER_API_KEY` Secret missing! OCR will fail.</h3>")
    if not DEEPFACE_AVAILABLE: gr.Markdown("<h3 style='color:orange;'>⚠️ WARNING: `deepface` library not installed. Facial recognition features are disabled.</h3>")

    with gr.Row():
        with gr.Column(scale=1):
            files_input = gr.Files(label="Upload Document Images (Bulk)", file_count="multiple", type="filepath")
            process_button = gr.Button("Process Uploaded Documents", variant="primary")
        with gr.Column(scale=2):
            overall_status_textbox = gr.Textbox(label="Current Task & Overall Progress", interactive=False, lines=2)
    
    gr.Markdown("---")
    gr.Markdown("## Document Processing Details")
    dataframe_headers = ["Doc ID", "Filename", "Status", "Type", "Face?", "Name", "DOB", "Main ID", "Person Key", "Linking Method"]
    document_status_df = gr.Dataframe(
        headers=dataframe_headers, datatype=["str"] * len(dataframe_headers),
        label="Individual Document Status & Extracted Entities",
        row_count=(1, "dynamic"), col_count=(len(dataframe_headers), "fixed"), wrap=True
    )
    
    with gr.Accordion("Selected Document Full OCR JSON", open=False):
        ocr_json_output = gr.Code(label="OCR JSON", language="json", interactive=False)

    gr.Markdown("---")
    person_classification_output_md = gr.Markdown("## Classified Persons & Documents\nNo persons identified yet.")

    process_button.click(
        fn=process_uploaded_files, inputs=[files_input],
        outputs=[document_status_df, person_classification_output_md, ocr_json_output, overall_status_textbox]
    )

    @document_status_df.select(inputs=None, outputs=ocr_json_output, show_progress="hidden")
    def display_selected_ocr(evt: gr.SelectData):
        if evt.index is None or evt.index[0] is None: return "{}"
        selected_row_index = evt.index[0]
        # Access global state. Be cautious with globals in complex apps.
        if 0 <= selected_row_index < len(processed_files_data):
            selected_doc_data = processed_files_data[selected_row_index]
            if selected_doc_data and selected_doc_data.get("ocr_json"):
                ocr_data_to_display = selected_doc_data["ocr_json"]
                return json.dumps(ocr_data_to_display, indent=2, ensure_ascii=False)
        return json.dumps({"message": "No OCR data or selection out of bounds."}, indent=2)

if __name__ == "__main__":
    demo.queue().launch(debug=True, share=os.environ.get("GRADIO_SHARE", "true").lower() == "true")