main.py

from dotenv import load_dotenv
load_dotenv()

from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from transformers import pipeline as hf_pipeline, AutoTokenizer, AutoModelForTokenClassification
from doctr.io import DocumentFile
from doctr.models import ocr_predictor
from img2table.document import Image as Img2TableImage
from img2table.ocr import DocTR
import cv2
import numpy as np
from PIL import Image
import io
import json
import os
import tempfile
import base64
from typing import Dict, Any, Optional, List
import difflib
import re
import httpx
from bs4 import BeautifulSoup

# Docling pipeline
from docling.document_converter import DocumentConverter, InputFormat, ImageFormatOption
from docling.datamodel.pipeline_options import PdfPipelineOptions
from docling_ocr_onnxtr import OnnxtrOcrOptions

# Chat router
from router_chat import router as chat_router
from faq_store import initialize_faq_store

app = FastAPI(title="ScanAssured OCR & NER API")

@app.on_event("startup")
async def startup_event():
    initialize_faq_store()

app.include_router(chat_router)

# --- DRUG INTERACTIONS DATABASE ---
DRUG_INTERACTIONS = {}
interactions_path = os.path.join(os.path.dirname(__file__), 'interactions_data.json')
if os.path.exists(interactions_path):
    with open(interactions_path, 'r') as f:
        DRUG_INTERACTIONS = json.load(f)
    print(f"Loaded {len(DRUG_INTERACTIONS)} drug interaction entries")

# --- LAB REFERENCE: MEDLINEPLUS MAPPINGS ---
MEDLINEPLUS_MAP = {}
medlineplus_map_path = os.path.join(os.path.dirname(__file__), 'medlineplus_map.json')
if os.path.exists(medlineplus_map_path):
    with open(medlineplus_map_path, 'r') as f:
        MEDLINEPLUS_MAP = json.load(f)
    print(f"Loaded {len(MEDLINEPLUS_MAP)} MedlinePlus test mappings")

MEDLINEPLUS_CACHE = {}
medlineplus_cache_path = os.path.join(os.path.dirname(__file__), 'medlineplus_cache.json')
if os.path.exists(medlineplus_cache_path):
    with open(medlineplus_cache_path, 'r') as f:
        MEDLINEPLUS_CACHE = json.load(f)
    print(f"Loaded {len(MEDLINEPLUS_CACHE)} MedlinePlus cached entries")

# Enable CORS for Flutter app
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- OCR PRESETS ---
OCR_PRESETS = {
    "high_accuracy": {
        "det": "db_resnet50",
        "reco": "crnn_vgg16_bn",
        "name": "High Accuracy",
        "description": "Best quality, slower processing"
    },
    "balanced": {
        "det": "db_resnet50",
        "reco": "crnn_mobilenet_v3_small",
        "name": "Balanced (Recommended)",
        "description": "Good quality and speed"
    },
    "fast": {
        "det": "db_mobilenet_v3_large",
        "reco": "crnn_mobilenet_v3_small",
        "name": "Fast",
        "description": "Fastest processing, slightly lower quality"
    },
}

OCR_DETECTION_MODELS = ["db_resnet50", "db_mobilenet_v3_large", "linknet_resnet18"]
OCR_RECOGNITION_MODELS = ["crnn_vgg16_bn", "crnn_mobilenet_v3_small", "parseq"]

# --- NER MODELS ---
NER_MODELS = {
    "Clinical-AI-Apollo/Medical-NER": {
        "name": "Medical NER (Recommended)",
        "description": "Medications, diseases, lab values, procedures, dosages",
        "entities": ["MEDICATION", "DOSAGE", "FREQUENCY", "DURATION",
                     "DISEASE_DISORDER", "SIGN_SYMPTOM", "DIAGNOSTIC_PROCEDURE",
                     "THERAPEUTIC_PROCEDURE", "LAB_VALUE", "SEVERITY"]
    },
    "samrawal/bert-base-uncased_clinical-ner": {
        "name": "Clinical Notes",
        "description": "Optimized for clinical/medical notes",
        "entities": ["PROBLEM", "TREATMENT", "TEST"]
    },
}

# --- GLOBAL MODEL CACHES ---
ner_model_cache: Dict[str, Any] = {}
ocr_model_cache: Dict[str, Any] = {}
# --- DOCLING CONVERTER CACHE ---
docling_converter_cache: Dict[str, Any] = {}

def get_docling_converter(det_arch: str = "db_mobilenet_v3_large", reco_arch: str = "crnn_vgg16_bn"):
    """Get or create a cached Docling DocumentConverter with OnnxTR OCR."""
    cache_key = f"docling_{det_arch}_{reco_arch}"

    if cache_key in docling_converter_cache:
        print(f"Using cached Docling converter: {cache_key}")
        return docling_converter_cache[cache_key]

    try:
        print(f"Initializing Docling converter: det={det_arch}, reco={reco_arch}...")

        ocr_options = OnnxtrOcrOptions(
            det_arch=det_arch,
            reco_arch=reco_arch,
        )

        pipeline_options = PdfPipelineOptions(ocr_options=ocr_options)
        pipeline_options.do_table_structure = True
        pipeline_options.do_ocr = True
        pipeline_options.allow_external_plugins = True

        converter = DocumentConverter(
            format_options={
                InputFormat.IMAGE: ImageFormatOption(pipeline_options=pipeline_options)
            }
        )

        docling_converter_cache[cache_key] = converter
        print(f"Docling converter {cache_key} initialized successfully!")
        return converter
    except Exception as e:
        print(f"ERROR: Failed to initialize Docling converter: {e}")
        import traceback
        traceback.print_exc()
        return None


def run_docling_pipeline(file_content: bytes) -> Dict[str, Any]:
    """
    Run the Docling pipeline on raw image bytes.
    Returns structured results for comparison with docTR.
    """
    try:
        converter = get_docling_converter()
        if converter is None:
            return {"error": "Docling converter not available", "success": False}

        # Docling needs a file path - write to temp file
        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
            tmp_file.write(file_content)
            tmp_path = tmp_file.name

        try:
            print("Running Docling pipeline...")
            result = converter.convert(source=tmp_path)

            # Extract markdown (preserves headings, tables, paragraphs)
            markdown_text = result.document.export_to_markdown()

            # Extract plain text
            if hasattr(result.document, 'export_to_text'):
                plain_text = result.document.export_to_text()
            else:
                plain_text = markdown_text

            # Extract tables
            docling_tables = []
            if hasattr(result.document, 'tables') and result.document.tables:
                for table in result.document.tables:
                    table_data = _parse_docling_table(table)
                    if table_data:
                        docling_tables.append(table_data)

            print(f"Docling: {len(markdown_text)} chars markdown, {len(docling_tables)} tables")

            return {
                "success": True,
                "markdown_text": markdown_text,
                "plain_text": plain_text,
                "tables": docling_tables,
                "primary_table": docling_tables[0] if docling_tables else None,
            }
        finally:
            try:
                os.unlink(tmp_path)
            except:
                pass

    except Exception as e:
        print(f"Docling pipeline error: {e}")
        import traceback
        traceback.print_exc()
        return {"error": str(e), "success": False}


def _parse_docling_table(table) -> Optional[Dict]:
    """Parse a Docling table into {cells, num_rows, num_columns} format."""
    try:
        if hasattr(table, 'export_to_dataframe'):
            df = table.export_to_dataframe()
            if df is not None and not df.empty:
                cells = []
                header = [str(col) if col is not None else '' for col in df.columns.tolist()]
                cells.append(header)
                for _, row in df.iterrows():
                    row_cells = [str(val).strip() if val is not None else '' for val in row.tolist()]
                    cells.append(row_cells)

                return {
                    "cells": cells,
                    "num_rows": len(cells),
                    "num_columns": len(header),
                    "method": "docling_tableformer"
                }

        if hasattr(table, 'export_to_markdown'):
            md = table.export_to_markdown()
            if md:
                return {
                    "cells": [],
                    "num_rows": 0,
                    "num_columns": 0,
                    "method": "docling_tableformer",
                    "markdown": md
                }

        return None
    except Exception as e:
        print(f"Docling table parse error: {e}")
        return None


# --- OCR MODEL LOADING ---
def get_ocr_predictor(det_arch: str, reco_arch: str):
    """Retrieves a loaded OCR predictor from cache or loads it if necessary."""
    cache_key = f"{det_arch}_{reco_arch}"

    if cache_key in ocr_model_cache:
        print(f"Using cached OCR model: {cache_key}")
        return ocr_model_cache[cache_key]

    try:
        print(f"Loading OCR model: det={det_arch}, reco={reco_arch}...")
        predictor = ocr_predictor(
            det_arch=det_arch,
            reco_arch=reco_arch,
            pretrained=True,
            assume_straight_pages=True,   # Assume pages are already straight
            straighten_pages=False,       # Don't auto-rotate (was causing issues)
            detect_orientation=False,     # Don't detect orientation (was inverting text)
            preserve_aspect_ratio=True    # Keep proportions
        )
        ocr_model_cache[cache_key] = predictor
        print(f"OCR model {cache_key} loaded successfully!")
        return predictor
    except Exception as e:
        print(f"ERROR: Failed to load OCR model {cache_key}: {e}")
        return None

# --- NER MODEL LOADING ---
def get_ner_pipeline(model_id: str):
    """Retrieves a loaded NER pipeline from cache or loads it if necessary."""
    if model_id not in NER_MODELS:
        raise ValueError(f"Unknown NER model ID: {model_id}")

    if model_id in ner_model_cache:
        print(f"Using cached NER model: {model_id}")
        return ner_model_cache[model_id]

    try:
        print(f"Loading NER model: {model_id}...")
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        model = AutoModelForTokenClassification.from_pretrained(model_id)

        ner_pipeline = hf_pipeline(
            "ner",
            model=model,
            tokenizer=tokenizer,
            aggregation_strategy="simple"
        )
        ner_model_cache[model_id] = ner_pipeline
        print(f"NER model {model_id} loaded successfully!")
        return ner_pipeline
    except Exception as e:
        print(f"ERROR: Failed to load NER model {model_id}: {e}")
        return None


