gradio_api.py

import os
import cv2
import numpy as np
import math
import time
import gc
import torch
import fitz  # PyMuPDF
from PIL import Image
import gradio as gr
from paddleocr import PaddleOCR
from google import genai
from google.genai import types
from pydantic import BaseModel
import concurrent.futures


# --- 1. Structured Output Definitions ---
class BoundingBox(BaseModel):
    box_2d: list[int]
    text: str

# API Key configuration (cleared for security, now loaded via environment variables)
API_KEYS = []

def get_gemini_api_keys():
    import os
    import json
    
    # 1. Check for GEMINI_API_KEYS (comma-separated env var)
    env_keys_str = os.environ.get("GEMINI_API_KEYS")
    if env_keys_str:
        keys = [k.strip() for k in env_keys_str.split(",") if k.strip()]
        if keys:
            return keys

    # 2. Check numbered env vars (GEMINI_API_KEY, GEMINI_API_KEY_2, etc.)
    env_keys = []
    default_key = os.environ.get("GEMINI_API_KEY")
    if default_key:
        env_keys.append(default_key.strip())
        
    for i in range(1, 10):
        k = os.environ.get(f"GEMINI_API_KEY_{i}") or os.environ.get(f"GEMINI_API_KEY{i}")
        if k:
            k_clean = k.strip()
            if k_clean and k_clean not in env_keys:
                env_keys.append(k_clean)
                
    if env_keys:
        return env_keys

    # 3. Fallback to hardcoded keys
    valid_keys = [k.strip() for k in API_KEYS if k.strip()]
    if valid_keys:
        return valid_keys
    return []

def fetch_gemini_ocr_for_page(page_num, img_bytes, api_keys, start_key_idx, prompt, mode):
    """
    Runs Gemini OCR call for a single page inside a thread.
    Tries different API keys in a loop if one hits a rate limit or quota.
    Handles fallback to gemini-2.5-flash.
    """
    from google import genai
    from google.genai import types
    import time
    
    num_keys = len(api_keys)
    max_attempts = num_keys * 3  # Try each key up to 3 times
    
    for attempt in range(max_attempts):
        key_idx = (start_key_idx + attempt) % num_keys
        api_key = api_keys[key_idx]
        
        # If we have cycled through all keys at least once, sleep 5 seconds
        if attempt >= num_keys and attempt % num_keys == 0:
            print(f"[API] Seite {page_num+1}: Alle Keys einmal versucht. Schlafe 5 Sekunden...")
            time.sleep(5)
            
        client = genai.Client(api_key=api_key)
        
        # Try gemini-3.1-flash-lite first, fallback to gemini-2.5-flash if needed
        for model in ['gemini-3.1-flash-lite', 'gemini-2.5-flash']:
            try:
                response = client.models.generate_content(
                    model=model,
                    contents=[prompt, types.Part.from_bytes(data=img_bytes, mime_type='image/png')],
                    config=types.GenerateContentConfig(
                        response_mime_type="application/json",
                        response_schema=list[BoundingBox],
                        temperature=0.0
                    )
                )
                if response.parsed:
                    # Filter out boxes that are invalid or don't have exactly 4 values in box_2d
                    valid_parsed = []
                    for box in response.parsed:
                        if hasattr(box, 'box_2d') and box.box_2d and len(box.box_2d) == 4:
                            valid_parsed.append(box)
                        else:
                            print(f"[API] Warning: Filtered out invalid box on page {page_num+1}: {box}")
                    return page_num, valid_parsed, None
                return page_num, response.parsed, None
            except Exception as e:
                error_msg = str(e)
                print(f"[API] Fehler auf Seite {page_num+1} mit Key-Index {key_idx} (Modell {model}): {error_msg}")
                
                # Check if it's a rate limit, quota issue, or server error
                is_quota_or_rate = any(code in error_msg for code in ["429", "Quota", "exhausted", "ResourceExhausted", "limit"])
                is_server_err = any(code in error_msg for code in ["503", "500", "502", "504", "unavailable"])
                
                if not (is_quota_or_rate or is_server_err):
                    # For non-retriable errors (like bad requests), fail this page immediately
                    return page_num, None, e
                
                # For rate limit or temporary server error, break the model loop to try the next key immediately
                break
                
    return page_num, None, Exception("Alle API-Schlüssel sind fehlgeschlagen oder im Limit.")

# Caching models for reuse
_PADDLE_OCR = None
def get_paddle_ocr():
    global _PADDLE_OCR
    if _PADDLE_OCR is None:
        print("[API] Initializing PaddleOCR...")
        _PADDLE_OCR = PaddleOCR(use_angle_cls=True, lang='de')
    return _PADDLE_OCR

_TROCR_PROCESSOR = None
_TROCR_MODEL = None
def get_trocr():
    global _TROCR_PROCESSOR, _TROCR_MODEL
    if _TROCR_PROCESSOR is None or _TROCR_MODEL is None:
        print("[API] Initializing TrOCR...")
        import logging as transformers_logging
        transformers_logging.getLogger("transformers").setLevel(transformers_logging.ERROR)
        
        onnx_path = "trocr_onnx"
        if os.path.exists(onnx_path):
            from transformers import TrOCRProcessor
            from optimum.onnxruntime import ORTModelForVision2Seq
            _TROCR_PROCESSOR = TrOCRProcessor.from_pretrained(onnx_path)
            _TROCR_MODEL = ORTModelForVision2Seq.from_pretrained(onnx_path, provider="CPUExecutionProvider")
        else:
            from transformers import TrOCRProcessor, VisionEncoderDecoderModel
            _TROCR_PROCESSOR = TrOCRProcessor.from_pretrained('microsoft/trocr-base-handwritten')
            _TROCR_MODEL = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-handwritten')
    return _TROCR_PROCESSOR, _TROCR_MODEL

