"""
Surya OCR Studio - Complete Implementation
Features: OCR, Text Detection, Layout Analysis, Table Recognition, LaTeX OCR
"""

# Import spaces FIRST before any CUDA-related packages
import spaces
import gradio as gr
import logging
import os
import json
from PIL import Image, ImageDraw, ImageFont
from typing import List, Optional
import torch

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Performance optimizations for ZeroGPU
os.environ["RECOGNITION_BATCH_SIZE"] = "64"
os.environ["DETECTOR_BATCH_SIZE"] = "8"
os.environ["LAYOUT_BATCH_SIZE"] = "8"
os.environ["TABLE_REC_BATCH_SIZE"] = "16"

# Surya imports
from surya.foundation import FoundationPredictor
from surya.recognition import RecognitionPredictor
from surya.detection import DetectionPredictor
from surya.layout import LayoutPredictor
from surya.table_rec import TableRecPredictor
from surya.texify import TexifyPredictor
from surya.settings import settings

logger.info("Loading Surya models...")

# Initialize predictors (lazy loading for faster startup)
_foundation_predictor = None
_detection_predictor = None
_recognition_predictor = None
_layout_predictor = None
_table_rec_predictor = None
_texify_predictor = None

def get_foundation_predictor():
    global _foundation_predictor
    if _foundation_predictor is None:
        _foundation_predictor = FoundationPredictor()
    return _foundation_predictor

def get_detection_predictor():
    global _detection_predictor
    if _detection_predictor is None:
        _detection_predictor = DetectionPredictor()
    return _detection_predictor

def get_recognition_predictor():
    global _recognition_predictor
    if _recognition_predictor is None:
        _recognition_predictor = RecognitionPredictor(get_foundation_predictor())
    return _recognition_predictor

