import gradio as gr
import os
import base64
import time
import json
import logging
import tempfile
import uuid
import io

from PIL import Image
from openai import OpenAI
from ultralytics import YOLO

from wrapper import process_image_description
from utils.pills import preprocess_image
import cv2
import cv2.dnn_superres as dnn_superres
import easyocr
from spellchecker import SpellChecker

logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')

GLOBAL_SR = None
GLOBAL_READER = None
GLOBAL_SPELL = None
YOLO_MODEL = None

def load_models():
    """
    Called once to load all necessary models into memory.
    """
    global GLOBAL_SR, GLOBAL_READER, GLOBAL_SPELL, YOLO_MODEL

    logging.info("Loading all models...")
    start_time_total = time.perf_counter()

    # Super-resolution
    logging.info("Loading super-resolution model...")
    start_time = time.perf_counter()
    sr = None
    model_path = "EDSR_x4.pb"
    if os.path.exists(model_path):
        if hasattr(cv2, 'dnn_superres'):
            try:
                sr = dnn_superres.DnnSuperResImpl_create()
            except AttributeError:
                sr = dnn_superres.DnnSuperResImpl()
            sr.readModel(model_path)
            sr.setModel('edsr', 4)
            GLOBAL_SR = sr
            logging.info("Super-resolution model loaded.")
        else:
            logging.warning("cv2.dnn_superres module not available.")
    else:
        logging.warning(f"Super-resolution model file not found: {model_path}. Skipping SR.")
    logging.info(f"Super-resolution init took {time.perf_counter()-start_time:.3f}s.")

    # EasyOCR + SpellChecker
    logging.info("Loading OCR + SpellChecker...")
    start_time = time.perf_counter()
    GLOBAL_READER = easyocr.Reader(['en'], gpu=True)
    GLOBAL_SPELL = SpellChecker()
    logging.info(f"OCR + SpellChecker init took {time.perf_counter()-start_time:.3f}s.")

    # YOLO Model
    logging.info("Loading YOLO model...")
    start_time = time.perf_counter()
    yolo_weights = "best.pt"
    if os.path.exists(yolo_weights):
        YOLO_MODEL = YOLO(yolo_weights)
        logging.info("YOLO model loaded.")
    else:
        logging.error(f"YOLO weights file '{yolo_weights}' not found! Endpoints will fail.")
    logging.info(f"YOLO init took {time.perf_counter()-start_time:.3f}s.")
    
    logging.info(f"Total model loading time: {time.perf_counter()-start_time_total:.3f}s.")


def pil_to_base64_str(pil_image, format="PNG"):
    """Converts a PIL Image to a base64 string with a data URI header."""
    buffered = io.BytesIO()
    pil_image.save(buffered, format=format)
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    return f"data:image/{format.lower()};base64,{img_str}"

def save_base64_image(image_data: str, file_path: str):
    """Saves a base64 encoded image to a file."""
    if image_data.startswith("data:image"):
        _, image_data = image_data.split(",", 1)
    img_bytes = base64.b64decode(image_data)
    with open(file_path, "wb") as f:
        f.write(img_bytes)
    return img_bytes

def run_wrapper(image_path: str, output_dir: str, skip_ocr: bool = False, skip_spell: bool = False, json_mini=False) -> str:
    """Calls the main processing script and returns the result."""
    process_image_description(
        input_image=image_path,
        weights_file="best.pt",
        output_dir=output_dir,
        no_captioning=True,
        output_json=True,
        json_mini=json_mini,
        model_obj=YOLO_MODEL,
        sr=GLOBAL_SR,
        spell=None if skip_ocr else GLOBAL_SPELL,
        reader=None if skip_ocr else GLOBAL_READER,
        skip_ocr=skip_ocr,
        skip_spell=skip_spell,
    )
    base_name = os.path.splitext(os.path.basename(image_path))[0]
    result_dir = os.path.join(output_dir, "result")
    json_file = os.path.join(result_dir, f"{base_name}.json")
    if os.path.exists(json_file):
        with open(json_file, "r", encoding="utf-8") as f:
            return f.read()
    else:
        raise FileNotFoundError(f"Result file not generated: {json_file}")

