import os

# Redirect cache to a writable path inside container
os.environ["XDG_CACHE_HOME"] = "/tmp/.cache"

import gradio as gr
from impresso_pipelines.ocrqa import OCRQAPipeline

pipeline = OCRQAPipeline()

LANGUAGES = ["en", "de", "fr"]

# Example OCR text (German text with typical OCR errors)
EXAMPLE_TEXT = """Vieles Seltsame geschieht auf Erden :
Nichts Seltsameres sieht der Mond
Als das Glück, das im Knopfloch wohnt.
Zaubrisch faßt es den ernsten Mann.
Ohne nach Weib u. Kinjd zu fragen
Reitet er aus, nach dem Glück zu jagen,
Nur nacb ihm war stets sein Vegehr.
Aber neben ihm 1reitet der Dämon her
Des Ehrgeizes mit finsterer Tücke,
Und so jagt er zuletzt auf die Brücke,
Die über dem Abgrund, d:m nächtlich schwarzen
Jählings abbricht."""

def process_ocr_qa(text, lang_choice):
    try:
        lang = None if lang_choice == "Auto-detect" else lang_choice
        result = pipeline(text, language=lang, diagnostics=True)

        # Format the output for better readability
        if isinstance(result, dict):
            output_lines = []
            
            # Language detection
            if 'language' in result:
                output_lines.append(f"🌍 **Language:** {result['language']}")
            
            # Quality score
            if 'score' in result:
                score = result['score']
                score_emoji = "🟢" if score >= 0.8 else "🟡" if score >= 0.5 else "🔴"
                output_lines.append(f"{score_emoji} **Quality Score:** {score:.2f}")
            
            # Diagnostics section
            if 'diagnostics' in result and result['diagnostics']:
                diagnostics = result['diagnostics']
                output_lines.append("📊 **Detailed Analysis:**")
                
                # Model information
                if 'model_id' in diagnostics:
                    output_lines.append(f"  🤖 Model: {diagnostics['model_id']}")
                
                # Known tokens
                if 'known_tokens' in diagnostics and diagnostics['known_tokens']:
                    output_lines.append(f"  ✅ Known tokens ({len(diagnostics['known_tokens'])}): {', '.join(diagnostics['known_tokens'][:10])}")
                    if len(diagnostics['known_tokens']) > 10:
                        output_lines.append(f"      ... and {len(diagnostics['known_tokens']) - 10} more")
                
                # Unknown tokens (potential OCR errors)
                if 'unknown_tokens' in diagnostics and diagnostics['unknown_tokens']:
                    output_lines.append(f"  ❌ Potential OCR errors ({len(diagnostics['unknown_tokens'])}): {', '.join(diagnostics['unknown_tokens'])}")
                elif 'unknown_tokens' in diagnostics:
                    output_lines.append("  ✨ No potential OCR errors detected!")
            
            # Other fields
            for key, value in result.items():
                if key not in ['language', 'score', 'diagnostics']:
                    output_lines.append(f"🔍 **{key.replace('_', ' ').title()}:** {value}")
            
            return "\n\n".join(output_lines)
        else:
            return f"✨ **Processed Result:**\n{result}"
            
    except Exception as e:
        print("❌ Pipeline error:", e)
        return f"Error: {e}"

# Create the interface with logo and improved description
with gr.Blocks(title="OCR QA Demo") as demo:
    # Add logo at the top
    gr.Image("logo.jpeg", label=None, show_label=False, container=False, height=100)
    
    gr.Markdown(
        """
        # 🔍 OCR Quality Assessment Pipeline Demo
        
        **OCR Quality Assessment** demonstrates how text extracted from OCR (Optical Character Recognition) 
        is analyzed in the **Impresso** project. This pipeline identifies OCR errors and
        assesses text quality, returning a score from 0.0 to 1.0.
        
        Try the example below (German text with typical OCR errors) or enter your own OCR text to see how it gets processed!
        """
    )
    
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(
                label="Enter OCR Text", 
                value=EXAMPLE_TEXT,
                lines=8,
                placeholder="Enter your OCR text here..."
            )
            lang_dropdown = gr.Dropdown(
                choices=["Auto-detect"] + LANGUAGES, 
                value="de", 
                label="Language"
            )
            submit_btn = gr.Button("🔍 Analyze OCR Quality", variant="primary")
        
        with gr.Column():
            with gr.Row():
                output = gr.Textbox(
                    label="Analysis Results", 
                    lines=15,
                    placeholder="Results will appear here...",
                    scale=10
                )
                info_btn = gr.Button("Pipeline Info", size="sm", scale=1)
    
    # Info modal/accordion for pipeline details
    with gr.Accordion("📝 About the OCR QA Pipeline", open=False, visible=False) as info_accordion:
        gr.Markdown(
            """
            - **OCR Error Detection**: Identifies common OCR mistakes and artifacts
            - **Quality Assessment**: Evaluates the overall quality of OCR text
            - **Language Processing**: Handles multilingual OCR text processing
            """
        )
    
    submit_btn.click(
        fn=process_ocr_qa,
        inputs=[text_input, lang_dropdown],
        outputs=output
    )
    
    # Toggle info visibility when info button is clicked
    info_btn.click(
        fn=lambda: gr.Accordion(visible=True, open=True),
        outputs=info_accordion
    )

demo.launch(server_name="0.0.0.0", server_port=7860)