def recursive_xy_cut(boxes_with_data):
    if len(boxes_with_data) <= 1:
        return boxes_with_data
        
    gap_threshold = 5  # Mindestabstand
    
    # 1. Horizontale Lücken berechnen
    y_intervals = sorted([(b[0][1], b[0][3]) for b in boxes_with_data])
    max_h_gap = 0
    h_gap_y = None
    max_y = y_intervals[0][1]
    for i in range(1, len(y_intervals)):
        if y_intervals[i][0] > max_y:
            gap = y_intervals[i][0] - max_y
            if gap > max_h_gap and gap > gap_threshold:
                max_h_gap = gap
                h_gap_y = (max_y + y_intervals[i][0]) / 2
        max_y = max(max_y, y_intervals[i][1])
        
    # 2. Vertikale Lücken berechnen
    x_intervals = sorted([(b[0][0], b[0][2]) for b in boxes_with_data])
    max_v_gap = 0
    v_gap_x = None
    max_x = x_intervals[0][1]
    for i in range(1, len(x_intervals)):
        if x_intervals[i][0] > max_x:
            gap = x_intervals[i][0] - max_x
            if gap > max_v_gap and gap > gap_threshold:
                max_v_gap = gap
                v_gap_x = (max_x + x_intervals[i][0]) / 2
        max_x = max(max_x, x_intervals[i][1])
        
    # 3. Entlang der GRÖSSTEN Lücke schneiden!
    if max_h_gap == 0 and max_v_gap == 0:
        return sorted(boxes_with_data, key=lambda b: (b[0][1], b[0][0]))
        
    if max_v_gap > max_h_gap:
        left_boxes = [b for b in boxes_with_data if (b[0][0]+b[0][2])/2 < v_gap_x]
        right_boxes = [b for b in boxes_with_data if b not in left_boxes]
        if len(left_boxes) > 0 and len(right_boxes) > 0:
            return recursive_xy_cut(left_boxes) + recursive_xy_cut(right_boxes)
    
    if h_gap_y is not None:
        top_boxes = [b for b in boxes_with_data if (b[0][1]+b[0][3])/2 < h_gap_y]
        bottom_boxes = [b for b in boxes_with_data if b not in top_boxes]
        if len(top_boxes) > 0 and len(bottom_boxes) > 0:
            return recursive_xy_cut(top_boxes) + recursive_xy_cut(bottom_boxes)
            
    return sorted(boxes_with_data, key=lambda b: (b[0][1], b[0][0]))

