import os
import re
import json
import time
import shutil
import uuid
import tempfile
import unicodedata
from io import BytesIO
from typing import Tuple, Optional, List, Iterable

import gradio as gr
import numpy as np
import torch
import spaces
from PIL import Image, ImageDraw, ImageFont

# Transformers & Qwen Utils
from transformers import (
    Qwen2_5_VLForConditionalGeneration,
    AutoProcessor,
)
from qwen_vl_utils import process_vision_info

# Gradio Theme Utils
from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes

colors.steel_blue = colors.Color(
    name="steel_blue",
    c50="#EBF3F8",
    c100="#D3E5F0",
    c200="#A8CCE1",
    c300="#7DB3D2",
    c400="#529AC3",
    c500="#4682B4",
    c600="#3E72A0",
    c700="#36638C",
    c800="#2E5378",
    c900="#264364",
    c950="#1E3450",
)

class SteelBlueTheme(Soft):
    def __init__(
        self,
        *,
        primary_hue: colors.Color | str = colors.gray,
        secondary_hue: colors.Color | str = colors.steel_blue,
        neutral_hue: colors.Color | str = colors.slate,
        text_size: sizes.Size | str = sizes.text_lg,
        font: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
        ),
        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
        ),
    ):
        super().__init__(
            primary_hue=primary_hue,
            secondary_hue=secondary_hue,
            neutral_hue=neutral_hue,
            text_size=text_size,
            font=font,
            font_mono=font_mono,
        )
        super().set(
            background_fill_primary="*primary_50",
            background_fill_primary_dark="*primary_900",
            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
            button_primary_text_color="white",
            button_primary_text_color_hover="white",
            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_800)",
            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_500)",
            block_title_text_weight="600",
            block_border_width="3px",
            block_shadow="*shadow_drop_lg",
            button_primary_shadow="*shadow_drop_lg",
            button_large_padding="11px",
        )

steel_blue_theme = SteelBlueTheme()

css = """
#main-title h1 { font-size: 2.3em !important; }
#out_img { height: 600px; object-fit: contain; }
"""

# -----------------------------------------------------------------------------
# 2. MODEL LOADING (Global Setup)
# -----------------------------------------------------------------------------

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# System Prompt
OS_SYSTEM_PROMPT = """You are a GUI agent. You are given a task and a screenshot of the current status. 
You need to generate the next action to complete the task.

Output your action inside a <tool_call> block using JSON format.
Include "coordinate": [x, y] in pixels for interactions.

Examples:
<tool_call>
{"name": "User", "arguments": {"action": "click", "coordinate": [400, 300]}}
</tool_call>

<tool_call>
{"name": "User", "arguments": {"action": "type", "coordinate": [100, 200], "text": "hello"}}
</tool_call>
"""

# Load Fara-7B
print("Loading Fara-7B...")
MODEL_ID_V = "microsoft/Fara-7B"
processor_v = AutoProcessor.from_pretrained(MODEL_ID_V, trust_remote_code=True)
model_v = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID_V,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16
).to(device).eval()

# Load UI-TARS-1.5-7B
print("Loading UI-TARS-1.5-7B...")
# Note: Using the official SFT repo. Adjust if you have a specific private repo.
MODEL_ID_X = "ByteDance-Seed/UI-TARS-1.5-7B" 
processor_x = AutoProcessor.from_pretrained(MODEL_ID_X, trust_remote_code=True)
model_x = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID_X,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).to(device).eval()

print("✅ All Models Loaded Successfully")

# -----------------------------------------------------------------------------
# 3. UTILS: IMAGE, PARSING, VISUALIZATION
# -----------------------------------------------------------------------------

def array_to_image(image_array: np.ndarray) -> Image.Image:
    if image_array is None:
        raise ValueError("No image provided. Please upload an image.")
    return Image.fromarray(np.uint8(image_array))

def get_navigation_prompt(task, image):
    return [
        {"role": "system", "content": [{"type": "text", "text": OS_SYSTEM_PROMPT}]},
        {"role": "user", "content": [
            {"type": "image", "image": image},
            {"type": "text", "text": f"Instruction: {task}"},
        ]},
    ]

def parse_tool_calls(response: str) -> list[dict]:
    """
    Parses the <tool_call>{JSON}</tool_call> format.
    """
    actions = []
    matches = re.findall(r"<tool_call>(.*?)</tool_call>", response, re.DOTALL)
    
    for match in matches:
        try:
            json_str = match.strip()
            data = json.loads(json_str)
            
            args = data.get("arguments", {})
            coords = args.get("coordinate", [])
            action_type = args.get("action", "unknown")
            text_content = args.get("text", "")
            
            if coords and isinstance(coords, list) and len(coords) == 2:
                actions.append({
                    "type": action_type,
                    "x": float(coords[0]),
                    "y": float(coords[1]),
                    "text": text_content,
                    "raw_json": data
                })
                print(f"Parsed Action: {action_type} at {coords}")
            else:
                # Handle actions without coordinates (like pressing enter generally)
                actions.append({
                    "type": action_type,
                    "text": text_content,
                    "raw_json": data
                })
                
        except json.JSONDecodeError:
            print(f"Failed to parse JSON: {match}")
            
    return actions