def handle_action(openai_key, image, prompt):
    if not openai_key: return "Error: OpenAI API Key is required for /action."
    if image is None: return "Error: Please upload an image."
    if not prompt: return "Error: Please provide a prompt."

    try:
        llm_client = OpenAI(api_key=openai_key)
        image_b64 = pil_to_base64_str(image)
        
        with tempfile.TemporaryDirectory() as temp_dir:
            request_id = str(uuid.uuid4())
            original_image_path = os.path.join(temp_dir, f"{request_id}.png")
            yolo_updated_image_path = os.path.join(temp_dir, f"{request_id}_yolo_updated.png")
            save_base64_image(image_b64, original_image_path)
            
            image_description = run_wrapper(original_image_path, temp_dir, skip_ocr=False, skip_spell=True, json_mini=True)
            
            with open(yolo_updated_image_path, "rb") as f:
                yolo_updated_img_bytes = f.read()
            
            _, new_b64 = preprocess_image(yolo_updated_img_bytes, threshold=2000, scale=0.5, fmt="png")

        base64_image_url = f"data:image/png;base64,{new_b64}"
        prompt_text = f"""You are an AI agent... (rest of your long prompt)
        The user said: "{prompt}"
        Description: "{image_description}" """
        
        messages = [{"role": "user", "content": [{"type": "text", "text": prompt_text}, {"type": "image_url", "image_url": {"url": base64_image_url, "detail": "high"}}]}]
        
        response = llm_client.chat.completions.create(model="gpt-4.1", messages=messages, temperature=0.2)
        return response.choices[0].message.content.strip()

    except Exception as e:
        logging.error(f"Error in /action endpoint: {e}", exc_info=True)
        return f"An error occurred: {e}"

def handle_analyze(image, output_style):
    if image is None: return "Error: Please upload an image."
    
    try:
        image_b64 = pil_to_base64_str(image)
        with tempfile.TemporaryDirectory() as temp_dir:
            image_path = os.path.join(temp_dir, "image_to_analyze.png")
            save_base64_image(image_b64, image_path)
            
            is_mini = (output_style == "Mini JSON")
            description_str = run_wrapper(image_path=image_path, output_dir=temp_dir, json_mini=is_mini)
            
            parsed_json = json.loads(description_str)
            return json.dumps(parsed_json, indent=2)

    except Exception as e:
        logging.error(f"Error in /analyze endpoint: {e}", exc_info=True)
        return f"An error occurred: {e}"

def handle_analyze_yolo(image, output_style):
    if image is None: return None, "Error: Please upload an image."

    try:
        image_b64 = pil_to_base64_str(image)
        with tempfile.TemporaryDirectory() as temp_dir:
            request_id = str(uuid.uuid4())
            image_path = os.path.join(temp_dir, f"{request_id}.png")
            yolo_image_path = os.path.join(temp_dir, f"{request_id}_yolo_updated.png")
            save_base64_image(image_b64, image_path)

            is_mini = (output_style == "Mini JSON")
            description_str = run_wrapper(image_path=image_path, output_dir=temp_dir, json_mini=is_mini)

            parsed_json = json.loads(description_str)
            description_output = json.dumps(parsed_json, indent=2)
            
            yolo_image_result = Image.open(yolo_image_path)
            return yolo_image_result, description_output

    except Exception as e:
        logging.error(f"Error in /analyze_and_get_yolo: {e}", exc_info=True)
        return None, f"An error occurred: {e}"

def handle_generate(openai_key, image, prompt):
    if not openai_key: return "Error: OpenAI API Key is required for /generate."
    if image is None: return "Error: Please upload an image."
    if not prompt: return "Error: Please provide a prompt."

    try:
        llm_client = OpenAI(api_key=openai_key)
        image_b64 = pil_to_base64_str(image)

        with tempfile.TemporaryDirectory() as temp_dir:
            request_id = str(uuid.uuid4())
            original_image_path = os.path.join(temp_dir, f"{request_id}.png")
            yolo_updated_image_path = os.path.join(temp_dir, f"{request_id}_yolo_updated.png")
            save_base64_image(image_b64, original_image_path)

            image_description = run_wrapper(image_path=original_image_path, output_dir=temp_dir, json_mini=False)
            
            with open(yolo_updated_image_path, "rb") as f:
                yolo_updated_img_bytes = f.read()
            
            _, new_b64 = preprocess_image(yolo_updated_img_bytes, threshold=1500, scale=0.5, fmt="png")

        base64_image_url = f"data:image/png;base64,{new_b64}"
        messages = [
            {"role": "user", "content": [
                {"type": "text", "text": f'"Prompt: {prompt}"\nImage description:\n"{image_description}"'},
                {"type": "image_url", "image_url": {"url": base64_image_url, "detail": "high"}}
            ]}
        ]
    
        response = llm_client.chat.completions.create(model="gpt-4.1", messages=messages, temperature=0.2)
        return response.choices[0].message.content.strip()

    except Exception as e:
        logging.error(f"Error in /generate endpoint: {e}", exc_info=True)
        return f"An error occurred: {e}"

default_image_1 = Image.open("./res/bb_1.jpeg")
default_image_2 = Image.open("./res/mfa_1.jpeg")

