app.py

import streamlit as st
import fitz  # PyMuPDF
import cv2
import numpy as np
import io
import math
from PIL import Image
import time
import os

# WICHTIG: torch MUSS vor paddleocr importiert werden unter Windows
import torch
from paddleocr import PaddleOCR

from google import genai
from google.genai import types
from pydantic import BaseModel

# --- 1. Strukturierter Output definieren (Pydantic) ---
class OCRResult(BaseModel):
    transcriptions: list[str]

class BoundingBox(BaseModel):
    box_2d: list[int]
    text: str

# Hardcoded API-Keys (cleared for security, now loaded via environment variables)
API_KEYS = []

def get_gemini_api_keys():
    import os
    import json
    
    # 1. Versuche aus config.json zu laden
    CONFIG_FILE = "config.json"
    if os.path.exists(CONFIG_FILE):
        try:
            with open(CONFIG_FILE, "r", encoding="utf-8") as f:
                config_data = json.load(f)
                config_keys = config_data.get("gemini_api_keys")
                if config_keys:
                    if isinstance(config_keys, list):
                        keys = [k.strip() for k in config_keys if isinstance(k, str) and k.strip()]
                        if keys:
                            return keys
                    elif isinstance(config_keys, str):
                        keys = [k.strip() for k in config_keys.split(",") if k.strip()]
                        if keys:
                            return keys
        except Exception:
            pass

    # 2. Versuche aus Umgebungsvariable GEMINI_API_KEYS (Komma-separiert) zu laden
    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

    # 3. Versuche nummerierte Umgebungsvariablen zu laden (GEMINI_API_KEY, GEMINI_API_KEY_2, GEMINI_API_KEY_3, 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

    # 4. Fallback auf hardcodierte Keys
    valid_keys = [k.strip() for k in API_KEYS if k.strip() and not k.startswith("ADD_YOUR_")]
    if valid_keys:
        return valid_keys
    return []

def fetch_gemini_ocr_for_page(page_num, img_bytes, api_key, prompt, mode):
    """
    Führt den Gemini-OCR Aufruf für eine einzelne Seite in einem Thread aus.
    Führt selbstständig Retries bei Rate Limits (429/503) und Fallback auf gemini-2.5-flash durch.
    """
    from google import genai
    from google.genai import types
    import time
    import re
    
    client = genai.Client(api_key=api_key)
    max_retries = 10 if mode == "Präzise (Hybrid: PaddleOCR + Gemini)" else 20
    current_model = 'gemini-3.1-flash-lite'
    
    for attempt in range(max_retries):
        try:
            response = client.models.generate_content(
                model=current_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
                )
            )
            return page_num, response.parsed, None
        except Exception as e:
            error_msg = str(e)
            is_retriable = any(code in error_msg for code in ["503", "429", "500", "502", "504", "Quota", "exhausted", "ResourceExhausted", "limit"])
            if is_retriable and attempt < max_retries - 1:
                wait_time = 5
                match = re.search(r"'retryDelay':\s*'(\d+(?:\.\d+)?)s'", error_msg)
                if match:
                    wait_time = int(float(match.group(1))) + 5
                
                if "503" in error_msg or "unavailable" in error_msg.lower():
                    current_model = 'gemini-2.5-flash'
                
                time.sleep(wait_time)
            else:
                return page_num, None, e
                
    return page_num, None, Exception("Maximale Anzahl an Retries überschritten")

# Cache die Modelle, damit sie nur bei Bedarf und nur einmal geladen werden
@st.cache_resource
def get_paddle_ocr():
    return PaddleOCR(use_angle_cls=True, lang='de')

@st.cache_resource
def get_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
        processor = TrOCRProcessor.from_pretrained(onnx_path)
        model = ORTModelForVision2Seq.from_pretrained(onnx_path, provider="DMLExecutionProvider")
        return processor, model
    else:
        from transformers import TrOCRProcessor, VisionEncoderDecoderModel
        processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-handwritten')
        model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-handwritten')
        return processor, 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:
        # Keine Lücken -> Zeilenweises sortieren (top to bottom, left to right)
        return sorted(boxes_with_data, key=lambda b: (b[0][1], b[0][0]))
        
    if max_v_gap > max_h_gap: # Spaltentrennung bevorzugen, wenn die vertikale Lücke größer ist
        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)
            
    # Fallback
    return sorted(boxes_with_data, key=lambda b: (b[0][1], b[0][0]))