def _edit_distance(s1: str, s2: str) -> int:
    """Compute Levenshtein edit distance between two strings."""
    if len(s1) < len(s2):
        return _edit_distance(s2, s1)
    if len(s2) == 0:
        return len(s1)

    prev_row = range(len(s2) + 1)
    for i, c1 in enumerate(s1):
        curr_row = [i + 1]
        for j, c2 in enumerate(s2):
            insertions = prev_row[j + 1] + 1
            deletions = curr_row[j] + 1
            substitutions = prev_row[j] + (c1 != c2)
            curr_row.append(min(insertions, deletions, substitutions))
        prev_row = curr_row
    return prev_row[-1]


# --- NER-INFORMED CORRECTION ---

_entity_dicts: dict[str, set] = {}


def _build_entity_dicts():
    """Build per-entity-type dictionaries from already-loaded DRUG_INTERACTIONS and MEDLINEPLUS_MAP."""
    global _entity_dicts

    med_dict: set[str] = set()
    for drug_name in DRUG_INTERACTIONS.keys():
        for part in str(drug_name).split(','):
            part = part.strip().lower()
            if len(part) >= 4:
                med_dict.add(part)

    lab_dict: set[str] = set()
    for test_name, data in MEDLINEPLUS_MAP.items():
        if len(test_name) >= 4:
            lab_dict.add(test_name.lower())
        for alias in data.get('aliases', []):
            if len(alias) >= 4:
                lab_dict.add(alias.lower())

    _entity_dicts = {
        'MEDICATION':            med_dict,
        'LAB_VALUE':             lab_dict,
        'DIAGNOSTIC_PROCEDURE':  lab_dict,
        'TREATMENT':             med_dict,
        'CHEM':                  med_dict,
        'CHEMICAL':              med_dict,
    }
    print(f"Entity dicts built: {len(med_dict)} medication terms, {len(lab_dict)} lab terms")


def _find_closest(word: str, dictionary: set) -> tuple:
    best_match, best_dist = None, 999
    word_lower = word.lower()
    for term in dictionary:
        if abs(len(term) - len(word_lower)) > 3:
            continue
        dist = _edit_distance(word_lower, term)
        if dist < best_dist:
            best_dist = dist
            best_match = term
    return best_match, best_dist


def _match_case(original: str, replacement: str) -> str:
    if original.isupper():
        return replacement.upper()
    if original[0].isupper():
        return replacement.capitalize()
    return replacement.lower()


def correct_with_ner_entities(
    words_with_boxes: list,
    ner_entities: list,
    text: str,
    confidence_threshold: float = 0.75,
) -> dict:
    """Second-pass correction using NER entity labels as context."""
    if not _entity_dicts:
        _build_entity_dicts()

    word_conf: dict[str, float] = {}
    for w in words_with_boxes:
        key = w['word'].lower()
        word_conf[key] = min(word_conf.get(key, 1.0), w.get('confidence', 1.0))

    corrections = []
    corrected_text = text

    for entity in ner_entities:
        entity_type = entity.get('entity_group', '')
        entity_word = entity.get('word', '').strip()
        lookup_dict = _entity_dicts.get(entity_type)
        if not lookup_dict or not entity_word:
            continue

        for token in entity_word.split():
            clean_token = re.sub(r'[^a-zA-Z]', '', token)
            if not clean_token.isalpha() or len(clean_token) < 4:
                continue

            ocr_conf = word_conf.get(clean_token.lower(), 1.0)
            if ocr_conf >= confidence_threshold:
                continue

            best_match, best_dist = _find_closest(clean_token, lookup_dict)
            if best_match is None or best_dist > 2:
                continue
            if best_match.lower() == clean_token.lower():
                continue

            replacement = _match_case(clean_token, best_match)
            match = re.search(r'\b' + re.escape(clean_token) + r'\b',
                              corrected_text, re.IGNORECASE)
            if not match:
                continue

            start, end = match.start(), match.end()
            corrected_text = corrected_text[:start] + replacement + corrected_text[end:]
            corrections.append({
                'original':       clean_token,
                'corrected':      replacement,
                'confidence':     round(1.0 - best_dist / max(len(clean_token), len(best_match)), 4),
                'ocr_confidence': round(ocr_conf, 4),
                'edit_distance':  best_dist,
                'source':         'ner',
                'entity_type':    entity_type,
            })
            word_conf[replacement.lower()] = 1.0

    return {'corrected_text': corrected_text, 'corrections': corrections}


