finalsite.py

import os
import cv2
import numpy as np
import gradio as gr
from ultralytics import YOLO
from pdf2image import convert_from_path
from PIL import Image
import easyocr
import uuid
import re
import difflib
import math
from faker import Faker
import datetime
import random
import csv
import tempfile
from io import StringIO

# Initialize Faker
fake = Faker()

# Conditional import for LLM
try:
    from llama_cpp import Llama
    LLAMA_AVAILABLE = True
except ImportError:
    print("Warning: llama_cpp not available. LLM functionality will be disabled.")
    LLAMA_AVAILABLE = False

# --- Configuration ---
# Folders for temporary files and results
UPLOAD_FOLDER = 'static/uploads/'
RESULTS_FOLDER = 'static/results/'

# --- Model Paths (Update these paths if necessary) ---
CUSTOM_MODEL_PATH = 'best.pt'
PRETRAINED_MODEL_PATH = 'yolov10s.pt'
SIGNATURE_MODEL_PATH = 'yolov8s.pt'
LLAMA_MODEL_PATH = "unsloth.F16.gguf"

# Detection Parameters
YOLO_CONFIDENCE_THRESHOLD = 0.5
OCR_CONFIDENCE_THRESHOLD = 0.5

# Create directories if they don't exist
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(RESULTS_FOLDER, exist_ok=True)

# --- Global Model Placeholders ---
custom_model, pretrained_model, signature_model, reader, llama_model = None, None, None, None, None

def load_models():
    """
    Loads all AI models into the global scope. This function is called on the first
    analysis request to avoid startup conflicts. It ensures models are only loaded once.
    """
    global custom_model, pretrained_model, signature_model, reader, llama_model
    # If models are already loaded, do nothing.
    if reader is not None and (llama_model is not None or not LLAMA_AVAILABLE):
        print("Models already loaded.")
        return
    print("=== Loading Models (this may take a moment) ===")
    
    # Helper function to check for model files
    def check_model_path(path, name):
        if not os.path.exists(path):
            print(f"✗ WARNING: {name} model not found at '{path}'. The application may not function correctly.")
            return False
        return True
    
    # YOLO Models
    if check_model_path(CUSTOM_MODEL_PATH, "Custom YOLO"):
        try:
            custom_model = YOLO(CUSTOM_MODEL_PATH)
            print("✓ Custom YOLO model loaded.")
        except Exception as e:
            print(f"✗ Error loading custom model: {e}")
    if check_model_path(PRETRAINED_MODEL_PATH, "Pre-trained YOLO"):
        try:
            pretrained_model = YOLO(PRETRAINED_MODEL_PATH)
            print("✓ Pre-trained YOLO model loaded.")
        except Exception as e:
            print(f"✗ Error loading pre-trained model: {e}")
    if check_model_path(SIGNATURE_MODEL_PATH, "Signature YOLO"):
        try:
            signature_model = YOLO(SIGNATURE_MODEL_PATH)
            print("✓ Signature YOLO model loaded.")
        except Exception as e:
            print(f"✗ Error loading signature model: {e}")
    
    # OCR Model
    try:
        reader = easyocr.Reader(['en'], gpu=True)
        print("✓ EasyOCR model loaded.")
    except Exception as e:
        print(f"✗ Error loading EasyOCR: {e}. Text detection will be unavailable.")
    
    # LLM Model - Only load if available
    if LLAMA_AVAILABLE and check_model_path(LLAMA_MODEL_PATH, "LLM"):
        try:
            llama_model = Llama(
                model_path=LLAMA_MODEL_PATH,
                n_gpu_layers=-1, n_ctx=4096, chat_format="llama-3", verbose=False
            )
            print("✓ LLM model loaded.")
        except Exception as e:
            print(f"✗ Error loading LLM model: {e}. Text analysis will be unavailable.")
    print("=== All Models Initialized ===")

# YOLO Class Mappings
CUSTOM_CLASS_NAMES = {0: 'face', 1: 'qr', 2: 'signature'}
PRETRAINED_CLASS_MAP = {0: 'face'}

# --- Core Detection & Processing Functions ---
def detect_visual_pii(image_data):
    """Runs the three-stage YOLO detection on a single image."""
    all_boxes = []
    all_classes = []
    if custom_model is None:
        print("Custom model not available for visual detection")
        return all_boxes, all_classes
    
    # Pass 1: Custom Model
    custom_results = custom_model.predict(source=image_data, conf=YOLO_CONFIDENCE_THRESHOLD, verbose=False)[0]
    detected_custom_classes = {CUSTOM_CLASS_NAMES[int(cls)] for cls in custom_results.boxes.cls}
    for box, cls in zip(custom_results.boxes.xyxy.cpu().numpy().astype(int), custom_results.boxes.cls):
        all_boxes.append(box)
        all_classes.append(CUSTOM_CLASS_NAMES[int(cls)])
    
    # Pass 2: Pre-trained Model (Face Fallback)
    if 'face' not in detected_custom_classes and pretrained_model is not None:
        print(" Custom model missed 'face'. Trying pre-trained model as fallback.")
        pretrained_results = pretrained_model.predict(source=image_data, conf=YOLO_CONFIDENCE_THRESHOLD, verbose=False)[0]
        for box in pretrained_results.boxes:
            if int(box.cls[0]) in PRETRAINED_CLASS_MAP:
                all_boxes.append(box.xyxy.cpu().numpy().astype(int)[0])
                all_classes.append("face (fallback)")
    
    # Pass 3: Specialized Model (Signature Fallback)
    if 'signature' not in detected_custom_classes and signature_model is not None:
        print(" Custom model missed 'signature'. Trying specialized signature model as fallback.")
        signature_results = signature_model.predict(source=image_data, conf=YOLO_CONFIDENCE_THRESHOLD, verbose=False)[0]
        for box in signature_results.boxes:
            all_boxes.append(box.xyxy.cpu().numpy().astype(int)[0])
            all_classes.append("signature (fallback)")
    
    return all_boxes, all_classes

# --- OCR + LLM Functions ---
def calculate_distance(bbox1, bbox2):
    """Calculates the Euclidean distance between the centers of two bounding boxes."""
    c1_x = (bbox1[0] + bbox1[2]) / 2
    c1_y = (bbox1[1] + bbox1[3]) / 2
    c2_x = (bbox2[0] + bbox2[2]) / 2
    c2_y = (bbox2[1] + bbox2[3]) / 2
    return math.sqrt((c2_x - c1_x)**2 + (c2_y - c1_y)**2)