def get_layout_predictor():
    global _layout_predictor
    if _layout_predictor is None:
        _layout_predictor = LayoutPredictor(
            FoundationPredictor(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
        )
    return _layout_predictor

def get_table_rec_predictor():
    global _table_rec_predictor
    if _table_rec_predictor is None:
        _table_rec_predictor = TableRecPredictor()
    return _table_rec_predictor

def get_texify_predictor():
    global _texify_predictor
    if _texify_predictor is None:
        _texify_predictor = TexifyPredictor()
    return _texify_predictor

logger.info("Models will be loaded on first use.")

# Layout labels and colors
LAYOUT_LABELS = {
    'Text': '#10B981',        # Green
    'Title': '#EF4444',       # Red
    'Section-header': '#F59E0B',  # Amber
    'Table': '#3B82F6',       # Blue
    'Figure': '#8B5CF6',      # Purple
    'Picture': '#8B5CF6',     # Purple
    'Caption': '#EC4899',     # Pink
    'Page-header': '#6366F1', # Indigo
    'Page-footer': '#6366F1', # Indigo
    'Footnote': '#84CC16',    # Lime
    'Formula': '#F97316',     # Orange
    'List-item': '#14B8A6',   # Teal
    'Form': '#A855F7',        # Fuchsia
    'Handwriting': '#64748B', # Slate
    'Table-of-contents': '#0EA5E9',  # Sky
}

# Supported languages
LANGUAGES = {
    "en": "English", "pt": "Portuguese", "es": "Spanish", "fr": "French",
    "de": "German", "it": "Italian", "nl": "Dutch", "ru": "Russian",
    "zh": "Chinese", "ja": "Japanese", "ko": "Korean", "ar": "Arabic",
    "hi": "Hindi", "bn": "Bengali", "tr": "Turkish", "vi": "Vietnamese",
    "th": "Thai", "id": "Indonesian", "pl": "Polish", "uk": "Ukrainian",
    "cs": "Czech", "sv": "Swedish", "da": "Danish", "no": "Norwegian",
    "fi": "Finnish", "el": "Greek", "he": "Hebrew", "hu": "Hungarian",
    "ro": "Romanian", "sk": "Slovak", "bg": "Bulgarian", "hr": "Croatian",
    "sl": "Slovenian", "et": "Estonian", "lv": "Latvian", "lt": "Lithuanian",
    "fa": "Persian", "ur": "Urdu", "ta": "Tamil", "te": "Telugu",
    "ml": "Malayalam", "kn": "Kannada", "gu": "Gujarati", "mr": "Marathi",
    "pa": "Punjabi", "ne": "Nepali", "si": "Sinhala", "my": "Burmese",
    "km": "Khmer", "lo": "Lao", "ka": "Georgian", "hy": "Armenian",
}


def prepare_image(image) -> Image.Image:
    """Prepare image for processing"""
    if isinstance(image, str):
        image = Image.open(image)
    elif hasattr(image, 'name'):
        image = Image.open(image.name)
    
    if image.mode != 'RGB':
        image = image.convert('RGB')
    return image


def draw_text_lines(image, text_lines, color=(0, 255, 0)):
    """Draw text line bounding boxes"""
    draw = ImageDraw.Draw(image)
    for line in text_lines:
        if hasattr(line, 'bbox'):
            bbox = line.bbox
            if len(bbox) == 4:
                draw.rectangle(bbox, outline=color, width=2)
    return image


def draw_layout_boxes(image, bboxes):
    """Draw layout boxes with labels and colors"""
    draw = ImageDraw.Draw(image)
    
    for bbox in bboxes:
        label = getattr(bbox, 'label', 'Text')
        color = LAYOUT_LABELS.get(label, '#FFFFFF')
        
        # Convert hex to RGB
        rgb = tuple(int(color.lstrip('#')[i:i+2], 16) for i in (0, 2, 4))
        
        box = bbox.bbox if hasattr(bbox, 'bbox') else bbox
        if len(box) == 4:
            draw.rectangle(box, outline=rgb, width=2)
            
            # Draw label
            label_text = label.replace('-', ' ').title()
            draw.text((box[0], box[1] - 12), label_text, fill=rgb)
    
    return image


def draw_table_cells(image, predictions):
    """Draw table cells with row/column info"""
    draw = ImageDraw.Draw(image)
    
    # Draw rows in blue
    for row in predictions.rows:
        draw.rectangle(row.bbox, outline=(0, 0, 255), width=2)
    
    # Draw columns in green
    for col in predictions.cols:
        draw.rectangle(col.bbox, outline=(0, 255, 0), width=2)
    
    # Draw cells in red
    for cell in predictions.cells:
        draw.rectangle(cell.bbox, outline=(255, 0, 0), width=1)
        # Draw cell text if available
        if hasattr(cell, 'text') and cell.text:
            draw.text((cell.bbox[0], cell.bbox[1]), cell.text[:10], fill=(100, 100, 100))
    
    return image


@spaces.GPU(duration=120)
def process_ocr(image, languages: str, disable_math: bool = False):
    """Run OCR on image"""
    logger.info(f"Running OCR with languages: {languages}")
    try:
        image = prepare_image(image)
        
        # Parse languages
        langs = [l.strip() for l in languages.split(',') if l.strip()]
        if not langs:
            langs = ['en']
        
        # Get predictors
        det_pred = get_detection_predictor()
        rec_pred = get_recognition_predictor()
        
        # Run OCR
        predictions = rec_pred(
            [image], 
            det_predictor=det_pred, 
            langs=[langs],
            disable_math=disable_math
        )
        
        if not predictions or len(predictions) == 0:
            return "", {}, None
        
        pred = predictions[0]
        
        # Extract text
        text_lines = pred.text_lines if hasattr(pred, 'text_lines') else []
        full_text = "\n".join([line.text for line in text_lines if hasattr(line, 'text')])
        
        # Build JSON result
        result = {
            "text": full_text,
            "languages": langs,
            "num_lines": len(text_lines),
            "lines": [
                {
                    "text": line.text,
                    "confidence": round(line.confidence, 3) if hasattr(line, 'confidence') else 1.0,
                    "bbox": list(line.bbox) if hasattr(line, 'bbox') else []
                }
                for line in text_lines
            ]
        }
        
        # Draw bounding boxes
        img_with_boxes = draw_text_lines(image.copy(), text_lines)
        
        return full_text, result, img_with_boxes
        
    except Exception as e:
        logger.error(f"OCR Error: {e}", exc_info=True)
        return f"Error: {str(e)}", {"error": str(e)}, None


@spaces.GPU(duration=60)
def process_detection(image):
    """Run text detection on image"""
    logger.info("Running text detection")
    try:
        image = prepare_image(image)
        
        det_pred = get_detection_predictor()
        predictions = det_pred([image])
        
        if not predictions or len(predictions) == 0:
            return {}, None
        
        pred = predictions[0]
        
        # Build result
        result = {
            "num_lines": len(pred.bboxes) if hasattr(pred, 'bboxes') else 0,
            "image_size": list(image.size),
            "bboxes": [
                {
                    "bbox": list(bbox.bbox) if hasattr(bbox, 'bbox') else list(bbox),
                    "confidence": round(bbox.confidence, 3) if hasattr(bbox, 'confidence') else 1.0
                }
                for bbox in (pred.bboxes if hasattr(pred, 'bboxes') else [])
            ]
        }
        
        # Draw boxes
        img_with_boxes = image.copy()
        draw = ImageDraw.Draw(img_with_boxes)
        for bbox in (pred.bboxes if hasattr(pred, 'bboxes') else []):
            box = bbox.bbox if hasattr(bbox, 'bbox') else bbox
            draw.rectangle(box, outline=(0, 255, 0), width=2)
        
        return result, img_with_boxes
        
    except Exception as e:
        logger.error(f"Detection Error: {e}", exc_info=True)
        return {"error": str(e)}, None


@spaces.GPU(duration=60)
def process_layout(image):
    """Run layout analysis on image"""
    logger.info("Running layout analysis")
    try:
        image = prepare_image(image)
        
        layout_pred = get_layout_predictor()
        predictions = layout_pred([image])
        
        if not predictions or len(predictions) == 0:
            return {}, None
        
        pred = predictions[0]
        
        # Count by label
        label_counts = {}
        for bbox in (pred.bboxes if hasattr(pred, 'bboxes') else []):
            label = getattr(bbox, 'label', 'Unknown')
            label_counts[label] = label_counts.get(label, 0) + 1
        
        # Build result
        result = {
            "num_elements": len(pred.bboxes) if hasattr(pred, 'bboxes') else 0,
            "label_counts": label_counts,
            "elements": [
                {
                    "label": getattr(bbox, 'label', 'Unknown'),
                    "confidence": round(getattr(bbox, 'confidence', 1.0), 3),
                    "position": getattr(bbox, 'position', -1),
                    "bbox": list(bbox.bbox) if hasattr(bbox, 'bbox') else []
                }
                for bbox in (pred.bboxes if hasattr(pred, 'bboxes') else [])
            ]
        }
        
        # Draw layout boxes
        img_with_boxes = draw_layout_boxes(image.copy(), pred.bboxes if hasattr(pred, 'bboxes') else [])
        
        return result, img_with_boxes
        
    except Exception as e:
        logger.error(f"Layout Error: {e}", exc_info=True)
        return {"error": str(e)}, None


@spaces.GPU(duration=90)
def process_table(image):
    """Run table recognition on image"""
    logger.info("Running table recognition")
    try:
        image = prepare_image(image)
        
        table_pred = get_table_rec_predictor()
        predictions = table_pred([image])
        
        if not predictions or len(predictions) == 0:
            return {}, None, ""
        
        pred = predictions[0]
        
        # Build markdown table
        md_table = ""
        if hasattr(pred, 'cells') and pred.cells:
            # Find max row and col
            max_row = max(c.row_id for c in pred.cells) if pred.cells else 0
            max_col = max(c.col_id for c in pred.cells) if pred.cells else 0
            
            # Create table data
            table_data = [["" for _ in range(max_col + 1)] for _ in range(max_row + 1)]
            
            for cell in pred.cells:
                text = getattr(cell, 'text', '')
                table_data[cell.row_id][cell.col_id] = text
            
            # Build markdown
            md_lines = []
            for i, row in enumerate(table_data):
                md_lines.append("| " + " | ".join(row) + " |")
                if i == 0:
                    md_lines.append("| " + " | ".join(["---"] * len(row)) + " |")
            
            md_table = "\n".join(md_lines)
        
        # Build result
        result = {
            "num_rows": len(pred.rows) if hasattr(pred, 'rows') else 0,
            "num_cols": len(pred.cols) if hasattr(pred, 'cols') else 0,
            "num_cells": len(pred.cells) if hasattr(pred, 'cells') else 0,
            "rows": [
                {"row_id": r.row_id, "is_header": getattr(r, 'is_header', False)}
                for r in (pred.rows if hasattr(pred, 'rows') else [])
            ],
            "cols": [
                {"col_id": c.col_id, "is_header": getattr(c, 'is_header', False)}
                for c in (pred.cols if hasattr(pred, 'cols') else [])
            ]
        }
        
        # Draw table cells
        img_with_boxes = draw_table_cells(image.copy(), pred)
        
        return result, img_with_boxes, md_table
        
    except Exception as e:
        logger.error(f"Table Recognition Error: {e}", exc_info=True)
        return {"error": str(e)}, None, ""


@spaces.GPU(duration=60)
def process_latex(image):
    """Run LaTeX OCR on image (equation)"""
    logger.info("Running LaTeX OCR")
    try:
        image = prepare_image(image)
        
        texify_pred = get_texify_predictor()
        predictions = texify_pred([image])
        
        if not predictions or len(predictions) == 0:
            return "", {}
        
        pred = predictions[0]
        latex = pred.text if hasattr(pred, 'text') else str(pred)
        
        result = {
            "latex": latex,
            "markdown": f"$$\n{latex}\n$$"
        }
        
        return latex, result
        
    except Exception as e:
        logger.error(f"LaTeX OCR Error: {e}", exc_info=True)
        return f"Error: {str(e)}", {"error": str(e)}


# ============== GRADIO UI ==============

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;500;600;700;800&family=Fira+Code:wght@400;500&display=swap');

:root {
    --bg: #0a0f1a;
    --surf: #0f1629;
    --card: #151d32;
    --border: #1e2a45;
    --border2: #2a3a5a;
    --green: #10b981;
    --blue: #3b82f6;
    --text: #e2e8f0;
    --muted: #64748b;
}

body, .gradio-container {
    background: var(--bg) !important;
    font-family: 'Outfit', sans-serif !important;
    color: var(--text) !important;
}

.gradio-container::before {
    content: '';
    position: fixed; inset: 0; pointer-events: none; z-index: 0;
    background: radial-gradient(ellipse 70% 50% at 50% -10%, rgba(16,185,129,0.08) 0%, transparent 65%);
}

.app-hero { padding: 40px 0 20px; text-align: center; }
.app-hero h1 {
    font-size: 2.8rem; font-weight: 800; letter-spacing: -0.04em;
    background: linear-gradient(135deg, #10b981, #06b6d4, #3b82f6);
    -webkit-background-clip: text; -webkit-text-fill-color: transparent;
    margin-bottom: 8px;
}
.app-hero .tagline { color: var(--muted); font-size: 1rem; }
.pills { display: flex; justify-content: center; gap: 8px; margin-top: 16px; flex-wrap: wrap; }
.pill {
    background: var(--card); border: 1px solid var(--border2); border-radius: 100px;
    padding: 5px 14px; font-size: 0.75rem; color: var(--muted); font-family: 'Fira Code', monospace;
}
.pill.green { color: var(--green); border-color: rgba(16,185,129,0.3); }

.tabs button { font-family: 'Outfit', sans-serif !important; font-weight: 500 !important; }

button.primary-btn { 
    background: linear-gradient(135deg, #10b981, #06b6d4) !important; 
    border: none !important; color: #000 !important; font-weight: 600 !important;
    padding: 12px 24px !important;
}

.output-image img { border-radius: 8px; }

footer { display: none !important; }
"""

# Language string for the dropdown
LANGUAGE_OPTIONS = [f"{code} - {name}" for code, name in sorted(LANGUAGES.items(), key=lambda x: x[1])]

with gr.Blocks(theme=gr.themes.Base(), css=CSS, title="Surya OCR Studio") as app:
    
    gr.HTML("""
    <div class="app-hero">
        <h1>📄 Surya OCR Studio</h1>
        <p class="tagline">Document OCR · Layout Analysis · Table Recognition · 90+ Languages</p>
        <div class="pills">
            <span class="pill green">ZeroGPU ⚡</span>
            <span class="pill">90+ Languages</span>
            <span class="pill">Layout Detection</span>
            <span class="pill">Table Recognition</span>
            <span class="pill">LaTeX OCR</span>
        </div>
    </div>
    """)
    
    with gr.Tabs() as tabs:
        
        # ============ OCR TAB ============
        with gr.TabItem("📝 OCR"):
            gr.Markdown("### Optical Character Recognition\nExtract text from images in 90+ languages.")
            
            with gr.Row():
                with gr.Column(scale=2):
                    ocr_input = gr.Image(label="📄 Upload Image", type="pil", height=400)
                    
                    with gr.Row():
                        ocr_langs = gr.Textbox(
                            label="Languages (comma-separated codes)",
                            value="en",
                            placeholder="en, pt, es, de, fr...",
                            info="Use language codes: en=English, pt=Portuguese, es=Spanish..."
                        )
                        ocr_disable_math = gr.Checkbox(label="Disable math detection", value=False)
                    
                    ocr_btn = gr.Button("🚀 Run OCR", variant="primary", elem_classes=["primary-btn"])
                
                with gr.Column(scale=3):
                    ocr_text = gr.Textbox(label="📝 Extracted Text", lines=12, show_copy_button=True)
                    ocr_json = gr.JSON(label="📊 Detailed Results")
            
            ocr_image = gr.Image(label="🖼️ Detected Text Lines", elem_classes=["output-image"])
            
            ocr_btn.click(
                process_ocr,
                [ocr_input, ocr_langs, ocr_disable_math],
                [ocr_text, ocr_json, ocr_image]
            )
        
        # ============ TEXT DETECTION TAB ============
        with gr.TabItem("🔍 Text Detection"):
            gr.Markdown("### Text Line Detection\nDetect text lines in documents without OCR.")
            
            with gr.Row():
                det_input = gr.Image(label="📄 Upload Image", type="pil", height=400)
                with gr.Column():
                    det_json = gr.JSON(label="📊 Detection Results")
                    det_btn = gr.Button("🔍 Detect Text Lines", variant="primary", elem_classes=["primary-btn"])
            
            det_image = gr.Image(label="🖼️ Detected Lines", elem_classes=["output-image"])
            
            det_btn.click(process_detection, [det_input], [det_json, det_image])
        
        # ============ LAYOUT ANALYSIS TAB ============
        with gr.TabItem("📊 Layout Analysis"):
            gr.Markdown("### Document Layout Analysis\nIdentify document structure: titles, tables, figures, etc.")
            
            with gr.Row():
                layout_input = gr.Image(label="📄 Upload Image", type="pil", height=400)
                with gr.Column():
                    layout_json = gr.JSON(label="📊 Layout Results")
                    layout_btn = gr.Button("📊 Analyze Layout", variant="primary", elem_classes=["primary-btn"])
            
            layout_image = gr.Image(label="🖼️ Layout Elements", elem_classes=["output-image"])
            
            # Legend
            gr.Markdown("""
            **Legend:** 
            🟢 Text | 🔴 Title | 🟡 Section Header | 🔵 Table | 🟣 Figure/Picture | 🩷 Caption | 🔷 Header/Footer
            """)
            
            layout_btn.click(process_layout, [layout_input], [layout_json, layout_image])
        
        # ============ TABLE RECOGNITION TAB ============
        with gr.TabItem("📋 Table Recognition"):
            gr.Markdown("### Table Recognition\nExtract table structure and convert to Markdown.")
            
            with gr.Row():
                table_input = gr.Image(label="📄 Upload Table Image", type="pil", height=400)
                with gr.Column():
                    table_json = gr.JSON(label="📊 Table Structure")
                    table_btn = gr.Button("📋 Recognize Table", variant="primary", elem_classes=["primary-btn"])
            
            table_image = gr.Image(label="🖼️ Table Cells", elem_classes=["output-image"])
            table_md = gr.Textbox(label="📝 Markdown Output", lines=10, show_copy_button=True)
            
            table_btn.click(process_table, [table_input], [table_json, table_image, table_md])
        
        # ============ LATEX OCR TAB ============
        with gr.TabItem("🔢 LaTeX OCR"):
            gr.Markdown("### LaTeX Equation OCR\nConvert equation images to LaTeX code.\n\n**Tip:** Crop the image to just the equation for best results.")
            
            with gr.Row():
                latex_input = gr.Image(label="📄 Upload Equation Image", type="pil", height=300)
                with gr.Column():
                    latex_code = gr.Textbox(label="🔢 LaTeX Code", lines=5, show_copy_button=True)
                    latex_json = gr.JSON(label="📊 Results")
                    latex_btn = gr.Button("🔢 Extract LaTeX", variant="primary", elem_classes=["primary-btn"])
            
            latex_btn.click(process_latex, [latex_input], [latex_code, latex_json])
    
    # ============ FOOTER ============
    gr.Markdown("""
    ---
    
    ### ℹ️ About Surya OCR
    
    | Feature | Description |
    |---------|-------------|
    | **OCR** | Text recognition in 90+ languages |
    | **Detection** | Line-level text detection |
    | **Layout** | Identify tables, figures, headers, etc. |
    | **Tables** | Extract table structure to Markdown |
    | **LaTeX** | Convert equations to LaTeX |
    
    **Performance Tips:**
    - Use higher resolution images for better accuracy
    - For blurry text, try preprocessing (binarization, deskewing)
    - Specify correct language codes for best OCR results
    
    **Model:** [Surya OCR](https://github.com/datalab-to/surya) by datalab-to  
    **Space by:** [@artificialguybr](https://twitter.com/artificialguybr)
    """)


if __name__ == "__main__":
    app.queue(max_size=20)
    app.launch()