def process_pdf(input_file_path, mode, smart_skip=True, progress=gr.Progress()):
    if not input_file_path:
        raise gr.Error("Bitte lade ein PDF hoch.")

    print(f"[API] Processing {input_file_path} in mode: {mode} (smart_skip: {smart_skip})")
    progress(0, desc="Initialisiere Datei und Modelle...")

    # Load required clients and models
    api_keys = None
    if "Gemini" in mode:
        api_keys = get_gemini_api_keys()
    
    paddle_ocr = None
    if "PaddleOCR" in mode:
        paddle_ocr = get_paddle_ocr()

    # Load PDF
    doc = fitz.open(input_file_path)
    num_pages = len(doc)
    
    # Pre-analysis: Determine which pages need OCR and which need redaction
    pages_to_ocr = []
    pages_to_redact = []
    for i in range(num_pages):
        page = doc.load_page(i)
        has_text = len(page.get_text().strip()) > 20
        if has_text and smart_skip:
            print(f"[API] Seite {i+1} hat bereits Text und Smart-Skip ist aktiv. Überspringe OCR.")
        else:
            if has_text and not smart_skip:
                print(f"[API] Seite {i+1} hat bereits Text und Smart-Skip ist inaktiv. Wird später redigiert (Force-OCR).")
                pages_to_redact.append(i)
            pages_to_ocr.append(i)
            
    output_filename = f"searchable_{os.path.basename(input_file_path)}"
    output_path = os.path.join(os.path.dirname(input_file_path), output_filename)

    # ==========================================
    # MODUS 1: Schnell (Gemini Full-Page) - Parallel
    # ==========================================
    if mode == "Schnell (Gemini Full-Page)":
        pages_img_bytes = {}
        progress(0.05, desc="Rendere PDF-Seiten...")
        for i in pages_to_ocr:
            page = doc.load_page(i)
            zoom = 150 / 72
            mat = fitz.Matrix(zoom, zoom)
            pix = page.get_pixmap(matrix=mat)
            pages_img_bytes[i] = pix.tobytes("png")

        results = {}
        errors = {}
        num_keys = len(api_keys)
        
        progress(0.1, desc=f"Starte parallele Gemini Semantic Analyse (mit {num_keys} API-Schlüsseln)...")
        
        prompt = """Du bist ein extrem präzises OCR-System für mathematische Vorlesungsskripte.
Extrahiere absolut JEDEN Text (sowohl handgeschrieben als auch Maschinenschrift / gedruckten Text). 
Verpasse kein einziges mathematisches Symbol, keinen Bruch und keinen Index.
WICHTIG FÜR FORMELN: Wandle ALLE mathematischen Formeln zwingend in eine saubere, einzeilige und logisch lesbare Schreibweise um! 
- Nutze Klammern und Schrägstriche für Brüche: (A)/(B)
- Nutze '^' für Exponenten und '_' für Indizes: x^(SV), q_BM
- Nutze korrekte Unicode-Sonderzeichen für alles andere: Wurzeln (√), Integrale (∫), Summen (∑), griechische Buchstaben (α, β, γ, μ) etc.
- ACHTUNG BEI EINHEITEN: Wenn Einheiten in eckigen Klammeln [...] neben einer Formel stehen, behalte die eckigen Klammern UNBEDINGT bei! Füge KEIN Multiplikationszeichen '*' dazwischen ein. Einheiten sind reine Beschriftungen, keine Faktoren!
- Versuche NICHT, das optische 2D-Layout von Formeln mit mehrzeiligen Leerzeichen nachzuahmen!
Fasse zusammenhängende Sätze, Absätze oder komplette mathematische Formeln in EINER GEMEINSAMEN BoundingBox zusammen. 
Zerstückele Formeln oder Brüche NICHT in Einzelteile! Eine komplette Formel = Eine BoundingBox.
Ignoriere Hintergrundmuster wie Punktraster komplett.
Gib für jeden Textblock/jede Formel eine BoundingBox zurück. box_2d ist [ymin, xmin, ymax, xmax] von 0 bis 1000.
Speichere den erkannten Text bzw. die Formel im Feld 'text' der BoundingBox.
WARNUNG: Es ist strengstens verboten, als Wert für das Feld 'text' einfach nur das Platzhalterwort 'text' einzutragen! Schreibe dort immer den tatsächlich erkannten Text hinein."""

        if pages_to_ocr:
            max_workers = min(len(pages_to_ocr), 12)
            with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
                futures = []
                for idx, i in enumerate(pages_to_ocr):
                    # Pass the api_keys list and start key index to allow rotation
                    futures.append(executor.submit(fetch_gemini_ocr_for_page, i, pages_img_bytes[i], api_keys, idx, prompt, mode))
                    
                    completed = 0
                for future in concurrent.futures.as_completed(futures):
                    p_num, parsed, err = future.result()
                    results[p_num] = parsed
                    errors[p_num] = err
                    completed += 1
                    progress(completed / len(pages_to_ocr), desc=f"Gemini Semantic Analyse: {completed} von {len(pages_to_ocr)} Seiten abgeschlossen...")

        progress(0.95, desc="Generiere durchsuchbares PDF...")
        for page_num in pages_to_ocr:
            page = doc.load_page(page_num)
            if page_num in pages_to_redact:
                print(f"[API] Bereinige alten Textlayer auf Seite {page_num+1} (Force-OCR)...")
                try:
                    traces = page.get_texttrace()
                    has_visible_text = any(t.get("type") != 3 for t in traces)
                except Exception:
                    has_visible_text = len(page.get_text().strip()) > 0
                
                if has_visible_text:
                    pix = page.get_pixmap(dpi=150)
                    img_bytes = pix.tobytes("png")
                    page.add_redact_annot(page.rect)
                    page.apply_redactions(images=1, graphics=1, text=0)
                    page.insert_image(page.rect, stream=img_bytes)
                else:
                    page.add_redact_annot(page.rect)
                    page.apply_redactions(images=0, graphics=0, text=0)
            err = errors.get(page_num)
            if err:
                print(f"[API] Fehler auf Seite {page_num+1}: {err}")
                continue
            parsed_boxes = results.get(page_num)
            if parsed_boxes:
                font = fitz.Font("helv")
                descender = font.descender

                boxes_with_data = []
                for box in parsed_boxes:
                    if not box.box_2d or len(box.box_2d) != 4:
                        print(f"[API] Warning: Invalid box_2d length {len(box.box_2d) if box.box_2d else 0} on page {page_num+1}: {box}")
                        continue
                    ymin, xmin, ymax, xmax = box.box_2d
                    x0 = (xmin / 1000) * page.rect.width
                    y0 = (ymin / 1000) * page.rect.height
                    x1 = (xmax / 1000) * page.rect.width
                    y1 = (ymax / 1000) * page.rect.height
                    boxes_with_data.append(([x0, y0, x1, y1], box))

                sorted_data = recursive_xy_cut(boxes_with_data)

                for coords, box in sorted_data:
                    text = box.text
                    if not text.strip() or text.strip() in [".", "..."]:
                        continue

                    x0, y0, x1, y1 = coords
                    rect = fitz.Rect(x0, y0, x1, y1)

                    fontsize = rect.height
                    text_length = fitz.get_text_length(text, fontname="helv", fontsize=fontsize)
                    scale_x = rect.width / text_length if text_length > 0 else 1.0
                    
                    y_baseline = rect.y1 + (descender * fontsize)
                    point = fitz.Point(rect.x0, y_baseline)
                    matrix = fitz.Matrix(scale_x, 1.0)

                    try:
                        page.insert_text(point, text, fontsize=fontsize, fontname="helv", render_mode=3, morph=(point, matrix))
                    except ValueError:
                        clean_text = text.encode("latin-1", "ignore").decode("latin-1")
                        if clean_text.strip():
                            try:
                                page.insert_text(point, clean_text, fontsize=fontsize, fontname="helv", render_mode=3, morph=(point, matrix))
                            except Exception:
                                pass
            gc.collect()

    # ==========================================
    # MODUS 2: Präzise (True Hybrid OCR) - Parallel
    # ==========================================
    elif mode == "Präzise (Hybrid: PaddleOCR + Gemini)":
        pages_img_bytes = {}
        pages_paddle_boxes = {}
        for idx, i in enumerate(pages_to_ocr):
            progress((idx + 1) / (len(pages_to_ocr) * 2) if pages_to_ocr else 0.5, desc=f"PaddleOCR Geometrie Analyse: Seite {i + 1} von {num_pages}...")
            page = doc.load_page(i)
            zoom = 150 / 72
            mat = fitz.Matrix(zoom, zoom)
            pix = page.get_pixmap(matrix=mat)
            img_bytes = pix.tobytes("png")
            pages_img_bytes[i] = img_bytes
            
            img_np = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.h, pix.w, pix.n)
            if pix.n == 4:
                img_np = cv2.cvtColor(img_np, cv2.COLOR_RGBA2BGR)
            else:
                img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)

            result = paddle_ocr.ocr(img_np)
            paddle_boxes = [(line[0], line[1][0]) for line in result[0]] if result and result[0] else []
            pages_paddle_boxes[i] = paddle_boxes

        results = {}
        errors = {}
        num_keys = len(api_keys)
        
        progress(0.5, desc=f"Starte parallele Gemini Semantic Analyse (mit {num_keys} API-Schlüsseln)...")
        prompt = """Du bist ein extrem präzises OCR-System für mathematische Vorlesungsskripte.
Extrahiere absolut JEDEN Text (sowohl handgeschrieben als auch Maschinenschrift / gedruckten Text). 
Verpasse kein einziges mathematisches Symbol, keinen Bruch und keinen Index.
WICHTIG FÜR FORMELN: Wandle ALLE mathematischen Formeln zwingend in eine saubere, einzeilige und logisch lesbare Schreibweise um! 
- Nutze Klammern und Schrägstriche für Brüche: (A)/(B)
- Nutze '^' für Exponenten und '_' für Indizes: x^(SV), q_BM
- Nutze korrekte Unicode-Sonderzeichen für alles andere: Wurzeln (√), Integrale (∫), Summen (∑), griechische Buchstaben (α, β, γ, μ) etc.
- ACHTUNG BEI EINHEITEN: Wenn Einheiten in eckigen Klammeln [...] neben einer Formel stehen, behalte die eckigen Klammern UNBEDINGT bei! Füge KEIN Multiplikationszeichen '*' dazwischen ein. Einheiten sind reine Beschriftungen, keine Faktoren!
- Versuche NICHT, das optische 2D-Layout von Formeln mit mehrzeiligen Leerzeichen nachzuahmen!
WICHTIG FÜR DAS LAYOUT (ABSOLUT KRITISCH!): 
1. NORMALE TEXTZEILEN: Du MUSST für JEDE physische Textzeile im Bild eine EIGENE, separate BoundingBox erstellen!
   - Es ist STRENGSTENS VERBOTEN, mehrere Zeilen zu einem Absatz zusammenzufassen!
   - Auch wenn eine Textzeile Variablen (wie f_A) enthält, ist sie eine normale Zeile und darf NICHT mit der Zeile darunter zusammengefasst werden.
2. MEHRZEILIGE BRÜCHE: NUR WIRKLICHE mehrzeilige Formeln (Zähler über Nenner) MÜSSEN in EINER gemeinsamen BoundingBox zusammengefasst werden.
Ignoriere Hintergrundmuster wie Punktraster komplett.
Gib für jeden Textblock/jede Formel eine BoundingBox zurück. box_2d ist [ymin, xmin, ymax, xmax] von 0 bis 1000.
Speichere den erkannten Text bzw. die Formel im Feld 'text' der BoundingBox.
WARNUNG: Es ist strengstens verboten, als Wert für das Feld 'text' einfach nur das Platzhalterwort 'text' einzutragen! Schreibe dort immer den tatsächlich erkannten Text hinein."""

        if pages_to_ocr:
            max_workers = min(len(pages_to_ocr), 12)
            with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
                futures = []
                for idx, i in enumerate(pages_to_ocr):
                    # Pass the api_keys list and start key index to allow rotation
                    futures.append(executor.submit(fetch_gemini_ocr_for_page, i, pages_img_bytes[i], api_keys, idx, prompt, mode))
                    
                completed = 0
                for future in concurrent.futures.as_completed(futures):
                    p_num, parsed, err = future.result()
                    results[p_num] = parsed
                    errors[p_num] = err
                    completed += 1
                    progress(0.5 + (completed / len(pages_to_ocr)) * 0.5, desc=f"Gemini Semantic Analyse: {completed} von {len(pages_to_ocr)} Seiten abgeschlossen...")

        progress(0.95, desc="Generiere durchsuchbares PDF...")
        zoom = 150 / 72
        for page_num in pages_to_ocr:
            page = doc.load_page(page_num)
            if page_num in pages_to_redact:
                print(f"[API] Bereinige alten Textlayer auf Seite {page_num+1} (Force-OCR)...")
                try:
                    traces = page.get_texttrace()
                    has_visible_text = any(t.get("type") != 3 for t in traces)
                except Exception:
                    has_visible_text = len(page.get_text().strip()) > 0
                
                if has_visible_text:
                    pix = page.get_pixmap(dpi=150)
                    img_bytes = pix.tobytes("png")
                    page.add_redact_annot(page.rect)
                    page.apply_redactions(images=1, graphics=1, text=0)
                    page.insert_image(page.rect, stream=img_bytes)
                else:
                    page.add_redact_annot(page.rect)
                    page.apply_redactions(images=0, graphics=0, text=0)
            err = errors.get(page_num)
            if err:
                print(f"[API] Fehler auf Seite {page_num+1}: {err}")
                continue
            parsed_boxes = results.get(page_num)
            paddle_boxes = pages_paddle_boxes.get(page_num, [])
            
            if parsed_boxes:
                boxes_with_data = []
                g_rects = []
                valid_g_boxes = []
                for g_box in parsed_boxes:
                    if not g_box.box_2d or len(g_box.box_2d) != 4:
                        print(f"[API] Warning: Invalid box_2d length {len(g_box.box_2d) if g_box.box_2d else 0} on page {page_num+1}: {g_box}")
                        continue
                    ymin, xmin, ymax, xmax = g_box.box_2d
                    x0 = (xmin / 1000) * page.rect.width
                    y0 = (ymin / 1000) * page.rect.height
                    x1 = (xmax / 1000) * page.rect.width
                    y1 = (ymax / 1000) * page.rect.height
                    g_rects.append(fitz.Rect(x0, y0, x1, y1))
                    valid_g_boxes.append(g_box)
                    
                parsed_boxes = valid_g_boxes
                assigned_p_boxes_per_g_idx = {idx: [] for idx in range(len(parsed_boxes))}
                
                for pb_data in paddle_boxes:
                    p_box, p_text = pb_data
                    p_xmin = min(p[0] for p in p_box) / zoom
                    p_ymin = min(p[1] for p in p_box) / zoom
                    p_xmax = max(p[0] for p in p_box) / zoom
                    p_ymax = max(p[1] for p in p_box) / zoom
                    p_rect = fitz.Rect(p_xmin, p_ymin, p_xmax, p_ymax)
                    
                    best_g_idx = -1
                    max_overlap = 0
                    for idx, g_rect in enumerate(g_rects):
                        overlap = g_rect.intersect(p_rect).get_area()
                        if overlap > max_overlap:
                            max_overlap = overlap
                            best_g_idx = idx
                            
                    if best_g_idx != -1 and max_overlap > 0.1 * p_rect.get_area():
                        assigned_p_boxes_per_g_idx[best_g_idx].append(pb_data)
                        
                for idx, g_box in enumerate(parsed_boxes):
                    g_text = g_box.text
                    assigned_p_boxes = assigned_p_boxes_per_g_idx[idx]
                    
                    if assigned_p_boxes:
                        assigned_p_boxes.sort(key=lambda b: min(p[1] for p in b[0]))
                        
                        clustered_p_boxes = []
                        for pb_tuple in assigned_p_boxes:
                            pb, pt = pb_tuple
                            y_center = (min(p[1] for p in pb) + max(p[1] for p in pb)) / 2
                            added_to_cluster = False
                            for cluster in clustered_p_boxes:
                                if abs(y_center - cluster['y_center']) < (10 / zoom):
                                    cluster['boxes'].append(pb_tuple)
                                    all_y = [min(p[1] for b in cluster['boxes'] for p in b[0]), max(p[1] for b in cluster['boxes'] for p in b[0])]
                                    cluster['y_center'] = sum(all_y) / 2
                                    added_to_cluster = True
                                    break
                            if not added_to_cluster:
                                clustered_p_boxes.append({'y_center': y_center, 'boxes': [pb_tuple]})
                                
                        for cluster in clustered_p_boxes:
                            cluster['boxes'].sort(key=lambda b: min(p[0] for p in b[0]))
                            
                        math_chars = sum(1 for c in g_text if c in ['=', '/', '^', '[', ']'])
                        is_formula = (math_chars >= 4 and "=" in g_text)
                        
                        if not is_formula:
                            g_words = g_text.split()
                            word_idx = 0
                            for c_idx, cluster in enumerate(clustered_p_boxes):
                                cluster_boxes = cluster['boxes']
                                cluster_word_count = sum(max(1, len(pt.split())) for pb, pt in cluster_boxes)
                                chunk = g_words[word_idx : word_idx + cluster_word_count]
                                line_text = " ".join(chunk)
                                word_idx += cluster_word_count
                                if c_idx == len(clustered_p_boxes) - 1 and word_idx < len(g_words):
                                    if line_text: line_text += " "
                                    line_text += " ".join(g_words[word_idx:])
                                if not line_text.strip(): continue
                                
                                total_dx, total_dy = 0, 0
                                all_points = []
                                for pb, pt in cluster_boxes:
                                    total_dx += pb[1][0] - pb[0][0]
                                    total_dy += pb[1][1] - pb[0][1]
                                    for p in pb: all_points.append((p[0]/zoom, p[1]/zoom))
                                    
                                angle_rad = math.atan2(total_dy, total_dx) if (total_dx != 0 or total_dy != 0) else 0
                                angle_deg = math.degrees(angle_rad)
                                cos_a, sin_a = math.cos(-angle_rad), math.sin(-angle_rad)
                                local_points = [(px * cos_a - py * sin_a, px * sin_a + py * cos_a) for px, py in all_points]
                                min_lx, max_lx = min(p[0] for p in local_points), max(p[0] for p in local_points)
                                min_ly, max_ly = min(p[1] for p in local_points), max(p[1] for p in local_points)
                                cos_inv, sin_inv = math.cos(angle_rad), math.sin(angle_rad)
                                merged_box = [[lx * cos_inv - ly * sin_inv, lx * sin_inv + ly * cos_inv] for lx, ly in [(min_lx, min_ly), (max_lx, min_ly), (max_lx, max_ly), (min_lx, max_ly)]]
                                p0, p1, p2, p3 = merged_box
                                dx_up, dy_up = p0[0] - p3[0], p0[1] - p3[1]
                                font = fitz.Font("helv")
                                pdf_baseline = fitz.Point(p3[0] + dx_up * -font.descender, p3[1] + dy_up * -font.descender)
                                boxes_with_data.append(([min(p[0] for p in merged_box), min(p[1] for p in merged_box), max(p[0] for p in merged_box), max(p[1] for p in merged_box)], (line_text, pdf_baseline, math.hypot(p1[0]-p0[0], p1[1]-p0[1]), math.hypot(dx_up, dy_up), angle_deg)))
                        else:
                            total_dx, total_dy = 0, 0
                            all_points = []
                            for pb, pt in assigned_p_boxes:
                                total_dx += pb[1][0] - pb[0][0]
                                total_dy += pb[1][1] - pb[0][1]
                                for p in pb: all_points.append((p[0]/zoom, p[1]/zoom))
                            angle_rad = math.atan2(total_dy, total_dx) if (total_dx != 0 or total_dy != 0) else 0
                            angle_deg = math.degrees(angle_rad)
                            cos_a, sin_a = math.cos(-angle_rad), math.sin(-angle_rad)
                            local_points = [(px * cos_a - py * sin_a, px * sin_a + py * cos_a) for px, py in all_points]
                            min_lx, max_lx = min(p[0] for p in local_points), max(p[0] for p in local_points)
                            min_ly, max_ly = min(p[1] for p in local_points), max(p[1] for p in local_points)
                            cos_inv, sin_inv = math.cos(angle_rad), math.sin(angle_rad)
                            merged_box = [[lx * cos_inv - ly * sin_inv, lx * sin_inv + ly * cos_inv] for lx, ly in [(min_lx, min_ly), (max_lx, min_ly), (max_lx, max_ly), (min_lx, max_ly)]]
                            p0, p1, p2, p3 = merged_box
                            dx_up, dy_up = p0[0] - p3[0], p0[1] - p3[1]
                            font = fitz.Font("helv")
                            pdf_baseline = fitz.Point(p3[0] + dx_up * -font.descender, p3[1] + dy_up * -font.descender)
                            boxes_with_data.append(([min(p[0] for p in merged_box), min(p[1] for p in merged_box), max(p[0] for p in merged_box), max(p[1] for p in merged_box)], (g_text.replace('\n', '  '), pdf_baseline, math.hypot(p1[0]-p0[0], p1[1]-p0[1]), math.hypot(dx_up, dy_up), angle_deg)))
                    else:
                        ymin, xmin, ymax, xmax = g_box.box_2d
                        x0, y0, x1, y1 = (xmin / 1000) * page.rect.width, (ymin / 1000) * page.rect.height, (xmax / 1000) * page.rect.width, (ymax / 1000) * page.rect.height
                        boxes_with_data.append(([x0, y0, x1, y1], (g_text.replace('\n', '  '), fitz.Point(x0, y1 - (y1-y0)*0.2), x1-x0, y1-y0, 0)))

                sorted_data = recursive_xy_cut(boxes_with_data)
                for coords, data in sorted_data:
                    text, pdf_baseline, box_width_pdf, box_height_pdf, angle_deg = data
                    fontsize = box_height_pdf
                    text_length = fitz.get_text_length(text, fontname="helv", fontsize=fontsize)
                    matrix = fitz.Matrix(box_width_pdf / text_length if text_length > 0 else 1.0, 1.0) * fitz.Matrix(-angle_deg)
                    try:
                        page.insert_text(pdf_baseline, text, fontsize=fontsize, fontname="helv", render_mode=3, morph=(pdf_baseline, matrix))
                    except Exception: pass
        gc.collect()

    # ==========================================
    # MODUS 3: Lokal Deep (TrOCR) - Sequentiell
    # ==========================================
    elif mode == "Lokal Deep (PaddleOCR + TrOCR)":
        trocr_processor, trocr_model = get_trocr()
        for idx, page_num in enumerate(pages_to_ocr):
            progress(idx / len(pages_to_ocr) if pages_to_ocr else 1.0, desc=f"Verarbeite Seite {page_num + 1} von {num_pages}...")
            page = doc.load_page(page_num)
            zoom = 3.0
            mat = fitz.Matrix(zoom, zoom)
            pix = page.get_pixmap(matrix=mat)
            img_np = cv2.cvtColor(np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.h, pix.w, pix.n), cv2.COLOR_RGBA2BGR if pix.n == 4 else cv2.COLOR_RGB2BGR)
            if page_num in pages_to_redact:
                print(f"[API] Bereinige alten Textlayer auf Seite {page_num+1} (Force-OCR)...")
                try:
                    traces = page.get_texttrace()
                    has_visible_text = any(t.get("type") != 3 for t in traces)
                except Exception:
                    has_visible_text = len(page.get_text().strip()) > 0
                
                if has_visible_text:
                    pix = page.get_pixmap(dpi=150)
                    img_bytes = pix.tobytes("png")
                    page.add_redact_annot(page.rect)
                    page.apply_redactions(images=1, graphics=1, text=0)
                    page.insert_image(page.rect, stream=img_bytes)
                else:
                    page.add_redact_annot(page.rect)
                    page.apply_redactions(images=0, graphics=0, text=0)

            result = paddle_ocr.ocr(img_np)
            if not result or not result[0]: continue
            
            crops, valid_boxes = [], []
            for line in [l for l in result[0] if l]:
                box = line[0]
                x_coords, y_coords = [int(p[0]) for p in box], [int(p[1]) for p in box]
                crop_img = img_np[max(0, min(y_coords) - 2):min(img_np.shape[0], max(y_coords) + 2), max(0, min(x_coords) - 2):min(img_np.shape[1], max(x_coords) + 2)]
                if crop_img.size > 0:
                    crops.append(Image.fromarray(cv2.cvtColor(crop_img, cv2.COLOR_BGR2RGB)))
                    valid_boxes.append((box, line[1][0]))
            
            results = []
            for b_idx in range(0, len(crops), 4):
                try:
                    batch_texts = trocr_processor.batch_decode(trocr_model.generate(trocr_processor(crops[b_idx:b_idx+4], return_tensors="pt").pixel_values, max_new_tokens=30), skip_special_tokens=True)
                    results.extend(batch_texts)
                except Exception: results.extend([vb[1] for vb in valid_boxes[b_idx:b_idx+4]])

            boxes_with_data = [([min(p[0] for p in vb[0]), min(p[1] for p in vb[0]), max(p[0] for p in vb[0]), max(p[1] for p in vb[0])], (vb[0], results[idx])) for idx, vb in enumerate(valid_boxes) if results[idx].strip()]
            for coords, (box, text) in recursive_xy_cut(boxes_with_data):
                p0, p1, p3 = box[0], box[1], box[3]
                angle_deg = math.degrees(math.atan2(p1[1]-p0[1], p1[0]-p0[0]))
                font = fitz.Font("helv")
                base_x, base_y = (p3[0] + (p0[0]-p3[0]) * -font.descender) / zoom, (p3[1] + (p0[1]-p3[1]) * -font.descender) / zoom
                fontsize = math.hypot(p0[0]-p3[0], p0[1]-p3[1]) / zoom
                matrix = fitz.Matrix((math.hypot(p1[0]-p0[0], p1[1]-p0[1]) / zoom) / fitz.get_text_length(text, fontname="helv", fontsize=fontsize), 1.0) * fitz.Matrix(-angle_deg)
                try:
                    page.insert_text(fitz.Point(base_x, base_y), text, fontsize=fontsize, fontname="helv", render_mode=3, morph=(fitz.Point(base_x, base_y), matrix))
                except Exception: pass

    # ==========================================
    # MODUS 4: Lokal Schnell (PaddleOCR) - Sequentiell
    # ==========================================
    elif mode == "Lokal Schnell (PaddleOCR)":
        for idx, page_num in enumerate(pages_to_ocr):
            progress(idx / len(pages_to_ocr) if pages_to_ocr else 1.0, desc=f"Verarbeite Seite {page_num + 1} von {num_pages} (Lokal Schnell)...")
            page = doc.load_page(page_num)
            zoom = 3.0
            mat = fitz.Matrix(zoom, zoom)
            pix = page.get_pixmap(matrix=mat)
            img_np = cv2.cvtColor(np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.h, pix.w, pix.n), cv2.COLOR_RGBA2BGR if pix.n == 4 else cv2.COLOR_RGB2BGR)
            if page_num in pages_to_redact:
                print(f"[API] Bereinige alten Textlayer auf Seite {page_num+1} (Force-OCR)...")
                try:
                    traces = page.get_texttrace()
                    has_visible_text = any(t.get("type") != 3 for t in traces)
                except Exception:
                    has_visible_text = len(page.get_text().strip()) > 0
                
                if has_visible_text:
                    pix = page.get_pixmap(dpi=150)
                    img_bytes = pix.tobytes("png")
                    page.add_redact_annot(page.rect)
                    page.apply_redactions(images=1, graphics=1, text=0)
                    page.insert_image(page.rect, stream=img_bytes)
                else:
                    page.add_redact_annot(page.rect)
                    page.apply_redactions(images=0, graphics=0, text=0)

            result = paddle_ocr.ocr(img_np)
            if not result or not result[0]: continue
            
            page_data = result[0]
            boxes_with_data = []
            for line in page_data:
                if not line: continue
                box = line[0]
                text = line[1][0]
                if not text.strip(): continue
                
                xmin = min(p[0] for p in box)
                ymin = min(p[1] for p in box)
                xmax = max(p[0] for p in box)
                ymax = max(p[1] for p in box)
                boxes_with_data.append(([xmin, ymin, xmax, ymax], (box, text)))
                
            sorted_data = recursive_xy_cut(boxes_with_data)
            
            for coords, (box, text) in sorted_data:
                p0, p1, p3 = box[0], box[1], box[3]
                angle_deg = math.degrees(math.atan2(p1[1]-p0[1], p1[0]-p0[0]))
                font = fitz.Font("helv")
                base_x, base_y = (p3[0] + (p0[0]-p3[0]) * -font.descender) / zoom, (p3[1] + (p0[1]-p3[1]) * -font.descender) / zoom
                fontsize = math.hypot(p0[0]-p3[0], p0[1]-p3[1]) / zoom
                text_length = fitz.get_text_length(text, fontname="helv", fontsize=fontsize)
                scale_x = (math.hypot(p1[0]-p0[0], p1[1]-p0[1]) / zoom) / text_length if text_length > 0 else 1.0
                matrix = fitz.Matrix(scale_x, 1.0) * fitz.Matrix(-angle_deg)
                try:
                    page.insert_text(fitz.Point(base_x, base_y), text, fontsize=fontsize, fontname="helv", render_mode=3, morph=(fitz.Point(base_x, base_y), matrix))
                except Exception: pass

    doc.save(output_path)
    doc.close()
    progress(1.0, desc="Fertig! PDF erfolgreich generiert.")
    print(f"[API] Searchable PDF saved to {output_path}")

    # Copy to history directory for persistent Web UI history
    try:
        history_dir = "output_files"
        os.makedirs(history_dir, exist_ok=True)
        dest_filename = f"ocr_{os.path.basename(input_file_path)}"
        if not dest_filename.lower().endswith(".pdf"):
            dest_filename += ".pdf"
        dest_path = os.path.join(history_dir, dest_filename)
        
        import shutil
        shutil.copy2(output_path, dest_path)
        print(f"[API] Copied output to history (overwritten if existed): {dest_path}")
    except Exception as e:
        print(f"[API] Warning: Failed to copy file to history: {e}")


    return output_path