def refine_pii_flags(ocr_results, isolation_threshold=150):
    """Post-processing step to unmark short, isolated PII detections."""
    pii_indices = [i for i, result in enumerate(ocr_results) if result["is_pii"]]
    if len(pii_indices) <= 1:
        return ocr_results
    indices_to_unmark = []
    for i in pii_indices:
        current_result = ocr_results[i]
        normalized_text = re.sub(r'[^a-zA-Z0-9]', '', current_result["text"])
    
        if len(normalized_text) <= 3:
            min_dist_to_neighbor = float('inf')
        
            for j in pii_indices:
                if i == j:
                    continue
            
                other_result = ocr_results[j]
                dist = calculate_distance(current_result["bbox"], other_result["bbox"])
                if dist < min_dist_to_neighbor:
                    min_dist_to_neighbor = dist
            
            if min_dist_to_neighbor > isolation_threshold:
                print(f" - Refining PII: Unmarking short ('{current_result['text']}') and isolated (min_dist: {min_dist_to_neighbor:.2f}px) PII.")
                indices_to_unmark.append(i)

    for i in indices_to_unmark:
        ocr_results[i]["is_pii"] = False
        
    return ocr_results

def parse_pii_output(generated_text):
    """Parse the new curly braces format PII output"""
    pii_list = []
    try:
        match = re.search(r'\{([^}]*)\}', generated_text)
        if match:
            content = match.group(1)
            items = re.findall(r'"([^"]*)"', content)
            pii_list = [item.strip() for item in items if item.strip()]
    except Exception as e:
        print(f"Error parsing PII output: {e}")
        pii_list = []
    return pii_list

def normalize_text(text):
    """Comprehensive text normalization for better matching"""
    if not text:
        return ""
    normalized = re.sub(r'[.,;:!?()"\'\-_/\\]', '', text)
    ocr_corrections = {
        '0': 'o', 'O': '0', '1': 'l', 'l': '1', '5': 's', 'S': '5',
        '8': 'b', 'B': '8', 'rn': 'm', 'RN': 'M', 'vv': 'w', 'VV': 'W',
        'cl': 'd', 'CL': 'D',
    }
    for wrong, correct in ocr_corrections.items():
        normalized = normalized.replace(wrong, correct)
    normalized = ' '.join(normalized.split()).lower()
    return normalized

def fuzzy_match_score(text1, text2, threshold=0.8):
    """Calculate fuzzy matching score between two strings"""
    if not text1 or not text2:
        return False
    return difflib.SequenceMatcher(None, text1.lower(), text2.lower()).ratio() >= threshold

def levenshtein_distance(s1, s2):
    """Calculate Levenshtein distance between two strings"""
    if len(s1) < len(s2):
        return levenshtein_distance(s2, s1)
    if len(s2) == 0:
        return len(s1)
    previous_row = list(range(len(s2) + 1))
    for i, c1 in enumerate(s1):
        current_row = [i + 1]
        for j, c2 in enumerate(s2):
            insertions = previous_row[j + 1] + 1
            deletions = current_row[j] + 1
            substitutions = previous_row[j] + (c1 != c2)
            current_row.append(min(insertions, deletions, substitutions))
        previous_row = current_row
    return previous_row[-1]

