import gradio as gr
import logging
import os
import json
from PIL import Image, ImageDraw
import torch
from surya.ocr import run_ocr
from surya.detection import batch_text_detection
from surya.layout import batch_layout_detection
from surya.ordering import batch_ordering
from surya.model.detection.model import load_model as load_det_model, load_processor as load_det_processor
from surya.model.recognition.model import load_model as load_rec_model
from surya.model.recognition.processor import load_processor as load_rec_processor
from surya.settings import settings

# Configuração de logging
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Carregamento de modelos
det_processor, det_model = load_det_processor(), load_det_model()
rec_model, rec_processor = load_rec_model(), load_rec_processor()

class CustomJSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, Image.Image):
            return "Image object (not serializable)"
        return str(obj)

def serialize_result(result):
    return json.dumps(result, cls=CustomJSONEncoder, indent=2)

def save_metadata(results):
    output_file = "/mnt/data/ocr_metadata.json"  # Caminho de armazenamento persistente
    with open(output_file, "w") as f:
        json.dump(results, f, cls=CustomJSONEncoder, indent=2)
    return output_file

def ocr_workflow(images, langs):
    logger.info(f"Iniciando workflow OCR para {len(images)} imagens com idiomas: {langs}")
    results = []
    for image_file in images:
        try:
            image = Image.open(image_file.name)
            predictions = run_ocr([image], [langs.split(',')], det_model, det_processor, rec_model, rec_processor)
            formatted_text = "\n".join([line.text for line in predictions[0].text_lines])
            results.append({
                "image": image_file.name,
                "text": formatted_text,
                "predictions": predictions[0].text_lines  # Assuming text_lines is serializable
            })
        except Exception as e:
            logger.error(f"Erro com {image_file.name}: {e}")
            results.append({"image": image_file.name, "error": str(e)})
    
    metadata_file = save_metadata(results)
    return serialize_result(results), metadata_file

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Análise de Documentos com Surya")
    
    with gr.Tab("OCR"):
        gr.Markdown("## Reconhecimento Óptico de Caracteres para Múltiplas Imagens")
        with gr.Row():
            ocr_input = gr.Files(label="Carregar Imagens ou PDFs (até 100)")
            ocr_langs = gr.Textbox(label="Idiomas (separados por vírgula)", value="en")
        ocr_button = gr.Button("Executar OCR")
        ocr_output = gr.JSON(label="Resultados OCR")
        ocr_file = gr.File(label="Baixar Metadata OCR")
        
        # Executa a função OCR e salva o arquivo de metadata para download
        ocr_button.click(ocr_workflow, inputs=[ocr_input, ocr_langs], outputs=[ocr_output, ocr_file])

if __name__ == "__main__":
    logger.info("Iniciando aplicativo Gradio...")
    demo.launch()