HISTORY_DIR = "output_files"

def get_history_files():
    if not os.path.exists(HISTORY_DIR):
        return []
    import glob
    # get all pdfs and zips
    files = glob.glob(os.path.join(HISTORY_DIR, "*.pdf")) + glob.glob(os.path.join(HISTORY_DIR, "*.zip"))
    # sort by modification time descending
    files.sort(key=os.path.getmtime, reverse=True)
    return files

def get_history_html():
    files = get_history_files()
    # Filter out zip files to prevent duplicate entries
    files = [f for f in files if not f.endswith(".zip")]
    
    if not files:
        return "<div style='color: #94a3b8; text-align: center; padding: 20px; font-family: system-ui, sans-serif;'>Keine verarbeiteten Dateien vorhanden.</div>"
    
    html = """
    <div style="max-height: 250px; overflow-y: auto; border: 1px solid rgba(255,255,255,0.1); border-radius: 8px; background: rgba(30,41,59,0.4); padding: 5px;">
      <table style="width: 100%; border-collapse: collapse; text-align: left; font-family: system-ui, sans-serif; color: white;">
        <thead>
          <tr style="border-bottom: 2px solid rgba(255,255,255,0.1); color: #94a3b8; font-size: 13px;">
            <th style="padding: 10px; font-weight: 600;">Dateiname</th>
            <th style="padding: 10px; text-align: right; font-weight: 600; width: 100px;">Größe</th>
            <th style="padding: 10px; text-align: right; font-weight: 600; width: 150px;">Aktion</th>
          </tr>
        </thead>
        <tbody>
    """
    for path in files:
        filename = os.path.basename(path)
        try:
            size_bytes = os.path.getsize(path)
            size_mb = size_bytes / (1024 * 1024)
            size_str = f"{size_mb:.1f} MB"
        except Exception:
            size_str = "unbekannt"
            
        # Point download url directly to the /file= endpoint served by Gradio (requires absolute path)
        import urllib.parse
        abs_path = os.path.abspath(path).replace(os.sep, '/')
        encoded_path = urllib.parse.quote(abs_path, safe="/")
        
        space_id = os.environ.get("SPACE_ID")
        if space_id:
            subdomain = space_id.lower().replace("/", "-")
            base_url = f"https://{subdomain}.hf.space/gradio_api"
        else:
            base_url = ""
            
        download_url = f"{base_url}/file={encoded_path}"
        
        html += f"""
          <tr style="border-bottom: 1px solid rgba(255,255,255,0.05); font-size: 14px;">
            <td style="padding: 10px; max-width: 320px; overflow: hidden; text-overflow: ellipsis; white-space: nowrap;">📄 {filename}</td>
            <td style="padding: 10px; text-align: right; color: #94a3b8;">{size_str}</td>
            <td style="padding: 10px; text-align: right;">
              <a href="{download_url}" target="_blank" download="{filename}" style="display: inline-flex; align-items: center; justify-content: center; background: #007AFF; color: white; padding: 6px 12px; border-radius: 6px; text-decoration: none; font-size: 12px; font-weight: bold; transition: all 0.2s;" onmouseover="this.style.background='#0051a8'" onmouseout="this.style.background='#007AFF'">
                📥 Herunterladen
              </a>
            </td>
          </tr>
        """
    html += """
        </tbody>
      </table>
    </div>
    """
    return html