def is_similar_by_edit_distance(text1, text2, max_distance=2):
    """Check if two texts are similar within edit distance threshold"""
    if not text1 or not text2:
        return False
    distance = levenshtein_distance(text1.lower(), text2.lower())
    max_len = max(len(text1), len(text2))
    if max_len <= 3:
        threshold = 1
    elif max_len <= 6:
        threshold = 2
    else:
        threshold = min(max_distance, max_len // 3)
    return distance <= threshold

def extract_sentence_text(ocr_results):
    """Extract sentence-based text for LLM input"""
    paragraph_text = ""
    for item in ocr_results:
        if len(item) == 3:
            _, text, _ = item
        elif len(item) == 2:
            _, text = item
        else:
            print(f"Unexpected OCR result format: {item}")
            continue
        if text.strip():
            paragraph_text += text + " "
    return paragraph_text.strip()

def extract_word_bboxes_improved(ocr_results):
    """Improved word extraction with better handling of punctuation and spacing"""
    word_bbox_map = []
    for item in ocr_results:
        if len(item) == 3:
            bbox, text, confidence = item
        elif len(item) == 2:
            bbox, text = item
            confidence = 1.0
        else:
            print(f"Unexpected OCR result format: {item}")
            continue
    
        original_text = text.strip()
        if not original_text:
            continue
    
        if isinstance(bbox[0], (list, tuple)):
            x_coords = [point[0] for point in bbox]
            y_coords = [point[1] for point in bbox]
            line_x1, line_y1 = min(x_coords), min(y_coords)
            line_x2, line_y2 = max(x_coords), max(y_coords)
        else:
            line_x1, line_y1, line_x2, line_y2 = bbox
    
        tokens = re.findall(r'\S+', original_text)
        if len(tokens) <= 1:
            padding = 1
            word_bbox_map.append({
                "word": original_text,
                "bbox": [
                    max(0, int(line_x1 - padding)),
                    max(0, int(line_y1 - padding)),
                    int(line_x2 + padding),
                    int(line_y2 + padding)
                ],
                "confidence": confidence,
                "original_line": original_text
            })
            continue
    
        full_width = line_x2 - line_x1
        text_without_spaces = original_text.replace(' ', '')
        total_chars = len(text_without_spaces)
    
        char_position = 0
        for i, token in enumerate(tokens):
            token_start_ratio = char_position / total_chars if total_chars > 0 else 0
            char_position += len(token)
            token_end_ratio = char_position / total_chars if total_chars > 0 else 1
        
            token_x1 = line_x1 + (full_width * token_start_ratio)
            token_x2 = line_x1 + (full_width * token_end_ratio)
        
            padding = 1
            word_bbox = [
                max(0, int(token_x1 - padding)),
                max(0, int(line_y1 - padding)),
                int(min(token_x2 + padding, line_x2)),
                int(line_y2 + padding)
            ]
        
            word_bbox_map.append({
                "word": token,
                "bbox": word_bbox,
                "confidence": confidence,
                "original_line": original_text
            })
    return sorted(word_bbox_map, key=lambda x: (x['bbox'][1], x['bbox'][0]))

def advanced_match_pii_to_words(pii_list, word_bbox_map):
    """Advanced multi-strategy PII matching with comprehensive fallbacks"""
    ocr_results_for_template = []
    words = [info['word'] for info in word_bbox_map]
    bboxes = [info['bbox'] for info in word_bbox_map]
    is_pii_flags = [False] * len(words)
    # Pre-process all words with different normalization strategies
    normalized_words = [normalize_text(word) for word in words]
    print(f"Processing {len(pii_list)} PII items against {len(words)} OCR words")
    for pii_idx, pii_item in enumerate(pii_list):
        if not pii_item.strip():
            continue
    
        print(f"Processing PII item {pii_idx + 1}: '{pii_item}'")
    
        # Normalize the PII item
        normalized_pii = normalize_text(pii_item)
        pii_words = normalized_pii.split()
    
        if not pii_words:
            continue
    
        matched = False
    
        # Strategy 1: Exact matching after normalization
        if len(pii_words) == 1:
            pii_word = pii_words[0]
            for idx, norm_word in enumerate(normalized_words):
                if norm_word == pii_word and not is_pii_flags[idx]:
                    is_pii_flags[idx] = True
                    matched = True
                    print(f" ✓ Exact match: '{words[idx]}' -> '{pii_item}'")
        else:
            # Multi-word exact matching
            pii_len = len(pii_words)
            start_idx = 0
            while start_idx < len(normalized_words) - pii_len + 1:
                exact_match = True
                for j in range(pii_len):
                    if normalized_words[start_idx + j] != pii_words[j]:
                        exact_match = False
                        break
                
                if exact_match:
                    # Check spatial proximity
                    spatial_ok = True
                    for j in range(1, pii_len):
                        prev_bbox = bboxes[start_idx + j - 1]
                        curr_bbox = bboxes[start_idx + j]
                    
                        horizontal_distance = curr_bbox[0] - prev_bbox[2]
                        vertical_alignment = (abs(prev_bbox[1] - curr_bbox[1]) < 30 and
                                              abs(prev_bbox[3] - curr_bbox[3]) < 30)
                    
                        if not (vertical_alignment and horizontal_distance <= 150):
                            spatial_ok = False
                            break
                    
                    if spatial_ok:
                        for j in range(pii_len):
                            if not is_pii_flags[start_idx + j]:
                                is_pii_flags[start_idx + j] = True
                        matched = True
                        matched_text = ' '.join(words[start_idx:start_idx + pii_len])
                        print(f" ✓ Multi-word exact: '{matched_text}' -> '{pii_item}'")
                        start_idx += pii_len
                        continue
                start_idx += 1
    
        # Strategy 2: Fuzzy matching if exact matching failed
        if not matched:
            if len(pii_words) == 1:
                pii_word = pii_words[0]
                for idx, norm_word in enumerate(normalized_words):
                    if (not is_pii_flags[idx] and
                        (fuzzy_match_score(norm_word, pii_word, 0.9) or
                         is_similar_by_edit_distance(norm_word, pii_word, 2))):
                        is_pii_flags[idx] = True
                        matched = True
                        print(f" ✓ Fuzzy match: '{words[idx]}' -> '{pii_item}'")
            else:
                # Multi-word fuzzy matching
                pii_len = len(pii_words)
                start_idx = 0
                while start_idx < len(normalized_words) - pii_len + 1:
                    fuzzy_match = True
                    for j in range(pii_len):
                        if not (fuzzy_match_score(normalized_words[start_idx + j], pii_words[j], 0.85) or
                                is_similar_by_edit_distance(normalized_words[start_idx + j], pii_words[j], 2)):
                            fuzzy_match = False
                            break
                    
                    if fuzzy_match:
                        # Check spatial proximity
                        spatial_ok = True
                        for j in range(1, pii_len):
                            prev_bbox = bboxes[start_idx + j - 1]
                            curr_bbox = bboxes[start_idx + j]
                        
                            horizontal_distance = curr_bbox[0] - prev_bbox[2]
                            vertical_alignment = (abs(prev_bbox[1] - curr_bbox[1]) < 30 and
                                                  abs(prev_bbox[3] - curr_bbox[3]) < 30)
                        
                            if not (vertical_alignment and horizontal_distance <= 150):
                                spatial_ok = False
                                break
                        
                        if spatial_ok:
                            for j in range(pii_len):
                                if not is_pii_flags[start_idx + j]:
                                    is_pii_flags[start_idx + j] = True
                            matched = True
                            matched_text = ' '.join(words[start_idx:start_idx + pii_len])
                            print(f" ✓ Multi-word fuzzy: '{matched_text}' -> '{pii_item}'")
                            start_idx += pii_len
                            continue
                    start_idx += 1
    
        # Strategy 3: Substring and partial matching
        if not matched:
            full_normalized_text = ' '.join(normalized_words)
        
            pos = 0
            while True:
                start_pos = full_normalized_text.find(normalized_pii, pos)
                if start_pos == -1:
                    break
                end_pos = start_pos + len(normalized_pii)
                
                char_count = 0
                start_word_idx = None
                end_word_idx = None
                
                for idx, norm_word in enumerate(normalized_words):
                    word_start = char_count
                    word_end = char_count + len(norm_word)
                    
                    if start_word_idx is None and word_end > start_pos:
                        start_word_idx = idx
                    
                    if word_start < end_pos:
                        end_word_idx = idx
                    
                    char_count += len(norm_word) + 1
                
                if start_word_idx is not None and end_word_idx is not None and end_word_idx - start_word_idx + 1 >= len(pii_words):
                    spatial_ok = True
                    for j in range(start_word_idx, end_word_idx):
                        if j + 1 <= end_word_idx:
                            prev_bbox = bboxes[j]
                            next_bbox = bboxes[j + 1]
                        
                            horizontal_distance = next_bbox[0] - prev_bbox[2]
                            vertical_alignment = (abs(prev_bbox[1] - next_bbox[1]) < 30 and
                                                  abs(prev_bbox[3] - next_bbox[3]) < 30)
                        
                            if not (vertical_alignment and horizontal_distance <= 200):
                                spatial_ok = False
                                break
                    
                    if spatial_ok:
                        for j in range(start_word_idx, end_word_idx + 1):
                            if not is_pii_flags[j]:
                                is_pii_flags[j] = True
                        matched = True
                        matched_text = ' '.join(words[start_word_idx:end_word_idx + 1])
                        print(f" ✓ Substring match: '{matched_text}' -> '{pii_item}'")
                pos = end_pos
    
        # Strategy 4: Individual word matching with relaxed criteria
        if not matched:
            for pii_word in pii_words:
                if len(pii_word) < 3:
                    continue
                
                for idx, norm_word in enumerate(normalized_words):
                    if not is_pii_flags[idx]:
                        if (norm_word == pii_word or
                            fuzzy_match_score(norm_word, pii_word, 0.8) or
                            is_similar_by_edit_distance(norm_word, pii_word, 2) or
                            (len(pii_word) > 5 and (pii_word in norm_word or norm_word in pii_word))):
                            is_pii_flags[idx] = True
                            print(f" ✓ Individual word match: '{words[idx]}' -> '{pii_word}' from '{pii_item}'")
    
        if not matched:
            print(f" ✗ No match found for: '{pii_item}'")
    for idx, word_info in enumerate(word_bbox_map):
        ocr_results_for_template.append({
            "text": word_info["word"],
            "bbox": word_info["bbox"],
            "is_pii": is_pii_flags[idx],
            "confidence": word_info.get("confidence", 1.0)
        })
    ocr_results_for_template = merge_horizontal_pii_boxes_improved(ocr_results_for_template)
    return ocr_results_for_template

def merge_horizontal_pii_boxes_improved(ocr_results, merge_distance=50):
    """Improved merging with better spatial awareness and tighter boxes"""
    if not ocr_results:
        return ocr_results
    merged_results = []
    i = 0
    while i < len(ocr_results):
        current_word = ocr_results[i]
    
        if not current_word["is_pii"]:
            merged_results.append(current_word)
            i += 1
            continue
    
        merge_group = [current_word]
        j = i + 1
    
        while j < len(ocr_results):
            next_word = ocr_results[j]
        
            if not next_word["is_pii"]:
                break
        
            current_bbox = merge_group[-1]["bbox"]
            next_bbox = next_word["bbox"]
        
            y_center_current = (current_bbox[1] + current_bbox[3]) / 2
            y_center_next = (next_bbox[1] + next_bbox[3]) / 2
            y_overlap = abs(y_center_current - y_center_next) < 20
        
            horizontal_distance = next_bbox[0] - current_bbox[2]
        
            if y_overlap and horizontal_distance <= merge_distance:
                merge_group.append(next_word)
                j += 1
            else:
                break
    
        if len(merge_group) > 1:
            min_x = min(word["bbox"][0] for word in merge_group)
            min_y = min(word["bbox"][1] for word in merge_group)
            max_x = max(word["bbox"][2] for word in merge_group)
            max_y = max(word["bbox"][3] for word in merge_group)
        
            merged_text = " ".join(word["text"] for word in merge_group)
        
            merged_word = {
                "text": merged_text,
                "bbox": [min_x, min_y, max_x, max_y],
                "is_pii": True,
                "confidence": max(word.get("confidence", 1.0) for word in merge_group)
            }
            merged_results.append(merged_word)
            print(f" ✓ Merged PII box: '{merged_text}' at [{min_x},{min_y},{max_x},{max_y}]")
        else:
            merged_results.append(current_word)
    
        i = j
    return merged_results

def post_process_pii_detection(ocr_results_for_template, pii_list):
    """Post-process to catch any missed PII using relaxed matching"""
    words = [result["text"] for result in ocr_results_for_template]
    for pii_item in pii_list:
        normalized_pii = normalize_text(pii_item)
        pii_words = normalized_pii.split()
    
        if not pii_words:
            continue
    
        pii_detected = False
        for result in ocr_results_for_template:
            if result["is_pii"]:
                result_normalized = normalize_text(result["text"])
                if (normalized_pii in result_normalized or
                    result_normalized in normalized_pii or
                    fuzzy_match_score(result_normalized, normalized_pii, 0.7)):
                    pii_detected = True
                    break
    
        if not pii_detected:
            print(f" ⚠ PII not detected, trying fallback matching: '{pii_item}'")
        
            for idx, result in enumerate(ocr_results_for_template):
                if result["is_pii"]:
                    continue
                
                word_normalized = normalize_text(result["text"])
                
                for pii_word in pii_words:
                    if (len(pii_word) > 3 and
                        (pii_word in word_normalized or
                         word_normalized in pii_word or
                         fuzzy_match_score(word_normalized, pii_word, 0.6) or
                         is_similar_by_edit_distance(word_normalized, pii_word, 3))):
                    
                        ocr_results_for_template[idx]["is_pii"] = True
                        print(f" ✓ Fallback match: '{result['text']}' -> '{pii_word}' from '{pii_item}'")
                        break
    return ocr_results_for_template

def detect_pii_from_combined_text(combined_text):
    """Detect PII from combined multi-page text using LLM"""
    if llama_model is None:
        print("LLM model not available for PII detection")
        return [], "LLM model not available"
    instruction = (
        "Extract all Personally Identifiable Information (PII) of the main subject from the given text. "
        "Include data like Name, Date of Birth, Gender, Address, Phone Number, Email, Social Security Number (SSN), Member ID, Group Number, or any other PII data available. "
        "Ignore any information about doctors, staff, providers, colleagues, organizations, companies, hospitals, educational institutes, or facilities. "
        "Return the results strictly as a flat set of strings enclosed in { } without labels."
    )
    prompt_content = f"{instruction}\n{combined_text}"
    pii_list = []
    llama_raw_output = ""
    try:
        messages = [{"role": "user", "content": prompt_content}]
        response = llama_model.create_chat_completion(
            messages=messages,
            max_tokens=512,
            temperature=0.1,
        )
        llama_raw_output = response['choices'][0]['message']['content']
        pii_list = parse_pii_output(llama_raw_output)
        print(f"LLM detected {len(pii_list)} PII items from combined text: {pii_list}")
    except Exception as e:
        print(f"Error during Llama PII detection: {e}")
        llama_raw_output = f"Error: {str(e)}"
        pii_list = []
    return pii_list, llama_raw_output

# --- Combined Processing Function ---
def process_page_combined(img_cv, global_pii_list):
    """Process a single page with both YOLO and OCR+LLM detection"""
    all_detections = []
    # Step 1: YOLO Visual Detection
    print(" Running YOLO visual detection...")
    visual_boxes, visual_classes = detect_visual_pii(img_cv)
    for box, cls in zip(visual_boxes, visual_classes):
        all_detections.append({
            "text": cls,
            "bbox": box.tolist() if hasattr(box, 'tolist') else box,
            "is_pii": True,
            "confidence": 1.0,
            "detection_type": "visual"
        })
    print(f" YOLO detected {len(visual_boxes)} visual elements")
    # Step 2: OCR + LLM Text Detection
    if reader is not None:
        print(" Running OCR text extraction...")
        word_ocr_results = reader.readtext(img_cv, paragraph=False, width_ths=0.7, height_ths=0.7)
    
        if word_ocr_results:
            word_bbox_map = extract_word_bboxes_improved(word_ocr_results)
        
            ocr_results_for_template = advanced_match_pii_to_words(global_pii_list, word_bbox_map)
            ocr_results_for_template = post_process_pii_detection(ocr_results_for_template, global_pii_list)
            ocr_results_for_template = refine_pii_flags(ocr_results_for_template)
            ocr_results_for_template = merge_horizontal_pii_boxes_improved(ocr_results_for_template)
        
            for result in ocr_results_for_template:
                if result["is_pii"]:
                    result["detection_type"] = "text"
                    all_detections.append(result)
            
            print(f" OCR detected {sum(1 for r in ocr_results_for_template if r['is_pii'])} text PII elements")
    return all_detections

def classify_pii(text):
    text = text.strip()
    clean_text = re.sub(r'\s+', '', text)
    if re.match(r'^\d{3}-\d{2}-\d{4}$', text) or re.match(r'^\d{3}-\d{2}-\d{4}$', clean_text):
        return 'ssn'
    elif re.match(r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b', text) or re.match(r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b', clean_text) or '@' in text:
        return 'email'
    elif re.match(r'^\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}$', text) or re.match(r'^\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}$', clean_text):
        return 'phone'
    elif re.match(r'^\d{1,2}/\d{1,2}/\d{4}$', text) or re.match(r'^\d{4}-\d{2}-\d{2}$', text) or re.match(r'^\d{1,2}/\d{1,2}/\d{4}$', clean_text) or re.match(r'^\d{1,2}-\d{1,2}-\d{4}$', text) or re.match(r'^\d{1,2}-\d{1,2}-\d{2}$', text) or re.match(r'^\d{2}/\d{2}/\d{4}$', text) or re.match(r'^\d{2}-\d{2}-\d{4}$', text) or re.match(r'^(January|February|March|April|May|June|July|August|September|October|November|December)\s+\d{1,2},\s+\d{4}$', text, re.IGNORECASE) or re.match(r'^(Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\s+\d{1,2},\s+\d{4}$', text, re.IGNORECASE) or re.match(r'^\d{1,2}\s+(January|February|March|April|May|June|July|August|September|October|November|December)\s+\d{4}$', text, re.IGNORECASE) or re.match(r'^\d{1,2}\s+(Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\s+\d{4}$', text, re.IGNORECASE):
        return 'dob'
    elif re.match(r'^(?=.*\d)[A-Za-z0-9]+$', text) and len(text) > 5:
        return 'id'
    elif ',' in text or 'St' in text or 'Ave' in text or re.search(r'\d{5}', text):
        return 'address'
    elif re.match(r'^(male|female|m|f|transgender|nonbinary|non-binary|other|unknown|u|o)$', text.lower()):
        return 'gender'
    else:
        return 'name'

def format_gender(base_gender, original_text):
    orig = original_text.strip()
    orig_lower = orig.lower()
    if orig_lower not in ['male', 'female', 'm', 'f', 'transgender', 'nonbinary', 'non-binary', 'other', 'unknown', 'u', 'o']:
        return base_gender.capitalize()
    if len(orig) == 1:
        char = 'M' if base_gender == 'male' else 'F'
        return char.lower() if orig.islower() else char
    else:
        if orig.isupper():
            return base_gender.upper()
        elif orig.islower():
            return base_gender.lower()
        else:
            return base_gender.capitalize()

def generate_fake(pii_type, length, original=None):
    max_attempts = 100
    if pii_type == 'dob' and original:
        formats = [
            '%m/%d/%Y', '%m/%d/%y', '%d/%m/%Y', '%d/%m/%y', '%Y-%m-%d', '%y-%m-%d', '%m-%d-%Y', '%m-%d-%y',
            '%d-%m-%Y', '%d-%m-%y', '%Y/%m/%d', '%y/%m/%d', '%d.%m.%Y', '%m.%d.%Y', '%B %d, %Y', '%b %d, %Y',
            '%d %B %Y', '%d %b %Y', '%B %d %Y', '%b %d %Y'
        ]
        for fmt in formats:
            try:
                datetime.datetime.strptime(original.strip(), fmt)
                fake_dt = fake.date_object()
                return fake_dt.strftime(fmt)
            except ValueError:
                pass
        return fake.date(pattern='%m/%d/%Y')
    elif pii_type == 'gender' and original:
        return format_gender(random.choice(['male', 'female']), original)
    elif pii_type in ['name', 'email', 'phone', 'address']:
        for _ in range(max_attempts):
            if pii_type == 'name': f = fake.name()
            elif pii_type == 'email': f = fake.email()
            elif pii_type == 'phone': f = fake.phone_number()
            elif pii_type == 'address': f = fake.address().replace('\n', ', ')
            if len(f) == length:
                return f
        
        closest = None
        min_diff = float('inf')
        for _ in range(50):
            if pii_type == 'name': f = fake.name()
            elif pii_type == 'email': f = fake.email()
            elif pii_type == 'phone': f = fake.phone_number()
            elif pii_type == 'address': f = fake.address().replace('\n', ', ')
            diff = abs(len(f) - length)
            if diff < min_diff:
                min_diff, closest = diff, f
        return closest
    elif pii_type == 'ssn':
        return fake.ssn()
    else: # id and others
        return fake.lexify(text='?' * length)

def detect_format(text):
    text = text.strip().strip('"')
    if text.startswith('ISA*'):
        return 'edi_x12'
    else:
        return 'plain'

def convert_to_readable(text, format_type):
    if format_type == 'edi_x12':
        parsed_transactions = parse_edi_fallback(text)
        output = StringIO()
        format_output(parsed_transactions, file=output)
        return output.getvalue()
    else:
        return text

def redact_text(original_text, pii_list):
    redacted = original_text
    for pii in pii_list:
        redacted = re.sub(re.escape(pii), '[REDACTED]', redacted, flags=re.IGNORECASE)
    return redacted

def anonymize_text(original_text, pii_list):
    anonymized = original_text
    pii_map = {}
    for pii in pii_list:
        normalized = normalize_text(pii)
        pii_type = classify_pii(pii)
        if normalized not in pii_map:
            if pii_type == 'gender':
                base_gender = random.choice(['male', 'female'])
                fake_val = format_gender(base_gender, pii)
                pii_map[normalized] = fake_val
            else:
                fake_val = generate_fake(pii_type, len(pii), pii)
                pii_map[normalized] = fake_val
        else:
            fake_val = pii_map[normalized]
        anonymized = re.sub(re.escape(pii), fake_val, anonymized, flags=re.IGNORECASE)
    return anonymized

def parse_edi_fallback(edi_content):
    """
    Fallback parser that combines address parts into a single line for easier redaction.
    """
    print("Using fallback parser...")
    
    segments = edi_content.replace('~', '\n').split('\n')
    segments = [seg.strip() for seg in segments if seg.strip()]
    
    parsed_data = {
        'transaction_info': {}, 'patient_info': {}, 'provider_info': {},
        'service_info': {}, 'diagnosis_info': {}
    }
    
    for segment in segments:
        elements = segment.split('*')
        segment_id = elements[0]
        
        if segment_id == 'ST':
            parsed_data['transaction_info']['transaction_type'] = elements[1]
        elif segment_id == 'NM1':
            entity_type = elements[1]
            if entity_type == 'IL':  # Patient
                last_name = elements[3] if len(elements) > 3 else ''
                first_name = elements[4] if len(elements) > 4 else ''
                parsed_data['patient_info']['name'] = f"{first_name} {last_name}".strip()
                if len(elements) > 8:
                    parsed_data['patient_info']['id'] = elements[9]
            elif entity_type == 'SJ':  # Provider
                last_name = elements[3] if len(elements) > 3 else ''
                first_name = elements[4] if len(elements) > 4 else ''
                parsed_data['provider_info']['name'] = f"{first_name} {last_name}".strip()
                if len(elements) > 8:
                    parsed_data['provider_info']['npi'] = elements[9]
        elif segment_id == 'N3':
            # Store the first line of the address
            parsed_data['patient_info']['address'] = elements[1]
        elif segment_id == 'N4':
            # Combine City, State, and Zip with the address line
            city = elements[1] if len(elements) > 1 else ''
            state = elements[2] if len(elements) > 2 else ''
            zip_code = elements[3] if len(elements) > 3 else ''
            
            full_address_parts = [city, state, zip_code]
            
            # If an address line already exists from N3, prepend it
            if 'address' in parsed_data['patient_info']:
                full_address_parts.insert(0, parsed_data['patient_info']['address'])
            
            # Join all parts with ", " and filter out any empty parts
            parsed_data['patient_info']['address'] = ", ".join(filter(None, full_address_parts))

        elif segment_id == 'DMG':
            parsed_data['patient_info']['dob'] = elements[2]
            parsed_data['patient_info']['gender'] = 'Female' if elements[3] == 'F' else 'Male'
        elif segment_id == 'UM':
            parsed_data['service_info']['service_type'] = elements[1]
            parsed_data['service_info']['request_category'] = elements[2]
            parsed_data['service_info']['service_code'] = elements[3]
            if len(elements) > 4:
                parsed_data['service_info']['quantity'] = elements[4]
        elif segment_id == 'HI':
            diagnosis_info = elements[1].split(':')
            if len(diagnosis_info) > 1:
                parsed_data['diagnosis_info']['code_qualifier'] = diagnosis_info[0]
                parsed_data['diagnosis_info']['diagnosis_code'] = diagnosis_info[1]
    
    return [{
        'transaction_type': parsed_data['transaction_info'].get('transaction_type', 'Unknown'),
        'parsed_data': {
            'description': 'Health Care Services Review',
            'patient_info': parsed_data['patient_info'],
            'provider_info': parsed_data['provider_info'],
            'service_info': parsed_data['service_info'],
            'diagnosis_info': parsed_data['diagnosis_info']
        }
    }]

def format_output(parsed_transactions, file=None):
    """Format parsed data for display"""
    output_lines = []
    for i, transaction in enumerate(parsed_transactions):
        output_lines.append(f"\n=== TRANSACTION {i+1} ===")
        output_lines.append(f"Transaction Type: {transaction['transaction_type']}")
        
        if 'parsed_data' in transaction:
            data = transaction['parsed_data']
            
            if 'description' in data:
                output_lines.append(f"Description: {data['description']}")
            
            # Patient Information
            if 'patient_info' in data and data['patient_info']:
                output_lines.append("\nPATIENT INFORMATION:")
                for key, value in data['patient_info'].items():
                    output_lines.append(f"  {key.replace('_', ' ').title()}: {value}")
            
            # Provider Information
            if 'provider_info' in data and data['provider_info']:
                output_lines.append("\nPROVIDER INFORMATION:")
                for key, value in data['provider_info'].items():
                    output_lines.append(f"  {key.replace('_', ' ').title()}: {value}")
            
            # Service Information
            if 'service_info' in data and data['service_info']:
                output_lines.append("\nSERVICE INFORMATION:")
                for key, value in data['service_info'].items():
                    output_lines.append(f"  {key.replace('_', ' ').title()}: {value}")
            
            # Diagnosis Information
            if 'diagnosis_info' in data and data['diagnosis_info']:
                output_lines.append("\nDIAGNOSIS INFORMATION:")
                for key, value in data['diagnosis_info'].items():
                    output_lines.append(f"  {key.replace('_', ' ').title()}: {value}")
    output_str = '\n'.join(output_lines)
    if file:
        file.write(output_str)
    else:
        print(output_str)
    return output_str

# --- Main Gradio Processing Function ---
def analyze_document(file, progress=gr.Progress()):
    """
    This function takes an uploaded file, processes it through the PII detection pipeline,
    and returns the annotated images, a redacted PDF, and a summary report.
    """
    load_models()
    if file is None:
        return None, None, None, None, None, "Please upload a document to begin."
    
    unique_id = uuid.uuid4().hex
    if hasattr(file, 'name'):
        filepath = file.name
    else:
        filepath = str(file)
    filename = os.path.basename(filepath)
    extension = os.path.splitext(filename)[1]
    
    if extension.lower() == '.csv':
        rows = []
        with open(filepath, 'r', newline='') as csvfile:
            csv_reader = csv.reader(csvfile)
            for row in csv_reader:
                if row:
                    rows.append(row[0])
        total_rows = len(rows)
        print(f"Processing {total_rows} rows for job {unique_id}...")
        progress(0.1, desc="Reading CSV rows...")
        report = f"## 🔍 Analysis Report for CSV\n**Total Rows:** {total_rows}\n\n---\n"
        redacted_rows = []
        anonymized_rows = []
        for i, text in enumerate(rows):
            progress(0.4 + (i / total_rows * 0.5), desc=f"Processing Row {i+1}/{total_rows}...")
            format_type = detect_format(text)
            readable_text = convert_to_readable(text, format_type)
            pii_list, llama_raw_output = detect_pii_from_combined_text(readable_text)
            redacted_text = redact_text(readable_text, pii_list)
            anonymized_text = anonymize_text(readable_text, pii_list)
            redacted_rows.append(redacted_text)
            anonymized_rows.append(anonymized_text)
            report += f"### 📄 Row {i+1}\n- **Text Detections:** {len(pii_list)}\n- **PII Found:** {', '.join(pii_list) if pii_list else 'None'}\n\n"
        progress(0.9, desc="Generating final CSVs...")
        redacted_csv_path = os.path.join(RESULTS_FOLDER, f"redacted_{unique_id}.csv")
        with open(redacted_csv_path, 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)
            for txt in redacted_rows:
                writer.writerow([txt])
        anonymized_csv_path = os.path.join(RESULTS_FOLDER, f"anonymized_{unique_id}.csv")
        with open(anonymized_csv_path, 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)
            for txt in anonymized_rows:
                writer.writerow([txt])
        progress(1, desc="Complete!")
        print("Processing Complete.")
        return [], [], [], redacted_csv_path, anonymized_csv_path, report
    
    progress(0, desc="Converting document to images...")
    images_to_process = []
    try:
        if extension.lower() == '.pdf':
            try:
                images_to_process = [cv2.cvtColor(np.array(page), cv2.COLOR_RGB2BGR) for page in convert_from_path(filepath, dpi=300)]
            except Exception as e:
                print(f"PDF conversion error: {e}. Trying fallback method...")
                try:
                    images_to_process = [cv2.cvtColor(np.array(page), cv2.COLOR_RGB2BGR) for page in convert_from_path(filepath, dpi=150)]
                except Exception as e2:
                    return None, None, None, None, None, f"🔴 **Error:** Could not process PDF. Please ensure Poppler is installed.\nDetails: {e2}"
        else:
            img = cv2.imread(filepath)
            if img is not None:
                images_to_process.append(img)
    except Exception as e:
        return None, None, None, None, None, f"🔴 **Error:** Could not process file. Details: {e}"
    
    if not images_to_process:
        return None, None, None, None, None, "🔴 **Error:** No pages could be extracted from the document."
    
    total_pages = len(images_to_process)
    print(f"Processing {total_pages} pages for job {unique_id}...")
    
    progress(0.1, desc="Extracting text from all pages (OCR)...")
    combined_text, all_pages_data = "", []
    for i, img_cv in enumerate(images_to_process):
        page_text = ""
        if reader:
            page_text = extract_sentence_text(reader.readtext(img_cv, paragraph=True))
            combined_text += f"\n--- Page {i+1} ---\n{page_text}\n"
        all_pages_data.append({"img_cv": img_cv, "page_num": i + 1})
    
    progress(0.4, desc="Analyzing text for PII with LLM...")
    global_pii_list, llama_raw_output = detect_pii_from_combined_text(combined_text)
    
    annotated_paths, redacted_paths, anonymized_paths = [], [], []
    redacted_pils, anonymized_pils = [], []
    report = f"## 🔍 Analysis Report\n**Global PII Found:** `{', '.join(global_pii_list) if global_pii_list else 'None'}`\n\n---\n"
    pii_map = {}
    
    for i, page_info in enumerate(all_pages_data):
        progress(0.5 + (i / total_pages * 0.4), desc=f"Processing Page {i+1}/{total_pages} (Visual & Text)...")
        img_cv, page_num = page_info["img_cv"], page_info["page_num"]
        detections = process_page_combined(img_cv, global_pii_list)
        
        annotated_img = img_cv.copy()
        redacted_img = img_cv.copy()
        anonymized_img = img_cv.copy()
        visual_count = sum(1 for d in detections if d["detection_type"] == "visual")
        text_count = sum(1 for d in detections if d.get("detection_type") == "text")
        
        for d in detections:
            bbox = d.get("bbox", [])
            if not bbox: continue
            x1, y1, x2, y2 = map(int, bbox)
            color = (0, 255, 0) if d.get("detection_type") == "visual" else (0, 0, 255)
            cv2.rectangle(annotated_img, (x1, y1), (x2, y2), color, 3)
            
            if d["detection_type"] == "visual":
                cv2.rectangle(redacted_img, (x1, y1), (x2, y2), (0, 0, 0), -1)
                cv2.rectangle(anonymized_img, (x1, y1), (x2, y2), (0, 0, 0), -1)
            else:
                bg_color = (255, 255, 255)
                height, width = img_cv.shape[:2]
                if x2 + 20 < width:
                    sample = img_cv[y1:y2, x2:x2+20]
                    if sample.size > 0:
                        bg_color = tuple(map(int, np.mean(sample, axis=(0,1))))
                else:
                    if x1 > 20:
                        sample = img_cv[y1:y2, x1-20:x1]
                        if sample.size > 0:
                            bg_color = tuple(map(int, np.mean(sample, axis=(0,1))))
                
                cv2.rectangle(redacted_img, (x1, y1), (x2, y2), (0, 0, 0), -1)
                cv2.rectangle(anonymized_img, (x1, y1), (x2, y2), bg_color, -1)
                
                original_text = d["text"]
                pii_type = classify_pii(original_text)
                normalized = normalize_text(original_text)
                key = 'gender' if pii_type == 'gender' else normalized
                
                if key not in pii_map:
                    if pii_type == 'gender':
                        base_gender = random.choice(['male', 'female'])
                        fake_text = format_gender(base_gender, original_text)
                        pii_map[key] = base_gender
                    else:
                        fake_text = generate_fake(pii_type, len(original_text), original_text)
                        pii_map[key] = fake_text
                else:
                    if pii_type == 'gender':
                        base_gender = pii_map[key]
                        fake_text = format_gender(base_gender, original_text)
                    else:
                        fake_text = pii_map[key]
                
                font = cv2.FONT_HERSHEY_SIMPLEX
                font_scale = (y2 - y1) / 40.0
                thickness = 2
                text_size, _ = cv2.getTextSize(fake_text, font, font_scale, thickness)
                box_width = x2 - x1
                if text_size[0] > box_width - 10:
                    font_scale *= (box_width - 10) / text_size[0]
                    text_size, _ = cv2.getTextSize(fake_text, font, font_scale, thickness)
                text_x = x1 + (box_width - text_size[0]) // 2
                text_y = y1 + ((y2 - y1) + text_size[1]) // 2
                bg_brightness = 0.299 * bg_color[2] + 0.587 * bg_color[1] + 0.114 * bg_color[0]
                text_color = (0, 0, 0) if bg_brightness > 128 else (255, 255, 255)
                cv2.putText(anonymized_img, fake_text, (text_x, text_y), font, font_scale, text_color, thickness)

        # Save all three versions of the image for the galleries
        annotated_path = os.path.join(RESULTS_FOLDER, f"annotated_{unique_id}_{page_num}.jpg")
        redacted_path = os.path.join(RESULTS_FOLDER, f"redacted_preview_{unique_id}_{page_num}.jpg")
        anonymized_path = os.path.join(RESULTS_FOLDER, f"anonymized_preview_{unique_id}_{page_num}.jpg")

        cv2.imwrite(annotated_path, annotated_img)
        cv2.imwrite(redacted_path, redacted_img)
        cv2.imwrite(anonymized_path, anonymized_img)

        annotated_paths.append(annotated_path)
        redacted_paths.append(redacted_path)
        anonymized_paths.append(anonymized_path)
        
        redacted_pils.append(Image.fromarray(cv2.cvtColor(redacted_img, cv2.COLOR_BGR2RGB)))
        anonymized_pils.append(Image.fromarray(cv2.cvtColor(anonymized_img, cv2.COLOR_BGR2RGB)))
        
        report += f"### 📄 Page {page_num}\n- **Visual Detections (🟩 Green):** {visual_count}\n- **Text Detections (🟥 Red):** {text_count}\n"
    
    progress(0.9, desc="Generating final PDFs...")
    redacted_pdf_path, anonymized_pdf_path = None, None
    if redacted_pils:
        pdf_path = os.path.join(RESULTS_FOLDER, f"redacted_{unique_id}.pdf")
        redacted_pils[0].save(pdf_path, "PDF", resolution=100.0, save_all=True, append_images=redacted_pils[1:])
        redacted_pdf_path = pdf_path
    if anonymized_pils:
        pdf_path = os.path.join(RESULTS_FOLDER, f"anonymized_{unique_id}.pdf")
        anonymized_pils[0].save(pdf_path, "PDF", resolution=100.0, save_all=True, append_images=anonymized_pils[1:])
        anonymized_pdf_path = pdf_path
        
    progress(1, desc="Complete!")
    print("Processing Complete.")
    
    return annotated_paths, redacted_paths, anonymized_paths, redacted_pdf_path, anonymized_pdf_path, report

# --- Gradio Interface Definition ---
title = "🔒 Combined PII Detection System"
description = """
### Advanced Multi-Modal PII Detection
This system uses a combination of visual and textual analysis to detect and redact Personally Identifiable Information from your documents.
- **🖼️ Visual Detection (YOLO):** Detects Faces, QR Codes, and Signatures.
- **📝 Text Detection (OCR + LLM):** Detects Names, Addresses, Phone Numbers, IDs, and other contextual PII.
**How to Use:**
1. Upload a document (PDF or image format).
2. The system will process each page and display annotated previews with colored boxes.
3. A fully redacted PDF with blacked-out PII is generated for you to download.
4. An analysis report summarizes the findings for each page.
"""
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(f"<h1 style='margin-bottom: 0.25rem;'>{title}</h1>")
    gr.Markdown(
        "<p style='color:#475569; line-height:1.6;'>"
        "Upload a PDF or image to detect and redact PII using visual detectors (🟩) and text analysis (🟥). "
        "Fixed for local environment with proper YOLO support."
        "</p>"
    )
    with gr.Accordion("About this tool", open=False):
        gr.Markdown(description)
    with gr.Tabs():
        with gr.Tab("Run"):
            with gr.Row():
                with gr.Column(scale=1):
                    file_input = gr.File(
                        label="Upload Document",
                        file_types=['.pdf', '.jpg', '.jpeg', '.png', '.bmp', '.csv'],
                        file_count="single",
                        height=100
                    )
                    submit_btn = gr.Button("🚀 Analyze Document", variant="primary")
                    with gr.Accordion("Tips", open=False):
                        gr.Markdown(
                            "- Prefer high-resolution files for better OCR results (300 DPI for PDFs).\n"
                            "- For PDFs, ensure Poppler is installed on your system.\n"
                            "- Visual detections are drawn in green; text-based detections are in red.\n"
                            "- Use the Previews tab to inspect annotated pages and the Report tab to download the redacted PDF."
                        )
                with gr.Column(scale=1):
                    gr.Markdown("### What happens during analysis")
                    gr.Markdown(
                        "- Convert pages to images\n"
                        "- Run global OCR to build combined text\n"
                        "- Use LLM to extract possible PII strings\n"
                        "- Match PII back to words and merge boxes\n"
                        "- Render annotated previews and build a redacted PDF"
                    )
                    clear_btn = gr.Button("🧹 Clear Results", variant="secondary")
        
        with gr.Tab("Annotated Preview (Detection)"):
            gr.Markdown("### Annotated Previews (🟩 Visual, 🟥 Text)")
            annotated_gallery_output = gr.Gallery(
                label="Annotated Pages", show_label=False, elem_id="gallery_annotated",
                columns=[2], rows=[1], object_fit="contain", height=480
            )
            
        with gr.Tab("Redacted Preview"):
            gr.Markdown("### Redacted Previews (Blacked Out)")
            redacted_gallery_output = gr.Gallery(
                label="Redacted Pages", show_label=False, elem_id="gallery_redacted",
                columns=[2], rows=[1], object_fit="contain", height=480
            )

        with gr.Tab("Anonymized Preview"):
            gr.Markdown("### Anonymized Previews (Fake Data)")
            anonymized_gallery_output = gr.Gallery(
                label="Anonymized Pages", show_label=False, elem_id="gallery_anonymized",
                columns=[2], rows=[1], object_fit="contain", height=480
            )

        with gr.Tab("Report & Download"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Downloads")
                    redacted_file_output = gr.File(label="Redacted PDF (Blacked Out)")
                    anonymized_file_output = gr.File(label="Anonymized PDF (Fake Data)")
                with gr.Column(scale=2):
                    gr.Markdown("### Analysis Report")
                    report_output = gr.Markdown(label="Analysis Report")

    outputs_list = [
        annotated_gallery_output, 
        redacted_gallery_output, 
        anonymized_gallery_output,
        redacted_file_output, 
        anonymized_file_output, 
        report_output
    ]

    submit_btn.click(
        fn=analyze_document,
        inputs=file_input,
        outputs=outputs_list
    )
    
    clear_btn.click(
        fn=lambda: ([], [], [], None, None, "Ready. Upload a document and click Analyze."),
        inputs=None,
        outputs=outputs_list
    )

if __name__ == "__main__":
    demo.queue().launch(server_port=8000)