# --- IMAGE PREPROCESSING ---
def deskew_image(image: np.ndarray) -> np.ndarray:
    """Deskew image using projection profile method."""
    try:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) == 3 else image
        edges = cv2.Canny(gray, 50, 150, apertureSize=3)
        lines = cv2.HoughLinesP(edges, 1, np.pi/180, threshold=100, minLineLength=100, maxLineGap=10)

        if lines is not None and len(lines) > 0:
            angles = []
            for line in lines:
                x1, y1, x2, y2 = line[0]
                angle = np.arctan2(y2 - y1, x2 - x1) * 180 / np.pi
                if abs(angle) < 45:
                    angles.append(angle)

            if angles:
                median_angle = np.median(angles)
                if abs(median_angle) > 0.5:
                    (h, w) = image.shape[:2]
                    center = (w // 2, h // 2)
                    M = cv2.getRotationMatrix2D(center, median_angle, 1.0)
                    rotated = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
                    return rotated

        return image
    except Exception as e:
        print(f"Deskew warning: {e}")
        return image

def preprocess_for_doctr(file_content: bytes) -> np.ndarray:
    """Automatic preprocessing pipeline optimized for docTR."""
    nparr = np.frombuffer(file_content, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

    if img is None:
        raise ValueError("Failed to decode image")

    img = deskew_image(img)

    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    lab[:, :, 0] = clahe.apply(lab[:, :, 0])
    img = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)

    img = cv2.fastNlMeansDenoisingColored(img, None, 6, 6, 7, 21)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    return img

def basic_cleanup(text: str) -> str:
    """Clean up OCR text for NER processing."""
    text = " ".join(text.split())
    return text


# --- TABLE DETECTION WITH IMG2TABLE ---

# Cache for img2table OCR instance
img2table_ocr_cache = {}

def get_img2table_ocr():
    """Get or create img2table DocTR OCR instance."""
    if 'doctr' not in img2table_ocr_cache:
        img2table_ocr_cache['doctr'] = DocTR()
    return img2table_ocr_cache['doctr']


def extract_tables_with_img2table(image_bytes: bytes, img_width: int, img_height: int) -> dict:
    """
    Use img2table to detect and extract table structure from image.
    Returns table data with properly structured cells.
    """
    try:
        # Save image to temp file (img2table needs file path)
        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
            tmp_file.write(image_bytes)
            tmp_path = tmp_file.name

        # Create img2table Image object
        img2table_img = Img2TableImage(src=tmp_path)

        # Get OCR instance
        ocr = get_img2table_ocr()

        # Extract tables with OCR
        tables = img2table_img.extract_tables(
            ocr=ocr,
            implicit_rows=True,  # Detect rows even without horizontal lines
            implicit_columns=True,  # Detect columns even without vertical lines
            borderless_tables=True,  # Detect tables without borders
            min_confidence=50  # Minimum OCR confidence
        )

        # Clean up temp file
        try:
            os.unlink(tmp_path)
        except:
            pass

        if not tables:
            return {'is_table': False, 'tables': []}

        # Process all detected tables
        all_tables = []
        for table in tables:
            cells = []

            # Method 1: Try using DataFrame (most reliable)
            if hasattr(table, 'df') and table.df is not None:
                df = table.df
                # Convert DataFrame to list of lists
                # Include header row
                header = [str(col) if col is not None else '' for col in df.columns.tolist()]
                cells.append(header)
                # Add data rows
                for _, row in df.iterrows():
                    row_cells = [str(val).strip() if val is not None else '' for val in row.tolist()]
                    cells.append(row_cells)

            # Method 2: Try content attribute with different structures
            elif hasattr(table, 'content') and table.content is not None:
                content = table.content
                if isinstance(content, list):
                    for row in content:
                        if isinstance(row, (list, tuple)):
                            row_cells = []
                            for cell in row:
                                if cell is None:
                                    row_cells.append('')
                                elif isinstance(cell, str):
                                    row_cells.append(cell.strip())
                                elif hasattr(cell, 'value'):
                                    row_cells.append(str(cell.value).strip() if cell.value else '')
                                elif hasattr(cell, 'text'):
                                    row_cells.append(str(cell.text).strip() if cell.text else '')
                                else:
                                    row_cells.append(str(cell).strip())
                            cells.append(row_cells)
                        elif isinstance(row, dict):
                            # Row might be a dict with cell data
                            row_cells = [str(v).strip() if v else '' for v in row.values()]
                            cells.append(row_cells)

            # Method 3: Try extracting from _content or other attributes
            elif hasattr(table, '_content'):
                print(f"Table has _content: {type(table._content)}")

            # Filter out empty tables
            if cells and len(cells) > 1:
                # Check if table has actual content (not just empty cells)
                has_content = any(any(c.strip() for c in row) for row in cells)
                if has_content:
                    num_cols = max(len(row) for row in cells) if cells else 0
                    all_tables.append({
                        'cells': cells,
                        'num_rows': len(cells),
                        'num_columns': num_cols
                    })
                    print(f"Extracted table with {len(cells)} rows and {num_cols} columns")

        if not all_tables:
            print("No valid tables extracted")
            return {'is_table': False, 'tables': []}

        # Return the largest table (most cells) as primary
        primary_table = max(all_tables, key=lambda t: t['num_rows'] * t['num_columns'])
        print(f"Primary table: {primary_table['num_rows']}x{primary_table['num_columns']}")

        return {
            'is_table': True,
            'cells': primary_table['cells'],
            'num_rows': primary_table['num_rows'],
            'num_columns': primary_table['num_columns'],
            'tables': all_tables,
            'total_tables': len(all_tables)
        }

    except Exception as e:
        print(f"img2table extraction error: {e}")
        import traceback
        traceback.print_exc()
        return {'is_table': False, 'error': str(e)}


def format_table_as_markdown(table_data: dict) -> str:
    """Format extracted table data as a markdown table."""
    if not table_data.get('is_table') or not table_data.get('cells'):
        return ''

    cells = table_data['cells']
    if not cells:
        return ''

    num_cols = max(len(row) for row in cells) if cells else 0
    if num_cols == 0:
        return ''

    lines = []
    col_widths = [3] * num_cols

    # Normalize rows to have same number of columns
    normalized_cells = []
    for row in cells:
        normalized_row = list(row) + [''] * (num_cols - len(row))
        normalized_cells.append(normalized_row)
        for i, cell in enumerate(normalized_row):
            if i < num_cols:
                col_widths[i] = max(col_widths[i], len(str(cell)))

    for row_idx, row in enumerate(normalized_cells):
        formatted_cells = []
        for i, cell in enumerate(row):
            if i < num_cols:
                formatted_cells.append(str(cell).ljust(col_widths[i]))

        line = '| ' + ' | '.join(formatted_cells) + ' |'
        lines.append(line)

        if row_idx == 0:
            separator = '|' + '|'.join(['-' * (w + 2) for w in col_widths]) + '|'
            lines.append(separator)

    return '\n'.join(lines)


def extract_text_with_table_detection(image_bytes: bytes, img_width: int, img_height: int) -> tuple:
    """
    Extract tables from image using img2table.
    Returns (markdown_text, table_data).
    """
    table_data = extract_tables_with_img2table(image_bytes, img_width, img_height)

    if table_data.get('is_table'):
        markdown_table = format_table_as_markdown(table_data)
        return markdown_table, table_data
    else:
        return '', {'is_table': False}


# --- TWO-STAGE TABLE EXTRACTION (Structure First, Then OCR) ---

def extract_tables_two_stage(image_bytes: bytes, img_width: int, img_height: int, ocr_predictor) -> dict:
    """
    Two-stage table extraction:
    1. Detect table structure (cells/grid) WITHOUT OCR
    2. Crop each cell and run docTR OCR individually

    This keeps multi-line text together within cells.
    """
    try:
        # Convert bytes to numpy array for cropping
        nparr = np.frombuffer(image_bytes, np.uint8)
        img_array = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        if img_array is None:
            return {'is_table': False, 'error': 'Failed to decode image'}

        actual_height, actual_width = img_array.shape[:2]

        # Save image to temp file for img2table
        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
            tmp_file.write(image_bytes)
            tmp_path = tmp_file.name

        # Create img2table Image object
        img2table_img = Img2TableImage(src=tmp_path)

        # Stage 1: Detect structure ONLY (no OCR)
        # This gives us cell bounding boxes without text
        tables = img2table_img.extract_tables(
            ocr=None,  # No OCR - structure detection only
            implicit_rows=True,
            implicit_columns=True,
            borderless_tables=True
        )

        # Clean up temp file
        try:
            os.unlink(tmp_path)
        except:
            pass

        if not tables:
            print("Two-stage: No tables detected")
            return {'is_table': False, 'tables': []}

        # Process detected tables
        all_tables = []

        for table_idx, table in enumerate(tables):
            print(f"Two-stage: Processing table {table_idx + 1}")

            # Get table bounding box and cells
            if not hasattr(table, 'bbox') or not hasattr(table, 'content'):
                continue

            table_bbox = table.bbox  # (x1, y1, x2, y2)

            # Get cell structure - extract cell positions
            cells_data = []

            # img2table stores cells with their positions
            if hasattr(table, '_items') and table._items:
                # _items contains Cell objects with bbox
                rows_dict = {}

                for cell in table._items:
                    if hasattr(cell, 'bbox'):
                        cell_bbox = cell.bbox  # (x1, y1, x2, y2)
                        row_key = cell_bbox[1]  # y1 as row identifier

                        # Find or create row
                        matched_row = None
                        for existing_row in rows_dict.keys():
                            if abs(existing_row - row_key) < 10:  # 10px tolerance
                                matched_row = existing_row
                                break

                        if matched_row is None:
                            matched_row = row_key
                            rows_dict[matched_row] = []

                        rows_dict[matched_row].append({
                            'bbox': cell_bbox,
                            'x': cell_bbox[0]
                        })

                # Sort rows by y position
                sorted_rows = sorted(rows_dict.items(), key=lambda x: x[0])

                # Stage 2: OCR each cell individually
                table_cells = []

                for row_y, row_cells in sorted_rows:
                    # Sort cells in row by x position
                    row_cells.sort(key=lambda c: c['x'])
                    row_texts = []

                    for cell_info in row_cells:
                        bbox = cell_info['bbox']
                        x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])

                        # Add padding
                        padding = 2
                        x1 = max(0, x1 - padding)
                        y1 = max(0, y1 - padding)
                        x2 = min(actual_width, x2 + padding)
                        y2 = min(actual_height, y2 + padding)

                        # Crop cell from image
                        cell_img = img_array[y1:y2, x1:x2]

                        if cell_img.size == 0:
                            row_texts.append('')
                            continue

                        # Convert to RGB for docTR
                        cell_img_rgb = cv2.cvtColor(cell_img, cv2.COLOR_BGR2RGB)

                        # Run OCR on this cell
                        cell_text = ocr_single_cell(cell_img_rgb, ocr_predictor)
                        row_texts.append(cell_text)

                    if row_texts:
                        table_cells.append(row_texts)

                if table_cells:
                    num_cols = max(len(row) for row in table_cells)
                    # Normalize column count
                    normalized_cells = []
                    for row in table_cells:
                        normalized_row = row + [''] * (num_cols - len(row))
                        normalized_cells.append(normalized_row)

                    all_tables.append({
                        'cells': normalized_cells,
                        'num_rows': len(normalized_cells),
                        'num_columns': num_cols,
                        'method': 'two_stage'
                    })
                    print(f"Two-stage: Extracted {len(normalized_cells)}x{num_cols} table")

        if not all_tables:
            return {'is_table': False, 'tables': []}

        # Return the largest table
        primary_table = max(all_tables, key=lambda t: t['num_rows'] * t['num_columns'])

        return {
            'is_table': True,
            'cells': primary_table['cells'],
            'num_rows': primary_table['num_rows'],
            'num_columns': primary_table['num_columns'],
            'tables': all_tables,
            'total_tables': len(all_tables),
            'method': 'two_stage'
        }

    except Exception as e:
        print(f"Two-stage extraction error: {e}")
        import traceback
        traceback.print_exc()
        return {'is_table': False, 'error': str(e), 'method': 'two_stage'}


def ocr_single_cell(cell_image: np.ndarray, ocr_predictor) -> str:
    """
    Run OCR on a single cell image using docTR.
    Returns the extracted text with lines joined.
    """
    try:
        if cell_image.size == 0:
            return ''

        # Convert to PIL and then to bytes for docTR
        pil_img = Image.fromarray(cell_image)
        img_byte_arr = io.BytesIO()
        pil_img.save(img_byte_arr, format='PNG')
        img_bytes = img_byte_arr.getvalue()

        # Run docTR
        doc = DocumentFile.from_images([img_bytes])
        result = ocr_predictor(doc)

        # Extract text - join all lines with space (keeps multi-line together)
        lines = []
        for page in result.pages:
            for block in page.blocks:
                for line in block.lines:
                    line_text = ' '.join(word.value for word in line.words)
                    if line_text.strip():
                        lines.append(line_text.strip())

        # Join lines with space (not newline) to keep cell content together
        return ' '.join(lines)

    except Exception as e:
        print(f"Cell OCR error: {e}")
        return ''


def extract_text_two_stage(image_bytes: bytes, img_width: int, img_height: int, ocr_predictor) -> tuple:
    """
    Two-stage table extraction wrapper.
    Returns (markdown_text, table_data).
    """
    table_data = extract_tables_two_stage(image_bytes, img_width, img_height, ocr_predictor)

    if table_data.get('is_table'):
        markdown_table = format_table_as_markdown(table_data)
        return markdown_table, table_data
    else:
        return '', {'is_table': False, 'method': 'two_stage'}


# --- METHOD 3: BORDERLESS TABLE DETECTION (Text Position Clustering) ---

def extract_tables_borderless(doctr_result, min_columns: int = 2, min_rows: int = 2) -> dict:
    """
    Detect borderless tables by analyzing text positions from docTR.
    Works when there are no visible grid lines - uses whitespace gaps to infer structure.

    Algorithm:
    1. Collect all words with positions
    2. Find column boundaries by detecting consistent vertical gaps
    3. Group words into rows by y-position clustering
    4. Handle multi-line cells by merging text within same cell bounds
    """
    try:
        # Step 1: Collect all words with full position data
        all_words = []
        for page in doctr_result.pages:
            for block in page.blocks:
                for line in block.lines:
                    for word in line.words:
                        x_min, y_min = word.geometry[0]
                        x_max, y_max = word.geometry[1]
                        all_words.append({
                            'text': word.value,
                            'x_min': x_min,
                            'x_max': x_max,
                            'y_min': y_min,
                            'y_max': y_max,
                            'x_center': (x_min + x_max) / 2,
                            'y_center': (y_min + y_max) / 2,
                            'height': y_max - y_min
                        })

        if len(all_words) < 4:
            return {'is_table': False, 'reason': 'Too few words'}

        print(f"Borderless: Analyzing {len(all_words)} words")

        # Step 2: Detect columns by finding consistent x-position clusters
        columns = detect_column_boundaries(all_words)

        if len(columns) < min_columns:
            print(f"Borderless: Only {len(columns)} columns detected, need {min_columns}")
            return {'is_table': False, 'reason': f'Only {len(columns)} columns found'}

        print(f"Borderless: Detected {len(columns)} columns")

        # Step 3: Detect rows by clustering y-positions
        rows = detect_row_boundaries(all_words)

        if len(rows) < min_rows:
            print(f"Borderless: Only {len(rows)} rows detected, need {min_rows}")
            return {'is_table': False, 'reason': f'Only {len(rows)} rows found'}

        print(f"Borderless: Detected {len(rows)} rows")

        # Step 4: Assign words to cells and build table
        cells = build_table_cells(all_words, columns, rows)

        # Validate table has meaningful content
        non_empty_cells = sum(1 for row in cells for cell in row if cell.strip())
        total_cells = len(cells) * len(columns)
        fill_ratio = non_empty_cells / total_cells if total_cells > 0 else 0

        if fill_ratio < 0.3:  # Less than 30% cells filled - probably not a table
            print(f"Borderless: Low fill ratio {fill_ratio:.2f}, probably not a table")
            return {'is_table': False, 'reason': f'Low fill ratio: {fill_ratio:.2f}'}

        print(f"Borderless: Built {len(cells)}x{len(columns)} table with {fill_ratio:.2f} fill ratio")

        return {
            'is_table': True,
            'cells': cells,
            'num_rows': len(cells),
            'num_columns': len(columns),
            'method': 'borderless',
            'fill_ratio': fill_ratio
        }

    except Exception as e:
        print(f"Borderless extraction error: {e}")
        import traceback
        traceback.print_exc()
        return {'is_table': False, 'error': str(e), 'method': 'borderless'}