def clear_history():
    if os.path.exists(HISTORY_DIR):
        import shutil
        for f in os.listdir(HISTORY_DIR):
            try:
                path = os.path.join(HISTORY_DIR, f)
                if os.path.isfile(path):
                    os.remove(path)
                elif os.path.isdir(path):
                    shutil.rmtree(path)
            except Exception:
                pass
    return []

def clear_history_and_get_html():
    clear_history()
    return get_history_html()

def download_all_as_zip():
    if not os.path.exists(HISTORY_DIR):
        return None
    import zipfile
    zip_path = os.path.join(HISTORY_DIR, "all_ocr_files.zip")
    if os.path.exists(zip_path):
        try:
            os.remove(zip_path)
        except Exception:
            pass
            
    files = [f for f in os.listdir(HISTORY_DIR) if f.lower().endswith(".pdf")]
    if not files:
        return None
        
    with zipfile.ZipFile(zip_path, "w") as zipf:
        for f in files:
            zipf.write(os.path.join(HISTORY_DIR, f), arcname=f)
    return zip_path

custom_css = """
body, .gradio-container {
    background-color: #0f172a !important;
}
.feedback {
    border-radius: 16px !important;
    border: 1px solid rgba(255, 255, 255, 0.1) !important;
    background: rgba(30, 41, 59, 0.7) !important;
    backdrop-filter: blur(12px) !important;
    box-shadow: 0 8px 32px 0 rgba(0, 0, 0, 0.37) !important;
}
button.primary {
    background: linear-gradient(90deg, #007AFF 0%, #00C6FF 100%) !important;
    border: none !important;
    color: white !important;
    border-radius: 8px !important;
    font-weight: bold !important;
    transition: all 0.3s ease !important;
}
button.primary:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 4px 20px rgba(0, 122, 255, 0.5) !important;
}
"""

