import gradio as gr
import cv2
import numpy as np
import aiohttp
import asyncio
from stripeRemover import StripeRemover

class OCRUI:
    def __init__(self):
        self.API_URL = "http://s15.serv00.com:9081/compareAnalyze"
        self.stripe_remover = StripeRemover()

    def process_image(self, image, method):
        if image is None:
            return None
            
        try:
            if method == "Original":
                return cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
            elif method == "Fourier":
                return self.stripe_remover.fourier_method(image)
            elif method == "Morphological":
                return self.stripe_remover.morphological_method(image)
            elif method == "Adaptive":
                return self.stripe_remover.adaptive_threshold_method(image)
            elif method == "Enhanced":
                gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
                return self.stripe_remover.enhance_image(gray)
        except Exception as e:
            return None

    async def send_to_api(self, image_bytes):
        async with aiohttp.ClientSession() as session:
            data = aiohttp.FormData()
            data.add_field('image', image_bytes, 
                          filename='image.jpg',
                          content_type='image/jpeg')
            data.add_field('model', 'GEMINI')
            
            try:
                async with session.post(self.API_URL, data=data) as response:
                    return await response.json()
            except Exception as e:
                return {"error": str(e)}

    def ocr_process(self, image, preprocess_method):
        if image is None:
            return None, "Please upload an image"

        # Preprocess image
        processed_img = self.process_image(image, preprocess_method)
        if processed_img is None:
            return None, "Image processing failed"

        # Prepare image for API
        encode_params = [cv2.IMWRITE_JPEG_QUALITY, 50]
        _, img_bytes = cv2.imencode('.jpg', processed_img, encode_params)
        
        # Call API
        result = asyncio.run(self.send_to_api(img_bytes.tobytes()))
        
        if result is None or "error" in result:
            return None, "API call failed"
            
        return result, None  # Return raw JSON for gradio.JSON component
    
    
    def preview_process(self, image, method):
        """Preview processed image without OCR"""
        if image is None:
            return None
        return self.process_image(image, method)

def create_ui():
    ui = OCRUI()
    
    with gr.Blocks() as demo:
        gr.Markdown("# 美宜佳DEMO")
        
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(
                    type="numpy",
                    label="Input Image",
                    height=200,  # Smaller initial height
                    show_download_button=False,
                    interactive=True,
                    container=True,
                )
                preprocess_dropdown = gr.Dropdown(
                    choices=["Original", "Fourier", "Morphological", "Adaptive", "Enhanced"],
                    label="Preprocessing Method",
                    value="Original"
                )
                process_btn = gr.Button("Process Image")

            with gr.Column():
                preview_output = gr.Image(
                    type="numpy",
                    label="Processed Preview",
                    height=200,
                    show_download_button=False,
                    interactive=False,
                )
                json_output = gr.JSON(label="OCR Results")  # Replace text_output
                status_output = gr.Text(label="Status")     # For error messages

        # Add preview update on image or method change
        image_input.change(
            fn=ui.preview_process,
            inputs=[image_input, preprocess_dropdown],
            outputs=preview_output
        )
        preprocess_dropdown.change(
            fn=ui.preview_process,
            inputs=[image_input, preprocess_dropdown],
            outputs=preview_output
        )

        process_btn.click(
            fn=ui.ocr_process,
            inputs=[image_input, preprocess_dropdown],
            outputs=[json_output, status_output]
        )
    
    return demo

if __name__ == "__main__":
    demo = create_ui()
    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)