def main():
    st.set_page_config(page_title="Multi-Mode Math OCR", page_icon="⚙️")
    st.title("⚙️ Multi-Mode OCR System")
    
    st.info("Wähle unten deinen bevorzugten OCR-Modus aus. Jeder Modus hat seine eigenen Stärken in Bezug auf Geschwindigkeit und Ausrichtungspräzision.")

    mode = st.radio(
        "Wähle den Verarbeitungsmodus:",
        ["Schnell (Gemini Full-Page)", "Präzise (Hybrid: PaddleOCR + Gemini)", "Lokal Schnell (PaddleOCR)", "Lokal Deep (PaddleOCR + TrOCR)"],
        index=0,
        help="Schnell: Nimmt die ganze Seite auf einmal. Präzise: Zerschneidet die Seite für perfekte schräge Ausrichtung. Lokal Schnell: Nur PaddleOCR, extrem schnell und offline. Lokal Deep: Mit TrOCR für handschriftliche Texte."
    )

    smart_skip = st.checkbox(
        "Bereits durchsuchbare Seiten überspringen (Smart-Skip)",
        value=True,
        help="Seiten, die bereits editierbaren Text enthalten, werden übersprungen. Falls deaktiviert, wird OCR erzwungen und der alte Text vorher entfernt."
    )

    # Datei-Upload
    uploaded_file = st.file_uploader("Ziehe dein PDF hierhin oder klicke zum Auswählen", type=["pdf"])

    if uploaded_file is not None:
        if st.button("🚀 OCR Starten"):
            status_text = st.empty()
            progress_bar = st.progress(0)

            try:
                # 1. API Clients und KI Modelle bedarfsgerecht laden
                if "Gemini" in mode:
                    api_keys = get_gemini_api_keys()
                
                if "PaddleOCR" in mode:
                    status_text.text("Lade PaddleOCR Modell (Geometrie-KI)...")
                    paddle_ocr = get_paddle_ocr()
                    
                if "Lokal Deep" in mode:
                    status_text.text("Lade TrOCR Modell (Microsoft Deep Handwriting)... Dies kann einen Moment dauern.")
                    trocr_processor, trocr_model = get_trocr()

                # PDF laden
                pdf_bytes = uploaded_file.read()
                doc = fitz.open(stream=pdf_bytes, filetype="pdf")
                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"[App] Seite {i+1} hat bereits Text und Smart-Skip ist aktiv. Überspringe OCR.")
                    else:
                        if has_text and not smart_skip:
                            print(f"[App] 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)

                # ==========================================
                # MODUS 1: Schnell (Gemini Full-Page) - Parallel
                # ==========================================
                if mode == "Schnell (Gemini Full-Page)":
                    pages_img_bytes = {}
                    status_text.text("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")

                    import concurrent.futures
                    results = {}
                    errors = {}
                    num_keys = len(api_keys)
                    
                    status_text.text(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:
                        with concurrent.futures.ThreadPoolExecutor(max_workers=num_keys) as executor:
                            futures = []
                            for idx, i in enumerate(pages_to_ocr):
                                key = api_keys[idx % num_keys]
                                futures.append(executor.submit(fetch_gemini_ocr_for_page, i, pages_img_bytes[i], key, 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
                                status_text.text(f"Gemini Semantic Analyse: {completed} von {len(pages_to_ocr)} Seiten abgeschlossen...")
                                progress_bar.progress(completed / len(pages_to_ocr))

                    status_text.text("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"[App] 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:
                            st.warning(f"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:
                                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

                # ==========================================
                # 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):
                        status_text.text(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
                        progress_bar.progress((idx + 1) / (len(pages_to_ocr) * 2) if pages_to_ocr else 0.5)

                    import concurrent.futures
                    results = {}
                    errors = {}
                    num_keys = len(api_keys)
                    
                    status_text.text(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:
                        with concurrent.futures.ThreadPoolExecutor(max_workers=num_keys) as executor:
                            futures = []
                            for idx, i in enumerate(pages_to_ocr):
                                key = api_keys[idx % num_keys]
                                futures.append(executor.submit(fetch_gemini_ocr_for_page, i, pages_img_bytes[i], key, 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
                                status_text.text(f"Gemini Semantic Analyse: {completed} von {len(pages_to_ocr)} Seiten abgeschlossen...")
                                progress_bar.progress(0.5 + (completed / len(pages_to_ocr)) * 0.5)

                    status_text.text("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"[App] 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:
                            st.warning(f"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 = []
                            for g_box in parsed_boxes:
                                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))
                                
                            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:
                                            c_y_center = cluster['y_center']
                                            if abs(y_center - c_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
                                            
                                            if len(cluster_boxes) == 1:
                                                pb, pt = cluster_boxes[0]
                                                p0 = [pb[0][0]/zoom, pb[0][1]/zoom]
                                                p1 = [pb[1][0]/zoom, pb[1][1]/zoom]
                                                p2 = [pb[2][0]/zoom, pb[2][1]/zoom]
                                                p3 = [pb[3][0]/zoom, pb[3][1]/zoom]
                                                
                                                dx = p1[0] - p0[0]
                                                dy = p1[1] - p0[1]
                                                dx_up = p0[0] - p3[0]
                                                dy_up = p0[1] - p3[1]
                                                angle_rad = math.atan2(dy, dx) if (dx != 0 or dy != 0) else 0
                                                angle_deg = math.degrees(angle_rad)
                                            else:
                                                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 = math.cos(-angle_rad)
                                                sin_a = math.sin(-angle_rad)
                                                
                                                local_points = []
                                                for px, py in all_points:
                                                    local_points.append((px * cos_a - py * sin_a, px * sin_a + py * cos_a))
                                                    
                                                min_lx = min(p[0] for p in local_points)
                                                max_lx = max(p[0] for p in local_points)
                                                min_ly = min(p[1] for p in local_points)
                                                max_ly = max(p[1] for p in local_points)
                                                
                                                lp0, lp1, lp2, lp3 = (min_lx, min_ly), (max_lx, min_ly), (max_lx, max_ly), (min_lx, max_ly)
                                                
                                                cos_inv = math.cos(angle_rad)
                                                sin_inv = math.sin(angle_rad)
                                                
                                                merged_box = []
                                                for lx, ly in [lp0, lp1, lp2, lp3]:
                                                    merged_box.append([lx * cos_inv - ly * sin_inv, lx * sin_inv + ly * cos_inv])
                                                    
                                                p0, p1, p2, p3 = merged_box
                                                dx = p1[0] - p0[0]
                                                dy = p1[1] - p0[1]
                                                dx_up = p0[0] - p3[0]
                                                dy_up = p0[1] - p3[1]
                                                
                                            box_width_pdf = math.hypot(dx, dy)
                                            box_height_pdf = math.hypot(dx_up, dy_up)
                                            
                                            font = fitz.Font("helv")
                                            shift_factor = -font.descender
                                            pdf_baseline = fitz.Point(p3[0] + dx_up * shift_factor, p3[1] + dy_up * shift_factor)
                                            
                                            merged_points = [p0, p1, p2, p3]
                                            coords = [min(p[0] for p in merged_points), min(p[1] for p in merged_points), max(p[0] for p in merged_points), max(p[1] for p in merged_points)]
                                            boxes_with_data.append((coords, (line_text, pdf_baseline, box_width_pdf, box_height_pdf, 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 = math.cos(-angle_rad)
                                        sin_a = math.sin(-angle_rad)
                                        
                                        local_points = []
                                        for px, py in all_points:
                                            local_points.append((px * cos_a - py * sin_a, px * sin_a + py * cos_a))
                                            
                                        min_lx = min(p[0] for p in local_points)
                                        max_lx = max(p[0] for p in local_points)
                                        min_ly = min(p[1] for p in local_points)
                                        max_ly = max(p[1] for p in local_points)
                                        
                                        lp0, lp1, lp2, lp3 = (min_lx, min_ly), (max_lx, min_ly), (max_lx, max_ly), (min_lx, max_ly)
                                        
                                        cos_inv = math.cos(angle_rad)
                                        sin_inv = math.sin(angle_rad)
                                        
                                        merged_box = []
                                        for lx, ly in [lp0, lp1, lp2, lp3]:
                                            merged_box.append([lx * cos_inv - ly * sin_inv, lx * sin_inv + ly * cos_inv])
                                            
                                        p0, p1, p2, p3 = merged_box
                                        dx = p1[0] - p0[0]
                                        dy = p1[1] - p0[1]
                                        dx_up = p0[0] - p3[0]
                                        dy_up = p0[1] - p3[1]
                                        
                                        box_width_pdf = math.hypot(dx, dy)
                                        box_height_pdf = math.hypot(dx_up, dy_up)
                                        
                                        font = fitz.Font("helv")
                                        shift_factor = -font.descender
                                        pdf_baseline = fitz.Point(p3[0] + dx_up * shift_factor, p3[1] + dy_up * shift_factor)
                                        
                                        coords = [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)]
                                        flat_text = g_text.replace('\n', '  ')
                                        boxes_with_data.append((coords, (flat_text, pdf_baseline, box_width_pdf, box_height_pdf, angle_deg)))
                                else:
                                    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
                                    pdf_baseline = fitz.Point(x0, y1 - (y1-y0)*0.2)
                                    boxes_with_data.append(([x0, y0, x1, y1], (g_text.replace('\n', '  '), pdf_baseline, 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)
                                scale_x = box_width_pdf / text_length if text_length > 0 else 1.0
                                
                                matrix = fitz.Matrix(scale_x, 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 ValueError:
                                    clean_text = text.encode("latin-1", "ignore").decode("latin-1")
                                    if clean_text.strip():
                                        try:
                                            page.insert_text(pdf_baseline, clean_text, fontsize=fontsize, fontname="helv", render_mode=3, morph=(pdf_baseline, matrix))
                                        except Exception:
                                            pass

                # ==========================================
                # MODUS 3: Lokal Deep (TrOCR) - Sequentiell
                # ==========================================
                elif mode == "Lokal Deep (PaddleOCR + TrOCR)":
                    for idx, page_num in enumerate(pages_to_ocr):
                        status_text.text(f"Verarbeite Seite {page_num + 1} von {num_pages} (Modus: {mode})...")
                        page = doc.load_page(page_num)
                        
                        zoom = 3.0
                        mat = fitz.Matrix(zoom, zoom)
                        pix = page.get_pixmap(matrix=mat)
                        
                        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)
                        if page_num in pages_to_redact:
                            print(f"[App] 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)

                        status_text.text(f"Seite {page_num + 1}: PaddleOCR Layout Analyse...")
                        result = paddle_ocr.ocr(img_np)
                        
                        if not result or not result[0]:
                            continue
                            
                        page_data = result[0]
                        valid_lines = [l for l in page_data if l]
                        crops = []
                        valid_boxes = []
                        
                        for line in valid_lines:
                            box = line[0]
                            x_coords = [int(p[0]) for p in box]
                            y_coords = [int(p[1]) for p in box]
                            x_min, x_max = max(0, min(x_coords) - 2), min(img_np.shape[1], max(x_coords) + 2)
                            y_min, y_max = max(0, min(y_coords) - 2), min(img_np.shape[0], max(y_coords) + 2)
                            
                            crop_img = img_np[y_min:y_max, x_min:x_max]
                            if crop_img.size > 0:
                                crop_rgb = cv2.cvtColor(crop_img, cv2.COLOR_BGR2RGB)
                                crops.append(Image.fromarray(crop_rgb))
                                valid_boxes.append((box, line[1][0]))
                        
                        BATCH_SIZE = 4 
                        results = []
                        
                        status_text.text(f"Seite {page_num + 1}: TrOCR liest {len(crops)} Text-Schnipsel...")
                        
                        for b_idx in range(0, len(crops), BATCH_SIZE):
                            batch_crops = crops[b_idx:b_idx+BATCH_SIZE]
                            batch_fallbacks = [vb[1] for vb in valid_boxes[b_idx:b_idx+BATCH_SIZE]]
                            
                            batch_texts = list(batch_fallbacks)
                            
                            if batch_crops:
                                try:
                                    pixel_values = trocr_processor(batch_crops, return_tensors="pt").pixel_values
                                    generated_ids = trocr_model.generate(pixel_values, max_new_tokens=30)
                                    texts = trocr_processor.batch_decode(generated_ids, skip_special_tokens=True)
                                    for map_idx, txt in enumerate(texts):
                                        batch_texts[map_idx] = txt
                                except Exception as e:
                                    print(f"TrOCR batch failed, using fallback. Error: {e}")
                                    
                            results.extend(batch_texts)

                        boxes_with_data = []
                        for i_res, text in enumerate(results):
                            if not text.strip(): continue
                            box = valid_boxes[i_res][0]
                            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, p2, p3 = box
                            dx = p1[0] - p0[0]
                            dy = p1[1] - p0[1]
                            angle_deg = math.degrees(math.atan2(dy, dx))
                            
                            dx_up = p0[0] - p3[0]
                            dy_up = p0[1] - p3[1]
                            
                            font = fitz.Font("helv")
                            shift_factor = -font.descender
                            
                            base_x = p3[0] + dx_up * shift_factor
                            base_y = p3[1] + dy_up * shift_factor
                            
                            pdf_baseline = fitz.Point(base_x / zoom, base_y / zoom)
                            box_width_pdf = math.hypot(dx, dy) / zoom
                            box_height_pdf = math.hypot(dx_up, dy_up) / zoom
                            
                            fontsize = box_height_pdf
                            text_length = fitz.get_text_length(text, fontname="helv", fontsize=fontsize)
                            scale_x = box_width_pdf / text_length if text_length > 0 else 1.0
                            
                            matrix = fitz.Matrix(scale_x, 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 ValueError:
                                clean_text = text.encode("latin-1", "ignore").decode("latin-1")
                                if clean_text.strip():
                                    try:
                                        page.insert_text(pdf_baseline, clean_text, fontsize=fontsize, fontname="helv", render_mode=3, morph=(pdf_baseline, matrix))
                                    except Exception:
                                        pass
                        progress_bar.progress((idx + 1) / len(pages_to_ocr) if pages_to_ocr else 1.0)

                # ==========================================
                # MODUS 4: Lokal Schnell (PaddleOCR) - Sequentiell
                # ==========================================
                elif mode == "Lokal Schnell (PaddleOCR)":
                    for idx, page_num in enumerate(pages_to_ocr):
                        status_text.text(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"[App] 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
                        progress_bar.progress((idx + 1) / len(pages_to_ocr) if pages_to_ocr else 1.0)

                status_text.text("🎉 Verarbeitung komplett! PDF wird generiert...")
                
                out_bytes = doc.tobytes()
                
                # 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_{uploaded_file.name}"
                    if not dest_filename.lower().endswith(".pdf"):
                        dest_filename += ".pdf"
                    dest_path = os.path.join(history_dir, dest_filename)
                    
                    with open(dest_path, "wb") as f_out:
                        f_out.write(out_bytes)
                    print(f"[App] Saved completed file to history (overwritten if existed): {dest_path}")
                except Exception as ex:
                    print(f"[App] Warning: Failed to save completed file to history: {ex}")


                st.success(f"Fertig! Dein PDF ({mode}) steht zum Download bereit.")
                st.download_button(
                    label="📥 Fertiges PDF herunterladen",
                    data=out_bytes,
                    file_name=f"searchable_{uploaded_file.name}",
                    mime="application/pdf"
                )

            except Exception as e:
                st.error(f"❌ Es ist ein Fehler aufgetreten: {e}")

    # Verlauf-Sektion in Streamlit
    st.write("---")
    st.header("📋 Verlauf / Abgeschlossene Dateien")
    
    history_dir = "output_files"
    os.makedirs(history_dir, exist_ok=True)
    
    # List files
    import glob
    files = glob.glob(os.path.join(history_dir, "*.pdf")) + glob.glob(os.path.join(history_dir, "*.zip"))
    files.sort(key=os.path.getmtime, reverse=True)
    
    # Remove ZIP from the listing itself to prevent recursion
    files = [f for f in files if os.path.basename(f) != "all_ocr_files.zip"]
    
    if not files:
        st.info("Keine verarbeiteten Dateien vorhanden.")
    else:
        selected_files = []
        
        # Display list with checkboxes in a neat layout
        for path in files:
            filename = os.path.basename(path)
            
            col1, col2, col3 = st.columns([0.1, 0.7, 0.2])
            with col1:
                checked = st.checkbox("", key=f"select_{filename}")
                if checked:
                    selected_files.append(path)
            with col2:
                # Format file size
                try:
                    size_bytes = os.path.getsize(path)
                    size_mb = size_bytes / (1024 * 1024)
                    size_str = f"({size_mb:.2f} MB)"
                except Exception:
                    size_str = ""
                st.write(f"📄 {filename} {size_str}")
            with col3:
                # Read file bytes for download
                try:
                    with open(path, "rb") as f:
                        file_data = f.read()
                    st.download_button(
                        label="📥",
                        data=file_data,
                        file_name=filename,
                        mime="application/pdf" if filename.endswith(".pdf") else "application/zip",
                        key=f"dl_{filename}"
                    )
                except Exception:
                    pass
                    
        # Bulk Actions
        if selected_files:
            col_bulk1, col_bulk2 = st.columns([0.5, 0.5])
            with col_bulk1:
                # Create ZIP
                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
                try:
                    with zipfile.ZipFile(zip_path, "w") as zipf:
                        for f in selected_files:
                            zipf.write(f, arcname=os.path.basename(f))
                    with open(zip_path, "rb") as f:
                        zip_data = f.read()
                    st.download_button(
                        label="📦 Ausgewählte als ZIP herunterladen",
                        data=zip_data,
                        file_name="ocr_archive.zip",
                        mime="application/zip",
                        use_container_width=True
                    )
                except Exception as e:
                    st.error(f"Fehler beim Erstellen des ZIPs: {e}")
            with col_bulk2:
                if st.button("🗑️ Ausgewählte löschen", use_container_width=True):
                    for path in selected_files:
                        try:
                            os.remove(path)
                        except Exception:
                            pass
                    try:
                        st.rerun()
                    except AttributeError:
                        st.experimental_rerun()

if __name__ == "__main__":
    main()