# Gradio Interface
with gr.Blocks(title="OCR app API & Web Interface", theme=gr.themes.Default(primary_hue="blue", neutral_hue="slate"), css=custom_css) as demo:
    gr.Markdown("# 📄 Multi-Mode OCR API & Web UI")
    gr.Markdown("Wähle eine PDF-Datei und einen Modus, um ein durchsuchbares PDF zu generieren. Dieser Space kann auch programmgesteuert aufgerufen werden.")
    
    with gr.Row():
        with gr.Column():
            file_input = gr.File(label="PDF Datei hochladen", file_types=[".pdf"])
            mode_input = gr.Radio(
                choices=["Schnell (Gemini Full-Page)", "Präzise (Hybrid: PaddleOCR + Gemini)", "Lokal Schnell (PaddleOCR)", "Lokal Deep (PaddleOCR + TrOCR)"],
                value="Schnell (Gemini Full-Page)",
                label="OCR Modus"
            )
            smart_skip_input = gr.Checkbox(value=True, label="Bereits durchsuchbare Seiten überspringen (Smart-Skip)")
            btn = gr.Button("🚀 OCR starten", variant="primary")
        
        with gr.Column():
            file_output = gr.File(label="Durchsuchbares PDF herunterladen", interactive=False)

    # Verlauf/History Section
    with gr.Accordion("📋 Verlauf / Abgeschlossene Dateien", open=True):
        gr.Markdown("Hier siehst du alle fertig verarbeiteten OCR-Dokumente der aktuellen Sitzung. Klicke auf ein PDF, um es herunterzuladen.")
        history_files = gr.HTML(
            value=get_history_html
        )
        with gr.Row():
            zip_btn = gr.Button("📦 Als ZIP herunterladen")
            refresh_btn = gr.Button("🔄 Verlauf aktualisieren")
            clear_btn = gr.Button("🗑️ Verlauf leeren", variant="stop")
        zip_output = gr.File(label="Erstelltes ZIP-Archiv", visible=False, interactive=False)

    # Dummy compatibility button to reserve fn_index: 0 for older clients
    compat_btn = gr.Button(visible=False)
    compat_btn.click(
        fn=process_pdf,
        inputs=[file_input, mode_input],
        outputs=file_output,
        api_name="process_pdf"
    )

    btn.click(
        fn=process_pdf,
        inputs=[file_input, mode_input, smart_skip_input],
        outputs=file_output,
        api_name="process_pdf_v2"
    ).then(
        fn=get_history_html,
        inputs=[],
        outputs=history_files
    )

    zip_btn.click(
        fn=download_all_as_zip,
        inputs=[],
        outputs=zip_output
    ).then(
        fn=lambda: gr.update(visible=True),
        inputs=[],
        outputs=zip_output
    )

    refresh_btn.click(
        fn=get_history_html,
        inputs=[],
        outputs=history_files
    )

    clear_btn.click(
        fn=clear_history_and_get_html,
        inputs=[],
        outputs=history_files
    )

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860, allowed_paths=[os.path.abspath("output_files")])