def detect_column_boundaries(words: list, min_gap: float = 0.03) -> list:
    """
    Detect column boundaries by finding consistent vertical gaps in text.
    Returns list of (x_start, x_end) tuples for each column.
    """
    if not words:
        return []

    # Get all unique x_min positions and sort them
    x_positions = sorted(set(w['x_min'] for w in words))

    if len(x_positions) < 2:
        return [(0, 1)]

    # Find gaps between x positions
    gaps = []
    for i in range(1, len(x_positions)):
        gap = x_positions[i] - x_positions[i-1]
        gaps.append((x_positions[i-1], x_positions[i], gap))

    # Find significant gaps (potential column separators)
    # A gap is significant if it's larger than min_gap and consistent across rows
    significant_gaps = []
    for x1, x2, gap in gaps:
        if gap >= min_gap:
            # Check if this gap is consistent (appears in multiple rows)
            gap_mid = (x1 + x2) / 2
            rows_with_gap = count_rows_with_gap(words, gap_mid, gap * 0.5)
            if rows_with_gap >= 2:  # Gap appears in at least 2 rows
                significant_gaps.append(gap_mid)

    # Build column boundaries from gaps
    if not significant_gaps:
        # No clear columns - try using x_min clustering
        return cluster_columns_by_alignment(words, min_gap)

    # Sort gaps and create column boundaries
    significant_gaps = sorted(set(significant_gaps))

    columns = []
    prev_x = 0
    for gap_x in significant_gaps:
        columns.append((prev_x, gap_x))
        prev_x = gap_x
    columns.append((prev_x, 1.0))

    return columns


def count_rows_with_gap(words: list, gap_x: float, tolerance: float) -> int:
    """Count how many rows have a gap at the given x position."""
    # Group words by approximate y position
    y_groups = {}
    for word in words:
        y_key = round(word['y_center'] * 20) / 20
        if y_key not in y_groups:
            y_groups[y_key] = []
        y_groups[y_key].append(word)

    rows_with_gap = 0
    for y_key, row_words in y_groups.items():
        # Check if there's a gap at gap_x in this row
        words_before = [w for w in row_words if w['x_max'] < gap_x - tolerance]
        words_after = [w for w in row_words if w['x_min'] > gap_x + tolerance]

        if words_before and words_after:
            rows_with_gap += 1

    return rows_with_gap


def cluster_columns_by_alignment(words: list, min_gap: float) -> list:
    """
    Cluster columns by finding words that align vertically.
    Used when gap detection doesn't find clear separators.
    """
    # Get x_min positions and cluster them
    x_mins = sorted(w['x_min'] for w in words)

    clusters = []
    current_cluster = [x_mins[0]]

    for i in range(1, len(x_mins)):
        if x_mins[i] - x_mins[i-1] <= min_gap:
            current_cluster.append(x_mins[i])
        else:
            clusters.append(current_cluster)
            current_cluster = [x_mins[i]]
    clusters.append(current_cluster)

    # Convert clusters to column boundaries
    if len(clusters) < 2:
        return [(0, 1)]

    columns = []
    for i, cluster in enumerate(clusters):
        x_start = min(cluster) - 0.01
        if i < len(clusters) - 1:
            x_end = (max(cluster) + min(clusters[i+1])) / 2
        else:
            x_end = 1.0
        columns.append((max(0, x_start), min(1, x_end)))

    return columns


def detect_row_boundaries(words: list, y_tolerance: float = 0.02) -> list:
    """
    Detect row boundaries by clustering y-positions.
    Returns list of (y_start, y_end) tuples for each row.
    """
    if not words:
        return []

    # Sort words by y_min
    sorted_by_y = sorted(words, key=lambda w: w['y_min'])

    # Cluster words into rows
    rows = []
    current_row = [sorted_by_y[0]]

    for i in range(1, len(sorted_by_y)):
        word = sorted_by_y[i]
        prev_word = current_row[-1]

        # Check if this word is on the same row
        # Words are on same row if their y ranges overlap significantly
        y_overlap = min(word['y_max'], prev_word['y_max']) - max(word['y_min'], prev_word['y_min'])
        min_height = min(word['height'], prev_word['height'])

        if y_overlap > min_height * 0.3 or abs(word['y_center'] - prev_word['y_center']) < y_tolerance:
            current_row.append(word)
        else:
            # New row
            row_y_min = min(w['y_min'] for w in current_row)
            row_y_max = max(w['y_max'] for w in current_row)
            rows.append((row_y_min, row_y_max, current_row))
            current_row = [word]

    # Don't forget last row
    if current_row:
        row_y_min = min(w['y_min'] for w in current_row)
        row_y_max = max(w['y_max'] for w in current_row)
        rows.append((row_y_min, row_y_max, current_row))

    return rows


def build_table_cells(words: list, columns: list, rows: list) -> list:
    """
    Build table cells by assigning words to their respective cells.
    Handles multi-line text within cells.
    """
    num_cols = len(columns)
    table = []

    for row_y_min, row_y_max, row_words in rows:
        row_cells = [''] * num_cols

        # Sort words in row by x position
        row_words_sorted = sorted(row_words, key=lambda w: w['x_min'])

        for word in row_words_sorted:
            # Find which column this word belongs to
            word_x = word['x_min']

            for col_idx, (col_start, col_end) in enumerate(columns):
                if col_start <= word_x < col_end:
                    # Add word to this cell
                    if row_cells[col_idx]:
                        row_cells[col_idx] += ' ' + word['text']
                    else:
                        row_cells[col_idx] = word['text']
                    break

        table.append(row_cells)

    return table


def extract_text_borderless(doctr_result) -> tuple:
    """
    Borderless table extraction wrapper.
    Returns (markdown_text, table_data).
    """
    table_data = extract_tables_borderless(doctr_result)

    if table_data.get('is_table'):
        markdown_table = format_table_as_markdown(table_data)
        return markdown_table, table_data
    else:
        return '', {'is_table': False, 'method': 'borderless'}


# --- METHOD 4: BLOCK-GEOMETRY TABLE DETECTION (docTR Block Grouping) ---

def extract_tables_block_geometry(doctr_result, min_columns: int = 2, min_rows: int = 2) -> dict:
    """
    Detect tables using docTR's block-level grouping from .export().
    If multiple blocks exist at similar y-positions but different x-positions,
    they likely represent table columns.
    """
    try:
        exported = doctr_result.export()

        if not exported or 'pages' not in exported or not exported['pages']:
            return {'is_table': False, 'reason': 'No pages in export', 'method': 'block_geometry'}

        page = exported['pages'][0]
        blocks = page.get('blocks', [])

        if len(blocks) < 2:
            return {'is_table': False, 'reason': f'Only {len(blocks)} blocks found', 'method': 'block_geometry'}

        print(f"Block-geometry: Analyzing {len(blocks)} blocks")

        # Step 1: Extract block positions and text content
        block_data = []
        for block in blocks:
            geometry = block.get('geometry', [])
            if len(geometry) < 2:
                continue

            x_min, y_min = geometry[0]
            x_max, y_max = geometry[1]

            block_text_parts = []
            for line in block.get('lines', []):
                line_words = []
                for word in line.get('words', []):
                    line_words.append(word.get('value', ''))
                if line_words:
                    block_text_parts.append(' '.join(line_words))

            block_text = ' '.join(block_text_parts).strip()

            if block_text:
                block_data.append({
                    'text': block_text,
                    'x_min': x_min,
                    'x_max': x_max,
                    'y_min': y_min,
                    'y_max': y_max,
                    'y_center': (y_min + y_max) / 2,
                    'x_center': (x_min + x_max) / 2,
                    'height': y_max - y_min,
                })

        if len(block_data) < min_columns:
            return {'is_table': False, 'reason': f'Only {len(block_data)} text blocks', 'method': 'block_geometry'}

        # Step 2: Group blocks into rows by y-position overlap
        block_data.sort(key=lambda b: b['y_min'])

        rows = []
        current_row = [block_data[0]]

        for i in range(1, len(block_data)):
            block = block_data[i]
            prev_block = current_row[-1]

            y_overlap = min(block['y_max'], prev_block['y_max']) - max(block['y_min'], prev_block['y_min'])
            min_height = min(block['height'], prev_block['height'])

            if min_height > 0 and y_overlap / min_height > 0.3:
                current_row.append(block)
            else:
                rows.append(current_row)
                current_row = [block]

        if current_row:
            rows.append(current_row)

        print(f"Block-geometry: Found {len(rows)} potential rows")

        # Step 3: Filter to rows that have multiple blocks (potential table rows)
        multi_block_rows = [row for row in rows if len(row) >= min_columns]

        if len(multi_block_rows) < min_rows:
            print(f"Block-geometry: Only {len(multi_block_rows)} multi-block rows, need {min_rows}")
            return {'is_table': False, 'reason': f'Only {len(multi_block_rows)} multi-block rows', 'method': 'block_geometry'}

        # Step 4: Check consistency - do multiple rows have similar column counts?
        col_counts = [len(row) for row in multi_block_rows]
        most_common_count = max(set(col_counts), key=col_counts.count)
        consistent_rows = [row for row in multi_block_rows if len(row) == most_common_count]

        if len(consistent_rows) < min_rows:
            print(f"Block-geometry: Only {len(consistent_rows)} rows with {most_common_count} columns")
            return {'is_table': False, 'reason': 'Inconsistent column counts', 'method': 'block_geometry'}

        print(f"Block-geometry: {len(consistent_rows)} rows with {most_common_count} columns")

        # Step 5: Sort blocks within each row by x-position and extract cell text
        table_cells = []
        for row in consistent_rows:
            row_sorted = sorted(row, key=lambda b: b['x_min'])
            row_texts = [b['text'] for b in row_sorted]
            table_cells.append(row_texts)

        # Normalize column count
        max_cols = max(len(row) for row in table_cells) if table_cells else 0
        normalized_cells = []
        for row in table_cells:
            normalized_row = row + [''] * (max_cols - len(row))
            normalized_cells.append(normalized_row)

        # Validate fill ratio
        non_empty = sum(1 for row in normalized_cells for cell in row if cell.strip())
        total = len(normalized_cells) * max_cols
        fill_ratio = non_empty / total if total > 0 else 0

        if fill_ratio < 0.3:
            print(f"Block-geometry: Low fill ratio {fill_ratio:.2f}")
            return {'is_table': False, 'reason': f'Low fill ratio: {fill_ratio:.2f}', 'method': 'block_geometry'}

        print(f"Block-geometry: Built {len(normalized_cells)}x{max_cols} table with {fill_ratio:.2f} fill ratio")

        return {
            'is_table': True,
            'cells': normalized_cells,
            'num_rows': len(normalized_cells),
            'num_columns': max_cols,
            'method': 'block_geometry',
            'fill_ratio': fill_ratio,
        }

    except Exception as e:
        print(f"Block-geometry extraction error: {e}")
        import traceback
        traceback.print_exc()
        return {'is_table': False, 'error': str(e), 'method': 'block_geometry'}