def create_localized_image(original_image: Image.Image, actions: list[dict]) -> Optional[Image.Image]:
    """Draws markers on the image based on parsed pixel coordinates."""
    if not actions:
        return None
    
    img_copy = original_image.copy()
    draw = ImageDraw.Draw(img_copy)
    width, height = img_copy.size
    
    try:
        font = ImageFont.load_default()
    except:
        font = None
    
    colors = {
        'type': 'blue', 
        'click': 'red', 
        'left_click': 'red',
        'right_click': 'purple',
        'double_click': 'orange',
        'unknown': 'green'
    }
    
    for act in actions:
        # Only draw if coordinates exist
        if 'x' not in act or 'y' not in act:
            continue

        x = act['x']
        y = act['y']
        
        # Check if Normalized (0.0 - 1.0) or Absolute (Pixels > 1.0)
        if x <= 1.0 and y <= 1.0 and x > 0:
            pixel_x = int(x * width)
            pixel_y = int(y * height)
        else:
            pixel_x = int(x)
            pixel_y = int(y)
            
        action_type = act['type']
        color = colors.get(action_type, 'green')
        
        # Draw Circle Target
        r = 12 
        draw.ellipse([pixel_x - r, pixel_y - r, pixel_x + r, pixel_y + r], outline=color, width=4)
        draw.ellipse([pixel_x - 3, pixel_y - 3, pixel_x + 3, pixel_y + 3], fill=color)
        
        # Draw Label text
        label_text = f"{action_type}"
        if act['text']:
            label_text += f": '{act['text']}'"
            
        text_pos = (pixel_x + 15, pixel_y - 10)
        bbox = draw.textbbox(text_pos, label_text, font=font)
        draw.rectangle(bbox, fill="black")
        draw.text(text_pos, label_text, fill="white", font=font)

    return img_copy

# -----------------------------------------------------------------------------
# 4. PROCESSING LOGIC
# -----------------------------------------------------------------------------

@spaces.GPU
def process_screenshot(input_numpy_image: np.ndarray, task: str, model_choice: str) -> Tuple[str, Optional[Image.Image]]:
    if input_numpy_image is None:
        return "⚠️ Please upload an image first.", None

    # 1. Select Model
    if model_choice == "Fara-7B":
        model = model_v
        processor = processor_v
    elif model_choice == "UI-TARS-1.5-7B":
        model = model_x
        processor = processor_x
    else:
        return "Invalid model selection", None

    # 2. Prepare Data
    input_pil_image = array_to_image(input_numpy_image)
    prompt = get_navigation_prompt(task, input_pil_image)

    # 3. Generate
    text_prompts = processor.apply_chat_template(
        prompt, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(prompt)
    
    inputs = processor(
        text=[text_prompts],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    )
    inputs = inputs.to(device)

    print(f"Generating with {model_choice}...")
    with torch.no_grad():
        generated_ids = model.generate(**inputs, max_new_tokens=512)
    
    generated_ids_trimmed = [
        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    
    raw_response = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0]
    
    print(f"Raw Output:\n{raw_response}")
    
    # 4. Parse & Visualize
    actions = parse_tool_calls(raw_response)
    
    output_image = input_pil_image
    if actions:
        visualized = create_localized_image(input_pil_image, actions)
        if visualized:
            output_image = visualized
            
    return raw_response, output_image

# -----------------------------------------------------------------------------
# 5. GRADIO UI
# -----------------------------------------------------------------------------

with gr.Blocks(theme=steel_blue_theme, css=css) as demo:
    gr.Markdown("# **CUA GUI Agent 🖥️**", elem_id="main-title")
    gr.Markdown("Upload a screenshot, select a model, and provide a task. The model will determine the precise UI coordinates and actions.")

    with gr.Row():
        with gr.Column(scale=2):
            input_image = gr.Image(label="Upload Screenshot", height=500)
            
            with gr.Row():
                model_choice = gr.Radio(
                    choices=["Fara-7B", "UI-TARS-1.5-7B"],
                    label="Select Model",
                    value="Fara-7B",
                    interactive=True
                )
            
            task_input = gr.Textbox(
                label="Task Instruction",
                placeholder="e.g. Input the server address readyforquantum.com...",
                lines=2
            )
            submit_btn = gr.Button("Analyze UI & Generate Action", variant="primary")

        with gr.Column(scale=3):
            output_image = gr.Image(label="Visualized Action Points", elem_id="out_img", height=500)
            output_text = gr.Textbox(label="Raw Model Output (JSON)", lines=8, show_copy_button=True)

    # Wire up the button
    submit_btn.click(
        fn=process_screenshot,
        inputs=[input_image, task_input, model_choice],
        outputs=[output_text, output_image]
    )
    
    # Examples
    gr.Examples(
        examples=[
            ["./assets/google.png", "Search for 'Hugging Face'", "Fara-7B"], 
        ],
        inputs=[input_image, task_input, model_choice],
        label="Quick Examples"
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch(show_error=True)