from typing import Optional
import spaces
import gradio as gr
import numpy as np
import torch
from PIL import Image
import io
import base64, os
from util.utils import check_ocr_box, get_yolo_model, get_caption_model_processor, get_som_labeled_img
from huggingface_hub import snapshot_download
import threading
import subprocess
import time

# Monkey patch for gradio_client JSON schema bug
try:
    from gradio_client import utils as gradio_client_utils
    original_json_schema_to_python_type = gradio_client_utils.json_schema_to_python_type
    
    def patched_json_schema_to_python_type(schema):
        """Patched version that handles boolean schemas (additionalProperties can be bool)"""
        try:
            if not isinstance(schema, dict):
                return "Any"
            return original_json_schema_to_python_type(schema)
        except (TypeError, AttributeError) as e:
            if "argument of type 'bool' is not iterable" in str(e):
                return "Any"
            raise
    
    gradio_client_utils.json_schema_to_python_type = patched_json_schema_to_python_type
except Exception as e:
    print(f"Warning: Could not apply gradio_client patch: {e}")

# Patch gradio blocks to handle schema generation errors
try:
    import gradio.blocks as gradio_blocks
    import warnings
    original_get_api_info = gradio_blocks.Blocks.get_api_info
    
    def patched_get_api_info(self):
        """Patched version that catches schema generation errors silently"""
        try:
            return original_get_api_info(self)
        except (TypeError, AttributeError) as e:
            if "argument of type 'bool' is not iterable" in str(e):
                # Silently skip - this is a known Gradio 5.16.0 bug
                return None
            raise
    
    gradio_blocks.Blocks.get_api_info = patched_get_api_info
except Exception as e:
    print(f"Warning: Could not patch gradio.blocks: {e}")

_yolo_model = None
_caption_model_processor = None

# Proper device handling
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {DEVICE}")

def load_models():
    global _yolo_model, _caption_model_processor
    if _yolo_model is None or _caption_model_processor is None:
        # Define repository and local directory
        repo_id = "microsoft/OmniParser-v2.0"  # HF repo
        local_dir = "weights"  # Target local directory
        # Download the entire repository
        print(f"Downloading repository to: {local_dir}...")
        snapshot_download(repo_id=repo_id, local_dir=local_dir, ignore_patterns=["*.msgpack", "*.h5", "*.ot"])
        print(f"Repository downloaded to: {local_dir}")
        _yolo_model = get_yolo_model(model_path='weights/icon_detect/model.pt')
        _caption_model_processor = get_caption_model_processor(
            model_name="florence2", 
            model_name_or_path="weights/icon_caption",
            device=DEVICE
        )
    return _yolo_model, _caption_model_processor

MARKDOWN = """
# OmniParser V2 for Pure Vision Based General GUI Agent 🔥
<div>
    <a href="https://arxiv.org/pdf/2408.00203">
        <img src="https://img.shields.io/badge/arXiv-2408.00203-b31b1b.svg" alt="Arxiv" style="display:inline-block;">
    </a>
</div>

OmniParser is a screen parsing tool to convert general GUI screen to structured elements. 
"""

# Only use @spaces.GPU on Hugging Face Spaces, not for local development
def process(
    image_input,
    box_threshold,
    iou_threshold,
    use_paddleocr,
    imgsz
):
    try:
        yolo_model, caption_model_processor = load_models()
        box_overlay_ratio = image_input.size[0] / 3200
        draw_bbox_config = {
            'text_scale': 0.8 * box_overlay_ratio,
            'text_thickness': max(int(2 * box_overlay_ratio), 1),
            'text_padding': max(int(3 * box_overlay_ratio), 1),
            'thickness': max(int(3 * box_overlay_ratio), 1),
        }
        
        # Use consistent OCR settings from omniparser.py
        ocr_bbox_rslt, is_goal_filtered = check_ocr_box(
            image_input, 
            display_img=False, 
            output_bb_format='xyxy', 
            goal_filtering=None, 
            easyocr_args={'paragraph': False, 'text_threshold': 0.9}, 
            use_paddleocr=use_paddleocr
        )
        text, ocr_bbox = ocr_bbox_rslt
        
        # Use consistent parameters from omniparser.py
        dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(
            image_input, 
            yolo_model, 
            BOX_TRESHOLD=box_threshold, 
            output_coord_in_ratio=True, 
            ocr_bbox=ocr_bbox,
            draw_bbox_config=draw_bbox_config, 
            caption_model_processor=caption_model_processor, 
            ocr_text=text,
            iou_threshold=iou_threshold, 
            imgsz=imgsz,
            use_local_semantics=True,
            scale_img=False,
            batch_size=32
        )
        
        image = Image.open(io.BytesIO(base64.b64decode(dino_labled_img)))
        print('finish processing')
        parsed_content_list = '\n'.join([f'icon {i}: ' + str(v) for i,v in enumerate(parsed_content_list)])
        return image, str(parsed_content_list)
    except Exception as e:
        print(f"Error during processing: {e}")
        import traceback
        traceback.print_exc()
        return None, f"Error: {str(e)}"

with gr.Blocks(analytics_enabled=False) as demo:
    gr.Markdown(MARKDOWN)
    with gr.Row():
        with gr.Column():
            image_input_component = gr.Image(
                type='pil', label='Upload image')
            box_threshold_component = gr.Slider(
                label='Box Threshold', minimum=0.01, maximum=1.0, step=0.01, value=0.05)
            iou_threshold_component = gr.Slider(
                label='IOU Threshold', minimum=0.01, maximum=1.0, step=0.01, value=0.1)
            use_paddleocr_component = gr.Checkbox(
                label='Use PaddleOCR', value=False)
            imgsz_component = gr.Slider(
                label='Icon Detect Image Size', minimum=640, maximum=1920, step=32, value=640)
            submit_button_component = gr.Button(
                value='Submit', variant='primary')
        with gr.Column():
            image_output_component = gr.Image(type='pil', label='Image Output')
            text_output_component = gr.Textbox(label='Parsed screen elements', placeholder='Text Output')
    submit_button_component.click(
        fn=process,
        inputs=[
            image_input_component,
            box_threshold_component,
            iou_threshold_component,
            use_paddleocr_component,
            imgsz_component
        ],
        outputs=[image_output_component, text_output_component]
    )

def start_fastapi_server():
    """Start FastAPI server in background"""
    try:
        import uvicorn
        print("Starting FastAPI server on port 8000...")
        uvicorn.run("server:app", host="0.0.0.0", port=8000, log_level="critical")
    except Exception as e:
        print(f"FastAPI server error: {e}")

# Start FastAPI server in a daemon thread (for local usage and external ports)
fastapi_thread = threading.Thread(target=start_fastapi_server, daemon=True)
fastapi_thread.start()
time.sleep(2)

print("\n" + "="*60)
print("OmniParser is ready!")
print("="*60)
print("Gradio UI:     http://localhost:7860")
print("FastAPI Docs:  http://localhost:8000/docs")
print("API Health:    http://localhost:8000/health")
print("="*60 + "\n")

# Use simple launch for HF Spaces, let it handle the configuration
# This avoids the explicit server_name/port which sometimes triggers the localhost check error
demo.queue().launch(show_api=False)