def extract_text_block_geometry(doctr_result) -> tuple:
    """Block-geometry table extraction wrapper."""
    table_data = extract_tables_block_geometry(doctr_result)

    if table_data.get('is_table'):
        markdown_table = format_table_as_markdown(table_data)
        return markdown_table, table_data
    else:
        return '', {'is_table': False, 'method': 'block_geometry'}


def extract_text_structured(result) -> str:
    """
    Extract text from docTR result preserving logical structure.
    Explicitly sorts words by x-coordinate and lines by y-coordinate.
    """
    all_lines = []

    for page in result.pages:
        for block in page.blocks:
            for line in block.lines:
                # Collect words with their x-positions
                words_data = []
                for word in line.words:
                    # geometry is ((x_min, y_min), (x_max, y_max))
                    x_pos = word.geometry[0][0]  # Left edge x-coordinate
                    y_pos = word.geometry[0][1]  # Top edge y-coordinate
                    words_data.append({
                        'text': word.value,
                        'x': x_pos,
                        'y': y_pos
                    })

                if not words_data:
                    continue

                # Sort words by x position (left to right)
                words_data.sort(key=lambda w: w['x'])

                line_text = " ".join([w['text'] for w in words_data])
                avg_y = sum(w['y'] for w in words_data) / len(words_data)
                min_x = min(w['x'] for w in words_data)

                if line_text.strip():
                    all_lines.append({
                        'text': line_text.strip(),
                        'y': avg_y,
                        'x': min_x
                    })

    # Sort lines by y position (top to bottom)
    all_lines.sort(key=lambda l: (round(l['y'] * 20) / 20, l['x']))

    # Group lines by y-position and build final text
    result_lines = []
    prev_y_group = -1
    current_line_parts = []

    for line_info in all_lines:
        current_y_group = round(line_info['y'] * 20) / 20

        if prev_y_group != -1 and current_y_group != prev_y_group:
            if current_line_parts:
                result_lines.append(" ".join(current_line_parts))
            current_line_parts = [line_info['text']]
        else:
            current_line_parts.append(line_info['text'])

        prev_y_group = current_y_group

    if current_line_parts:
        result_lines.append(" ".join(current_line_parts))

    return "\n".join(result_lines)

def generate_synthesized_image(doctr_result) -> Optional[str]:
    """
    Generate a reconstructed document image using docTR's synthesize() method.
    Returns a base64-encoded PNG string, or None if synthesis fails.
    """
    try:
        synthetic_pages = doctr_result.synthesize()

        if not synthetic_pages or len(synthetic_pages) == 0:
            print("Synthesize: No pages returned")
            return None

        # Take the first page (single-page processing)
        synth_img = synthetic_pages[0]

        # synth_img is a numpy array (H, W, C) in uint8 format
        pil_img = Image.fromarray(synth_img)
        img_byte_arr = io.BytesIO()
        pil_img.save(img_byte_arr, format='PNG')
        img_bytes = img_byte_arr.getvalue()

        b64_string = base64.b64encode(img_bytes).decode('utf-8')
        print(f"Synthesize: Generated image ({len(b64_string)} chars base64)")
        return b64_string

    except Exception as e:
        print(f"Synthesize error: {e}")
        return None


def extract_words_with_boxes(result) -> list:
    """
    Extract all words with their bounding boxes and confidence from docTR result.
    Returns list of {word, confidence, bbox} where bbox is [[x0,y0], [x1,y1]] normalized 0-1.
    """
    words_with_boxes = []

    for page in result.pages:
        for block in page.blocks:
            for line in block.lines:
                for word in line.words:
                    # geometry is ((x0, y0), (x1, y1)) normalized
                    bbox = [
                        [word.geometry[0][0], word.geometry[0][1]],
                        [word.geometry[1][0], word.geometry[1][1]]
                    ]
                    words_with_boxes.append({
                        'word': word.value,
                        'confidence': word.confidence,
                        'bbox': bbox
                    })

    return words_with_boxes

def check_drug_interactions(detected_drugs: List[str]) -> List[Dict]:
    """
    Check for known interactions between detected drugs.
    Returns list of interaction warnings.
    """
    interactions = []
    drugs_lower = [d.lower().strip() for d in detected_drugs]

    # Check each pair of drugs
    for i, drug1 in enumerate(drugs_lower):
        for drug2 in drugs_lower[i+1:]:
            # Check if drug1 interacts with drug2
            if drug1 in DRUG_INTERACTIONS:
                if drug2 in DRUG_INTERACTIONS[drug1]:
                    interaction = DRUG_INTERACTIONS[drug1][drug2]
                    interactions.append({
                        'drug1': detected_drugs[i],
                        'drug2': detected_drugs[drugs_lower.index(drug2)],
                        'severity': interaction.get('severity', 'info'),
                        'description': interaction.get('description', ''),
                        'recommendation': interaction.get('recommendation'),
                    })
            # Check reverse (drug2 interacts with drug1)
            elif drug2 in DRUG_INTERACTIONS:
                if drug1 in DRUG_INTERACTIONS[drug2]:
                    interaction = DRUG_INTERACTIONS[drug2][drug1]
                    interactions.append({
                        'drug1': detected_drugs[drugs_lower.index(drug2)],
                        'drug2': detected_drugs[i],
                        'severity': interaction.get('severity', 'info'),
                        'description': interaction.get('description', ''),
                        'recommendation': interaction.get('recommendation'),
                    })

    return interactions

# ==================== LAB VALUE ANALYSIS ====================

def parse_reference_range(range_str: str):
    """
    Parse reference range strings from lab documents.
    Formats: "(13.5 - 18.0)", "(< 200)", "(> 60)", "(< 0.61)"
    Returns: (low, high) where either can be None
    """
    if not range_str:
        return None, None

    # Clean up the string
    s = range_str.strip().strip('()').strip()

    # Pattern: "< value" (upper limit only)
    m = re.match(r'^[<\u2264]\s*(\d+\.?\d*)$', s)
    if m:
        return None, float(m.group(1))

    # Pattern: "> value" (lower limit only)
    m = re.match(r'^[>\u2265]\s*(\d+\.?\d*)$', s)
    if m:
        return float(m.group(1)), None

    # Pattern: "low - high"
    m = re.match(r'(\d+\.?\d*)\s*[-\u2013]\s*(\d+\.?\d*)', s)
    if m:
        return float(m.group(1)), float(m.group(2))

    return None, None