def load_example_action_1(): return default_image_1, "Open and read Umico partner"
def load_example_action_2(): return default_image_2, "Sign up in the application"
def load_example_analyze_1(): return default_image_1
def load_example_analyze_2(): return default_image_2
def load_example_yolo_1(): return default_image_1
def load_example_yolo_2(): return default_image_2
def load_example_generate_1(): return default_image_1, "Generate the code for this screen for Android XML. Try to use constraint layout"
def load_example_generate_2(): return default_image_2, "Generate the code for this screen for Android XML. Try to use constraint layout"


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Deki Automata: UI Analysis and Generation")
    gr.Markdown("Provide your API keys below. The OpenAI key is only required for the 'Action' and 'Generate' tabs.")
    
    with gr.Row():
        openai_key_input = gr.Textbox(label="OpenAI API Key", placeholder="Enter your OpenAI API Key", type="password", scale=1)
    
    with gr.Tabs():
        with gr.TabItem("Action"):
            gr.Markdown("### Control a device with natural language.")
            with gr.Row():
                image_input_action = gr.Image(type="pil", label="Upload Screen Image")
                prompt_input_action = gr.Textbox(lines=2, placeholder="e.g., 'Open whatsapp and text my friend...'", label="Prompt")
            action_output = gr.Textbox(label="Response Command")
            action_button = gr.Button("Run Action", variant="primary")
            with gr.Row():
                example_action_btn1 = gr.Button("Load Example 1")
                example_action_btn2 = gr.Button("Load Example 2")
        
        with gr.TabItem("Analyze"):
            gr.Markdown("### Get a structured JSON description of the UI elements.")
            with gr.Row():
                image_input_analyze = gr.Image(type="pil", label="Upload Screen Image")
                with gr.Column():
                    output_style_analyze = gr.Radio(["Standard JSON", "Mini JSON"], label="Output Format", value="Standard JSON")
                    analyze_button = gr.Button("Analyze Image", variant="primary")
            analyze_output = gr.JSON(label="JSON Description")
            with gr.Row():
                example_analyze_btn1 = gr.Button("Load Example 1")
                example_analyze_btn2 = gr.Button("Load Example 2")

        with gr.TabItem("Analyze & Get YOLO"):
            gr.Markdown("### Get a JSON description and the image with detected elements.")
            with gr.Row():
                image_input_yolo = gr.Image(type="pil", label="Upload Screen Image")
                with gr.Column():
                     output_style_yolo = gr.Radio(["Standard JSON", "Mini JSON"], label="Output Format", value="Standard JSON")
                     yolo_button = gr.Button("Analyze and Visualize", variant="primary")
            with gr.Row():
                yolo_image_output = gr.Image(label="YOLO Annotated Image")
                description_output_yolo = gr.JSON(label="JSON Description")
            with gr.Row():
                example_yolo_btn1 = gr.Button("Load Example 1")
                example_yolo_btn2 = gr.Button("Load Example 2")

        with gr.TabItem("Generate"):
            gr.Markdown("### Generate code or text based on a screenshot.")
            with gr.Row():
                image_input_generate = gr.Image(type="pil", label="Upload Screen Image")
                prompt_input_generate = gr.Textbox(lines=2, placeholder="e.g., 'Generate the Android XML for this screen'", label="Prompt")
            generate_output = gr.Code(label="Generated Output", language="xml")
            generate_button = gr.Button("Generate", variant="primary")
            with gr.Row():
                example_generate_btn1 = gr.Button("Load Example 1")
                example_generate_btn2 = gr.Button("Load Example 2")

    action_button.click(fn=handle_action, inputs=[openai_key_input, image_input_action, prompt_input_action], outputs=action_output)
    analyze_button.click(fn=handle_analyze, inputs=[image_input_analyze, output_style_analyze], outputs=analyze_output)
    yolo_button.click(fn=handle_analyze_yolo, inputs=[image_input_yolo, output_style_yolo], outputs=[yolo_image_output, description_output_yolo])
    generate_button.click(fn=handle_generate, inputs=[openai_key_input, image_input_generate, prompt_input_generate], outputs=generate_output)

    example_action_btn1.click(fn=load_example_action_1, outputs=[image_input_action, prompt_input_action])
    example_action_btn2.click(fn=load_example_action_2, outputs=[image_input_action, prompt_input_action])
    example_analyze_btn1.click(fn=load_example_analyze_1, outputs=image_input_analyze)
    example_analyze_btn2.click(fn=load_example_analyze_2, outputs=image_input_analyze)
    example_yolo_btn1.click(fn=load_example_yolo_1, outputs=image_input_yolo)
    example_yolo_btn2.click(fn=load_example_yolo_2, outputs=image_input_yolo)
    example_generate_btn1.click(fn=load_example_generate_1, outputs=[image_input_generate, prompt_input_generate])
    example_generate_btn2.click(fn=load_example_generate_2, outputs=[image_input_generate, prompt_input_generate])

load_models()
demo.launch()