import gradio as gr
from modules.element_processing import process_screenshot
from modules.element_detector import initialize_models
import asyncio
import io
import os
import base64
import numpy as np
from PIL import Image

# Initialize models on startup
async def init():
    await initialize_models()

loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
loop.run_until_complete(init())

# Create a mock background tasks class
class MockBackgroundTasks:
    def add_task(self, *args, **kwargs):
        pass

# Define interface for data URI processing
def process_image_data_uri(data_uri):
    if not data_uri.startswith('data:image'):
        return None, "Invalid data URI format"
    
    try:
        # Split the header and the base64 data
        header, encoded = data_uri.split(",", 1)
        image_data = base64.b64decode(encoded)
        
        # Process the screenshot
        background_tasks = MockBackgroundTasks()
        elements, image_path = loop.run_until_complete(
            process_screenshot(image_data, background_tasks)
        )
        
        # Format output
        result_lines = []
        for element in elements:
            # Determine content based on element type
            if element["type"] == "text":
                content_value = element.get("text_content", "")
            else:
                content_value = element.get("object_label", "")
            
            # Format the element string
            element_str = (f"icon {element['code']}: {{"
                          f"'type': '{element['type']}', "
                          f"'centerX': {element['center_x']}, "
                          f"'centerY': {element['center_y']}, "
                          f"'content': '{content_value}'}}")
            
            result_lines.append(element_str)
        
        # Join all lines
        result_string = "\n".join(result_lines)
        
        # Load annotated image
        annotated_img = None
        if image_path and os.path.exists(image_path):
            annotated_img = Image.open(image_path)
        
        return annotated_img, result_string
        
    except Exception as e:
        return None, f"Error: {str(e)}"

# Also allow uploading an image directly
def process_image_upload(image):
    if image is None:
        return None, "No image provided"
    
    try:
        # Convert image to bytes
        img_byte_arr = io.BytesIO()
        image.save(img_byte_arr, format='PNG')
        image_data = img_byte_arr.getvalue()
        
        # Process the screenshot
        background_tasks = MockBackgroundTasks()
        elements, image_path = loop.run_until_complete(
            process_screenshot(image_data, background_tasks)
        )
        
        # Format output
        result_lines = []
        for element in elements:
            # Determine content based on element type
            if element["type"] == "text":
                content_value = element.get("text_content", "")
            else:
                content_value = element.get("object_label", "")
            
            # Format the element string
            element_str = (f"icon {element['code']}: {{"
                          f"'type': '{element['type']}', "
                          f"'centerX': {element['center_x']}, "
                          f"'centerY': {element['center_y']}, "
                          f"'content': '{content_value}'}}")
            
            result_lines.append(element_str)
        
        # Join all lines
        result_string = "\n".join(result_lines)
        
        # Load annotated image
        annotated_img = None
        if image_path and os.path.exists(image_path):
            annotated_img = Image.open(image_path)
        
        return annotated_img, result_string
        
    except Exception as e:
        return None, f"Error: {str(e)}"

# Create tabbed interface
with gr.Blocks(title="UI Element Detection") as demo:
    gr.Markdown("# UI Element Detection")
    gr.Markdown("Upload a screenshot or provide a data URI to detect UI elements with bounding boxes and text annotations")
    
    with gr.Tabs():
        with gr.TabItem("Upload Image"):
            with gr.Row():
                with gr.Column():
                    image_input = gr.Image(type="pil", label="Upload Screenshot")
                    upload_button = gr.Button("Process Image")
                with gr.Column():
                    image_output = gr.Image(type="pil", label="Annotated Image")
                    text_output = gr.Textbox(label="Detected Elements", lines=10)
            
            upload_button.click(
                fn=process_image_upload,
                inputs=image_input,
                outputs=[image_output, text_output]
            )
        
        with gr.TabItem("Data URI"):
            with gr.Row():
                with gr.Column():
                    data_uri_input = gr.Textbox(label="Image Data URI", lines=5, placeholder="data:image/png;base64,...")
                    data_uri_button = gr.Button("Process Data URI")
                with gr.Column():
                    data_uri_image_output = gr.Image(type="pil", label="Annotated Image")
                    data_uri_text_output = gr.Textbox(label="Detected Elements", lines=10)
            
            data_uri_button.click(
                fn=process_image_data_uri,
                inputs=data_uri_input,
                outputs=[data_uri_image_output, data_uri_text_output]
            )

# Ensure annotated directory exists
os.makedirs("annotated", exist_ok=True)

# Launch the app
demo.launch(share=True, enable_api=True)