def extract_lab_values_from_words(words_with_boxes: List[Dict]) -> List[Dict]:
    """
    Extract lab values using word positions from docTR.
    Groups words into rows by y-coordinate, then identifies columns
    (test name, value, unit, range) by x-position within each row.
    This is the most reliable method since it uses spatial layout.
    """
    extracted = []
    if not words_with_boxes:
        return extracted

    # 1. Group words into rows by y-center (within tolerance)
    ROW_TOLERANCE = 0.015  # Words within 1.5% of page height = same row
    rows = []
    sorted_words = sorted(words_with_boxes, key=lambda w: (w['bbox'][0][1], w['bbox'][0][0]))

    current_row = []
    current_y = None

    for word_info in sorted_words:
        y_center = (word_info['bbox'][0][1] + word_info['bbox'][1][1]) / 2
        if current_y is None or abs(y_center - current_y) < ROW_TOLERANCE:
            current_row.append(word_info)
            if current_y is None:
                current_y = y_center
            else:
                current_y = (current_y + y_center) / 2  # Running average
        else:
            if current_row:
                rows.append(sorted(current_row, key=lambda w: w['bbox'][0][0]))
            current_row = [word_info]
            current_y = y_center

    if current_row:
        rows.append(sorted(current_row, key=lambda w: w['bbox'][0][0]))

    # 2. For each row, classify words into: test_name, value, unit, range
    UNITS = {'mg/dl', 'mmol/l', 'g/dl', 'u/l', 'miu/l', 'ng/dl', 'pg/ml',
             'ug/dl', 'ng/ml', 'fl', 'pg', '%', 'mm/hr', 'mg/l', 'mg/mmol',
             'ug/l', 'ml/min/1.73m2'}

    SKIP_WORDS = {'result', 'unit', 'ref.range', 'ref', 'range', 'reference',
                  'date', 'request', 'no', 'no:'}

    for row in rows:
        words_text = [w['word'] for w in row]
        row_str = ' '.join(words_text).lower()

        # Skip header rows
        if 'result' in row_str and ('unit' in row_str or 'ref' in row_str):
            continue
        if 'profile' in row_str and len(words_text) <= 3:
            continue
        if 'function' in row_str and len(words_text) <= 3:
            continue

        # Classify each word
        name_parts = []
        value = None
        unit = ''
        range_parts = []
        is_flagged = False
        in_range = False

        for w in row:
            word = w['word'].strip()
            word_lower = word.lower().strip('()')

            if not word:
                continue

            # Check if this starts/continues a range (in parentheses)
            if '(' in word or in_range:
                in_range = True
                range_parts.append(word)
                if ')' in word:
                    in_range = False
                continue

            # Check for flagged marker
            if word == '*':
                is_flagged = True
                continue

            # Check if it's a unit
            if word_lower in UNITS or word_lower.replace('/', '').replace('.', '').replace('1', '').replace('3', '').replace('7', '').replace('m', '').replace('2', '') == '':
                cleaned_unit = word_lower
                if cleaned_unit in UNITS:
                    unit = word
                    continue

            # Check if unit with superscript like x10⁹/L or x10^9/L
            if 'x10' in word_lower or '10⁹' in word or '10¹²' in word:
                unit = word
                continue

            # Check if it's a number (the result value)
            cleaned_word = word.lstrip('*').strip()
            try:
                num = float(cleaned_word)
                if value is None:
                    value = num
                    if '*' in word:
                        is_flagged = True
                continue
            except ValueError:
                pass

            # Check if it's a skip word
            if word_lower in SKIP_WORDS:
                continue

            # Check if it's Chinese characters only — skip
            if all('\u4e00' <= c <= '\u9fff' or c in '()（）' for c in word):
                continue

            # Otherwise it's part of the test name
            if any(c.isalpha() for c in word):
                name_parts.append(word)

        # Parse the range
        range_str = ' '.join(range_parts).strip('() ')
        ref_low, ref_high = parse_reference_range(range_str)

        test_name = ' '.join(name_parts).strip()

        # Validate: need at least a name, a value, and a range
        if test_name and value is not None and (ref_low is not None or ref_high is not None):
            # Filter out section headers that slipped through
            if test_name.upper() == test_name and len(test_name.split()) > 2:
                continue  # ALL CAPS multi-word = likely a section header

            extracted.append({
                'test_name': test_name,
                'value': value,
                'unit': unit,
                'ref_low': ref_low,
                'ref_high': ref_high,
                'ref_range_str': range_str,
                'is_flagged_in_document': is_flagged,
            })

    return extracted


def extract_lab_values_from_text(structured_text: str) -> List[Dict]:
    """
    Extract test name, value, unit, and reference range from OCR structured text.
    Handles the document format: TestName [ChineseName] Result Unit (Range)
    """
    extracted = []
    if not structured_text:
        return extracted

    lines = structured_text.split('\n')
    for line in lines:
        line = line.strip()
        if not line or len(line) < 5:
            continue

        # Pattern to match lab result lines
        # Test Name [optional Chinese] [optional *] Number [Unit] [(Range)]
        # e.g., "Sodium  钠  142  mmol/L  (136 - 145)"
        # e.g., "HDL Cholesterol  高脂蛋白(好)胆固醇  * 38  mg/dL  (40 - 59)"

        # Try to find a reference range at the end of the line
        range_match = re.search(r'\(([<>\u2264\u2265]?\s*\d+\.?\d*(?:\s*[-\u2013]\s*\d+\.?\d*)?)\)\s*$', line)
        ref_range_str = None
        ref_low, ref_high = None, None
        if range_match:
            ref_range_str = range_match.group(1)
            ref_low, ref_high = parse_reference_range(ref_range_str)
            line = line[:range_match.start()].strip()

        # Skip if no reference range found (likely not a lab result line)
        if ref_low is None and ref_high is None:
            continue

        # Try to extract: test name, numeric value, unit
        # Look for the numeric value (possibly preceded by *)
        value_match = re.search(r'\*?\s*(\d+\.?\d*)\s+(mg/dL|mmol/L|g/dL|U/L|mIU/L|ng/dL|pg/mL|ug/dL|ng/mL|fL|pg|%|mm/hr|mg/L|mg/mmol|x10\^?\d+/L|mL/min/1\.73m2|ug/L)', line, re.IGNORECASE)

        if not value_match:
            # Try without unit (some lines have unit in separate position)
            value_match = re.search(r'\*?\s*(\d+\.?\d+)\s*$', line)
            if not value_match:
                # Try matching just a number in the middle of the line
                value_match = re.search(r'(?:[\u4e00-\u9fff\s]+|\s+)\*?\s*(\d+\.?\d*)\s', line)

        if not value_match:
            continue

        try:
            value = float(value_match.group(1))
        except (ValueError, IndexError):
            continue

        # Extract unit if captured
        unit = ''
        if value_match.lastindex and value_match.lastindex >= 2:
            unit = value_match.group(2)

        # Extract test name (everything before the Chinese characters or the value)
        # Find where Chinese characters start
        chinese_start = re.search(r'[\u4e00-\u9fff]', line)
        if chinese_start:
            test_name = line[:chinese_start.start()].strip()
        else:
            test_name = line[:value_match.start()].strip()

        # Clean up test name
        test_name = test_name.strip().rstrip(':').strip()

        # Check if it's flagged with * in the original document
        is_flagged = '*' in line[:value_match.end()]

        if test_name and len(test_name) >= 2:
            extracted.append({
                'test_name': test_name,
                'value': value,
                'unit': unit,
                'ref_low': ref_low,
                'ref_high': ref_high,
                'ref_range_str': ref_range_str or '',
                'is_flagged_in_document': is_flagged,
            })

    return extracted


def extract_lab_values_from_table(table_data: Dict) -> List[Dict]:
    """
    Extract lab values from structured table data.
    table_data has 'cells' (list of rows, each row is list of cell strings).
    """
    extracted = []
    cells = table_data.get('cells', [])
    if not cells or len(cells) < 2:
        return extracted

    # Try to identify column indices by looking at the header row
    header_row = cells[0] if cells else []
    name_col = -1
    value_col = -1
    unit_col = -1
    range_col = -1

    for i, cell in enumerate(header_row):
        cell_lower = cell.strip().lower()
        if any(kw in cell_lower for kw in ['test', 'name', 'parameter', 'investigation']):
            name_col = i
        elif 'result' in cell_lower:
            value_col = i
        elif 'unit' in cell_lower:
            unit_col = i
        elif any(kw in cell_lower for kw in ['ref', 'range', 'normal', 'reference']):
            range_col = i

    # If we couldn't identify columns from headers, use heuristics
    if name_col == -1 or value_col == -1:
        # Fall back to row-by-row parsing
        for row in cells[1:]:  # Skip header
            if len(row) < 2:
                continue

            test_name = None
            value = None
            unit = ''
            ref_range_str = ''
            is_flagged = False

            for cell in row:
                cell_stripped = cell.strip()
                if not cell_stripped:
                    continue

                # Check for reference range (in parentheses)
                range_m = re.match(r'^\(([<>\u2264\u2265]?\s*\d+\.?\d*(?:\s*[-\u2013]\s*\d+\.?\d*)?)\)$', cell_stripped)
                if range_m:
                    ref_range_str = range_m.group(1)
                    continue

                # Check for numeric value (possibly with *)
                val_m = re.match(r'^\*?\s*(\d+\.?\d*)$', cell_stripped)
                if val_m and test_name is not None and value is None:
                    value = float(val_m.group(1))
                    is_flagged = '*' in cell_stripped
                    continue

                # Check for unit
                if cell_stripped.lower() in ['mg/dl', 'mmol/l', 'g/dl', 'u/l', 'miu/l', 'ng/dl',
                                              'pg/ml', 'ug/dl', 'ng/ml', 'fl', 'pg', '%', 'mm/hr',
                                              'mg/l', 'mg/mmol', 'ug/l', 'ml/min/1.73m2']:
                    unit = cell_stripped
                    continue

                # Check for unit with superscript patterns
                if re.match(r'x10\^?\d+/L', cell_stripped, re.IGNORECASE):
                    unit = cell_stripped
                    continue

                # Otherwise, if it has alphabetic content, treat as test name
                if any(c.isalpha() for c in cell_stripped) and test_name is None:
                    # Skip Chinese-only cells
                    if not all('\u4e00' <= c <= '\u9fff' or c.isspace() for c in cell_stripped):
                        test_name = cell_stripped

            if test_name and value is not None and ref_range_str:
                ref_low, ref_high = parse_reference_range(ref_range_str)
                if ref_low is not None or ref_high is not None:
                    extracted.append({
                        'test_name': test_name,
                        'value': value,
                        'unit': unit,
                        'ref_low': ref_low,
                        'ref_high': ref_high,
                        'ref_range_str': ref_range_str,
                        'is_flagged_in_document': is_flagged,
                    })
        return extracted

    # Parse using identified column indices
    for row in cells[1:]:
        if len(row) <= max(name_col, value_col):
            continue

        test_name = row[name_col].strip() if name_col < len(row) else ''
        value_str = row[value_col].strip() if value_col < len(row) else ''
        unit = row[unit_col].strip() if unit_col >= 0 and unit_col < len(row) else ''
        ref_range_str = row[range_col].strip().strip('()') if range_col >= 0 and range_col < len(row) else ''

        if not test_name or not value_str:
            continue

        # Parse value (handle * prefix)
        is_flagged = '*' in value_str
        val_m = re.search(r'(\d+\.?\d*)', value_str)
        if not val_m:
            continue

        value = float(val_m.group(1))
        ref_low, ref_high = parse_reference_range(ref_range_str)

        if ref_low is not None or ref_high is not None:
            extracted.append({
                'test_name': test_name,
                'value': value,
                'unit': unit,
                'ref_low': ref_low,
                'ref_high': ref_high,
                'ref_range_str': ref_range_str,
                'is_flagged_in_document': is_flagged,
            })

    return extracted


def classify_lab_value(value: float, ref_low, ref_high) -> str:
    """
    Classify a lab value against reference range.
    Returns: 'critical_low', 'low', 'normal', 'high', 'critical_high'
    """
    if ref_low is not None and value < ref_low:
        # Check if critically low (below 70% of lower limit)
        if value < ref_low * 0.7:
            return 'critical_low'
        return 'low'

    if ref_high is not None and value > ref_high:
        # Check if critically high (above 150% of upper limit)
        if value > ref_high * 1.5:
            return 'critical_high'
        return 'high'

    return 'normal'


def match_test_to_medlineplus(test_name: str) -> Optional[Dict]:
    """
    Fuzzy-match a test name against the MedlinePlus map.
    Returns the map entry if matched, None otherwise.
    """
    if not MEDLINEPLUS_MAP:
        return None

    name_lower = test_name.lower().strip()

    # 1. Exact match on key
    if name_lower in MEDLINEPLUS_MAP:
        return MEDLINEPLUS_MAP[name_lower]

    # 2. Exact match on aliases
    for key, data in MEDLINEPLUS_MAP.items():
        aliases = [a.lower() for a in data.get('aliases', [])]
        if name_lower in aliases:
            return data

    # 3. Partial match (test name contains or is contained in a key/alias)
    for key, data in MEDLINEPLUS_MAP.items():
        if key in name_lower or name_lower in key:
            return data
        for alias in data.get('aliases', []):
            if alias.lower() in name_lower or name_lower in alias.lower():
                return data

    # 4. Fuzzy match using difflib
    all_names = list(MEDLINEPLUS_MAP.keys())
    for key, data in MEDLINEPLUS_MAP.items():
        all_names.extend([a.lower() for a in data.get('aliases', [])])

    close = difflib.get_close_matches(name_lower, all_names, n=1, cutoff=0.7)
    if close:
        matched_name = close[0]
        if matched_name in MEDLINEPLUS_MAP:
            return MEDLINEPLUS_MAP[matched_name]
        for key, data in MEDLINEPLUS_MAP.items():
            if matched_name in [a.lower() for a in data.get('aliases', [])]:
                return data

    return None


def get_medlineplus_info(slug: str, status: str) -> Dict:
    """
    Get educational info from MedlinePlus cache for a given test slug and status.
    Falls back to fetching from MedlinePlus if not cached.
    """
    url = f"https://medlineplus.gov/lab-tests/{slug}/"

    # Check cache
    if slug in MEDLINEPLUS_CACHE:
        cached = MEDLINEPLUS_CACHE[slug]
        direction = 'high' if 'high' in status else 'low'
        return {
            'url': cached.get('url', url),
            'description': cached.get(direction, ''),
        }

    # Try to fetch from MedlinePlus (with timeout)
    try:
        response = httpx.get(url, timeout=5.0, follow_redirects=True)
        if response.status_code == 200:
            soup = BeautifulSoup(response.text, 'html.parser')

            # Find "What do the results mean?" section
            results_section = None
            for heading in soup.find_all(['h2', 'h3']):
                if 'results' in heading.get_text().lower() and 'mean' in heading.get_text().lower():
                    results_section = heading
                    break

            description = ''
            if results_section:
                # Get all text until the next heading
                content_parts = []
                for sibling in results_section.find_next_siblings():
                    if sibling.name in ['h2', 'h3']:
                        break
                    text = sibling.get_text(strip=True)
                    if text:
                        content_parts.append(text)
                description = ' '.join(content_parts[:3])  # First 3 paragraphs

            # Cache the result
            MEDLINEPLUS_CACHE[slug] = {
                'url': url,
                'high': description,
                'low': description,
                'fetched_at': 'runtime'
            }

            return {
                'url': url,
                'description': description,
            }
    except Exception as e:
        print(f"MedlinePlus fetch failed for {slug}: {e}")

    return {'url': url, 'description': ''}


def check_lab_values(structured_text: str, table_data: Optional[Dict], words_with_boxes: Optional[List[Dict]] = None) -> List[Dict]:
    """
    Extract lab values from OCR output and check against reference ranges.
    Uses three extraction methods in priority order:
    1. Word-position-based (most reliable — uses spatial layout from docTR)
    2. Table-based (if table was detected)
    3. Text regex-based (fallback)
    Returns list of lab anomaly results.
    """
    # Method 1: Word-position-based extraction (best for columnar lab reports)
    extracted = []
    if words_with_boxes:
        extracted = extract_lab_values_from_words(words_with_boxes)
        print(f"Lab extraction (word-position): found {len(extracted)} values")

    # Method 2: Table-based extraction
    if table_data and table_data.get('is_table'):
        table_extracted = extract_lab_values_from_table(table_data)
        print(f"Lab extraction (table): found {len(table_extracted)} values")
        existing_names = {e['test_name'].lower() for e in extracted}
        for te in table_extracted:
            if te['test_name'].lower() not in existing_names:
                extracted.append(te)
                existing_names.add(te['test_name'].lower())

    # Method 3: Text regex fallback
    text_extracted = extract_lab_values_from_text(structured_text)
    print(f"Lab extraction (text-regex): found {len(text_extracted)} values")

    # Merge: add text-extracted values if test name not already found
    existing_names = {e['test_name'].lower() for e in extracted}
    for te in text_extracted:
        if te['test_name'].lower() not in existing_names:
            extracted.append(te)
            existing_names.add(te['test_name'].lower())

    # Check each extracted value
    results = []
    for item in extracted:
        status = classify_lab_value(item['value'], item['ref_low'], item['ref_high'])

        # Build reference range display string
        if item['ref_low'] is not None and item['ref_high'] is not None:
            range_display = f"{item['ref_low']} - {item['ref_high']}"
        elif item['ref_high'] is not None:
            range_display = f"< {item['ref_high']}"
        elif item['ref_low'] is not None:
            range_display = f"> {item['ref_low']}"
        else:
            range_display = item.get('ref_range_str', '')

        # Look up MedlinePlus info
        medlineplus_entry = match_test_to_medlineplus(item['test_name'])
        description = ''
        medlineplus_url = None
        category = 'General'

        if medlineplus_entry:
            slug = medlineplus_entry.get('slug', '')
            category = medlineplus_entry.get('category', 'General')
            medlineplus_url = f"https://medlineplus.gov/lab-tests/{slug}/"

            if status != 'normal' and slug:
                info = get_medlineplus_info(slug, status)
                description = info.get('description', '')
                medlineplus_url = info.get('url', medlineplus_url)

        results.append({
            'test_name': item['test_name'],
            'value': item['value'],
            'unit': item['unit'],
            'status': status,
            'ref_low': item['ref_low'],
            'ref_high': item['ref_high'],
            'reference_range': range_display,
            'category': category,
            'description': description,
            'medlineplus_url': medlineplus_url,
            'is_flagged_in_document': item.get('is_flagged_in_document', False),
        })

    return results

def map_entities_to_boxes(entities: list, words_with_boxes: list, cleaned_text: str) -> list:
    """
    Map NER entities back to word bounding boxes.
    Uses fuzzy matching to find entity words in OCR words.
    """
    entities_with_boxes = []

    for entity in entities:
        entity_word = entity['word'].lower().strip()
        entity_parts = entity_word.split()

        # Find matching word(s) in OCR output
        matched_boxes = []
        for word_info in words_with_boxes:
            ocr_word = word_info['word'].lower().strip()
            # Check if OCR word matches any part of entity
            for part in entity_parts:
                if part in ocr_word or ocr_word in part:
                    matched_boxes.append(word_info['bbox'])
                    break

        # Combine bounding boxes if multiple matches
        if matched_boxes:
            # Get bounding box that encompasses all matched words
            min_x = min(box[0][0] for box in matched_boxes)
            min_y = min(box[0][1] for box in matched_boxes)
            max_x = max(box[1][0] for box in matched_boxes)
            max_y = max(box[1][1] for box in matched_boxes)
            combined_bbox = [[min_x, min_y], [max_x, max_y]]
        else:
            combined_bbox = None

        entities_with_boxes.append({
            'entity_group': entity['entity_group'],
            'score': entity['score'],
            'word': entity['word'],
            'bbox': combined_bbox
        })

    return entities_with_boxes

# --- FastAPI Routes ---

@app.get("/")
async def root():
    """Health check endpoint."""
    return {"status": "running", "message": "ScanAssured OCR & NER API"}

@app.get("/models")
async def get_available_models():
    """Return all available OCR and NER models."""
    return {
        "ocr_presets": [
            {
                "id": preset_id,
                "name": preset_data["name"],
                "description": preset_data["description"]
            }
            for preset_id, preset_data in OCR_PRESETS.items()
        ],
        "ocr_detection_models": OCR_DETECTION_MODELS,
        "ocr_recognition_models": OCR_RECOGNITION_MODELS,
        "ner_models": {
            model_id: {
                "name": model_data["name"],
                "description": model_data["description"],
                "entities": model_data["entities"]
            }
            for model_id, model_data in NER_MODELS.items()
        },
        "ocr_correction_model": {
            "id": "ner-dictionary",
            "name": "NER-Informed Dictionary Correction",
            "description": "Edit-distance correction against medical entity dictionaries, guided by NER entity labels",
        }
    }

@app.post("/process")
async def process_image(
    file: UploadFile = File(...),
    ner_model_id: str = Form(...),
    ocr_preset: str = Form("balanced"),
    ocr_det_model: Optional[str] = Form(None),
    ocr_reco_model: Optional[str] = Form(None),
    enable_correction: str = Form("false"),
    correction_threshold: str = Form("0.75"),
):
    """Process an image with OCR and NER."""

    # Determine OCR models
    if ocr_det_model and ocr_reco_model:
        det_arch = ocr_det_model
        reco_arch = ocr_reco_model
    else:
        preset = OCR_PRESETS.get(ocr_preset, OCR_PRESETS["balanced"])
        det_arch = preset["det"]
        reco_arch = preset["reco"]

    # Validate NER model
    if ner_model_id not in NER_MODELS:
        return JSONResponse(
            status_code=400,
            content={"detail": f"Unknown NER model: {ner_model_id}"}
        )

    # Get OCR predictor
    ocr_predictor_instance = get_ocr_predictor(det_arch, reco_arch)
    if not ocr_predictor_instance:
        return JSONResponse(
            status_code=503,
            content={"detail": f"Failed to load OCR model: {det_arch}/{reco_arch}"}
        )

    # Get NER pipeline
    ner_pipeline = get_ner_pipeline(ner_model_id)
    if not ner_pipeline:
        return JSONResponse(
            status_code=503,
            content={"detail": f"Failed to load NER model: {ner_model_id}"}
        )

    try:
        # Read and preprocess image
        file_content = await file.read()
        preprocessed_img = preprocess_for_doctr(file_content)

        # Perform OCR with docTR
        print("Running docTR OCR...")
        # Convert numpy array to bytes for docTR
        pil_img = Image.fromarray(preprocessed_img)
        img_byte_arr = io.BytesIO()
        pil_img.save(img_byte_arr, format='PNG')
        img_bytes = img_byte_arr.getvalue()

        doc = DocumentFile.from_images([img_bytes])
        result = ocr_predictor_instance(doc)

        # Get image dimensions for frontend highlighting
        img_height, img_width = preprocessed_img.shape[:2]

        # Extract text and word bounding boxes using docTR
        structured_text = extract_text_structured(result)
        cleaned_text = basic_cleanup(structured_text)
        words_with_boxes = extract_words_with_boxes(result)

        print(f"OCR Structured Text:\n{structured_text[:500]}...")
        print(f"Extracted {len(words_with_boxes)} words with bounding boxes")

        # Generate synthesized (reconstructed) image
        print("Generating synthesized document image...")
        synthesized_image = generate_synthesized_image(result)

        # --- DOCLING PIPELINE (runs on raw bytes, not preprocessed) ---
        print("Running Docling pipeline for comparison...")
        docling_result = run_docling_pipeline(file_content)

        # Method 1: img2table with built-in OCR
        print("Running img2table for table detection (Method 1: integrated OCR)...")
        table_formatted_text, table_data = extract_text_with_table_detection(
            img_bytes, img_width, img_height
        )

        # Method 2: Two-stage (structure first, then OCR each cell)
        print("Running two-stage table detection (Method 2: structure + cell OCR)...")
        two_stage_text, two_stage_data = extract_text_two_stage(
            img_bytes, img_width, img_height, ocr_predictor_instance
        )

        # Method 3: Borderless table detection (text position clustering)
        print("Running borderless table detection (Method 3: text position analysis)...")
        borderless_text, borderless_data = extract_text_borderless(result)

        # Method 4: Block-geometry table detection (docTR block grouping)
        print("Running block-geometry table detection (Method 4: docTR block analysis)...")
        block_geo_text, block_geo_data = extract_text_block_geometry(result)

        # Determine which result to use as primary display
        # Priority: two-stage > img2table > borderless > block_geometry > plain text
        if two_stage_data.get('is_table'):
            display_text = two_stage_text
            primary_table_data = two_stage_data
            print(f"Using Two-stage: {two_stage_data.get('num_rows', 0)}x{two_stage_data.get('num_columns', 0)} table")
        elif table_data.get('is_table'):
            display_text = table_formatted_text
            primary_table_data = table_data
            print(f"Using img2table: {table_data.get('num_rows', 0)}x{table_data.get('num_columns', 0)} table")
        elif borderless_data.get('is_table'):
            display_text = borderless_text
            primary_table_data = borderless_data
            print(f"Using Borderless: {borderless_data.get('num_rows', 0)}x{borderless_data.get('num_columns', 0)} table")
        elif block_geo_data.get('is_table'):
            display_text = block_geo_text
            primary_table_data = block_geo_data
            print(f"Using Block-geometry: {block_geo_data.get('num_rows', 0)}x{block_geo_data.get('num_columns', 0)} table")
        else:
            display_text = structured_text
            primary_table_data = {'is_table': False}
            print("No table detected by any method, using regular OCR text")

        # OCR Text Correction (NER-informed dictionary pass)
        correction_enabled = enable_correction.lower() == "true"
        correction_result = {'corrected_text': cleaned_text, 'corrections': []}

        # Use cleaned text for NER input (NER correction runs after NER, see below)
        ner_input_text = cleaned_text

        # Perform NER on text
        print("Running NER...")
        entities = ner_pipeline(ner_input_text)

        # Structure entities (return all with score > 0.1, let frontend filter)
        structured_entities = []
        for entity in entities:
            if entity.get('score', 0.0) > 0.1:
                structured_entities.append({
                    'entity_group': entity['entity_group'],
                    'score': float(entity['score']),
                    'word': entity['word'].strip(),
                })

        # Map entities to bounding boxes
        entities_with_boxes = map_entities_to_boxes(structured_entities, words_with_boxes, ner_input_text)

        # NER-informed correction (second pass: fix low-confidence tokens matching entity dicts)
        if correction_enabled:
            ner_corr = correct_with_ner_entities(
                words_with_boxes, structured_entities,
                correction_result['corrected_text'], confidence_threshold=float(correction_threshold))
            if ner_corr['corrections']:
                correction_result['corrections'].extend(ner_corr['corrections'])
                correction_result['corrected_text'] = ner_corr['corrected_text']
                print(f"NER-informed correction: {len(ner_corr['corrections'])} additional fix(es)")

        # Check for drug interactions
        detected_drugs = []
        for entity in structured_entities:
            if entity['entity_group'] in ['CHEM', 'CHEMICAL', 'TREATMENT', 'MEDICATION']:
                detected_drugs.append(entity['word'])

        interactions = check_drug_interactions(detected_drugs) if detected_drugs else []
        print(f"Found {len(interactions)} drug interactions")

        # Check lab values against reference ranges
        lab_anomalies = check_lab_values(structured_text, primary_table_data, words_with_boxes)
        print(f"Found {len(lab_anomalies)} lab values ({sum(1 for a in lab_anomalies if a['status'] != 'normal')} abnormal)")

        return {
            "structured_text": display_text,  # Table-formatted if detected, otherwise regular
            "cleaned_text": cleaned_text,
            "corrected_text": correction_result['corrected_text'] if correction_enabled else None,
            "corrections": correction_result['corrections'] if correction_enabled else [],
            "medical_entities": entities_with_boxes,
            "interactions": interactions,  # Drug interaction warnings
            "lab_anomalies": lab_anomalies,  # Lab value reference range checks
            "model_id": NER_MODELS[ner_model_id]["name"],
            "ocr_model": f"{det_arch} + {reco_arch}",
            "image_width": img_width,
            "image_height": img_height,
            "synthesized_image": synthesized_image,  # Base64 PNG reconstructed from docTR
            # Primary table result (best of all methods)
            "table_detected": primary_table_data.get('is_table', False),
            "table_data": {
                "num_columns": primary_table_data.get('num_columns', 0),
                "num_rows": primary_table_data.get('num_rows', 0),
                "cells": primary_table_data.get('cells', []),
                "method": primary_table_data.get('method', 'unknown')
            } if primary_table_data.get('is_table') else None,
            # Comparison: All four methods' results
            "table_comparison": {
                "method1_img2table": {
                    "name": "img2table (line detection + integrated OCR)",
                    "detected": table_data.get('is_table', False),
                    "num_columns": table_data.get('num_columns', 0),
                    "num_rows": table_data.get('num_rows', 0),
                    "cells": table_data.get('cells', []),
                    "formatted_text": table_formatted_text if table_data.get('is_table') else None
                },
                "method2_two_stage": {
                    "name": "Two-stage (structure detection + cell-by-cell OCR)",
                    "detected": two_stage_data.get('is_table', False),
                    "num_columns": two_stage_data.get('num_columns', 0),
                    "num_rows": two_stage_data.get('num_rows', 0),
                    "cells": two_stage_data.get('cells', []),
                    "formatted_text": two_stage_text if two_stage_data.get('is_table') else None
                },
                "method3_borderless": {
                    "name": "Borderless (text position clustering)",
                    "detected": borderless_data.get('is_table', False),
                    "num_columns": borderless_data.get('num_columns', 0),
                    "num_rows": borderless_data.get('num_rows', 0),
                    "cells": borderless_data.get('cells', []),
                    "formatted_text": borderless_text if borderless_data.get('is_table') else None,
                    "fill_ratio": borderless_data.get('fill_ratio', 0)
                },
                "method4_block_geometry": {
                    "name": "Block-geometry (docTR block grouping)",
                    "detected": block_geo_data.get('is_table', False),
                    "num_columns": block_geo_data.get('num_columns', 0),
                    "num_rows": block_geo_data.get('num_rows', 0),
                    "cells": block_geo_data.get('cells', []),
                    "formatted_text": block_geo_text if block_geo_data.get('is_table') else None,
                    "fill_ratio": block_geo_data.get('fill_ratio', 0)
                }
            },
            # Docling pipeline results (side-by-side comparison)
            "docling_result": {
                "available": docling_result.get("success", False),
                "markdown_text": docling_result.get("markdown_text", ""),
                "plain_text": docling_result.get("plain_text", ""),
                "table_detected": bool(docling_result.get("tables")),
                "table_data": docling_result.get("primary_table"),
                "error": docling_result.get("error"),
            } if docling_result else {
                "available": False,
                "error": "Docling pipeline did not run",
            }
        }

    except Exception as e:
        print(f"Processing error: {e}")
        import traceback
        traceback.print_exc()
        return JSONResponse(
            status_code=500,
            content={"detail": f"An error occurred during processing: {str(e)}"}
        )