import gradio as gr
from PIL import Image, ImageDraw
import tempfile
import json
import os
import requests
from gradio_client import utils as gradio_client_utils


_original_json_schema_to_python_type = gradio_client_utils._json_schema_to_python_type


def _safe_json_schema_to_python_type(schema, defs=None):
    """Compatibility patch for Gradio 4.44 / gradio_client 1.3 schema parsing.

    Some component schemas contain additionalProperties as a boolean. Older
    gradio_client assumes that value is always a dict and crashes while building
    API info during Space startup.
    """
    if isinstance(schema, bool):
        return "Any"
    if isinstance(schema, dict) and isinstance(schema.get("additionalProperties"), bool):
        schema = dict(schema)
        schema["additionalProperties"] = {}
    return _original_json_schema_to_python_type(schema, defs)


gradio_client_utils._json_schema_to_python_type = _safe_json_schema_to_python_type

# Roboflow API configuration
ROBOFLOW_API_KEY = os.getenv("ROBOFLOW_API_KEY")
ROOM_MODEL = "room-segmentation-frntt/1"
DOOR_WINDOW_MODEL = "door-detection-model/2"

color_options = ["Red", "Green", "Blue", "Yellow"]
layer_options = ["Room Detection", "Doors and Windows Detection"]

def roboflow_infer(image_path, model_id, api_key):
    """Direct HTTP call to Roboflow inference API"""
    if not api_key:
        return {"predictions": [], "error": "Missing ROBOFLOW_API_KEY Space secret"}
    url = f"https://detect.roboflow.com/{model_id}"
    params = {"api_key": api_key}
    
    with open(image_path, "rb") as f:
        files = {"file": f}
        response = requests.post(url, params=params, files=files)
    
    if response.status_code == 200:
        return response.json()
    else:
        print(f"API Error: {response.status_code} - {response.text}")
        return {"predictions": []}

def apply_zoom(image, zoom_factor):
    width, height = image.size
    new_width = int(width * zoom_factor)
    new_height = int(height * zoom_factor)
    return image.resize((new_width, new_height))

def detect_and_draw(image_path, model_id, filter_classes=None, color_choice=None):
    result = roboflow_infer(image_path, model_id, ROBOFLOW_API_KEY)
    
    original_img = Image.open(image_path)
    overlayed_img = original_img.copy()
    draw = ImageDraw.Draw(overlayed_img)
    counts = {}

    for prediction in result.get('predictions', []):
        pred_class = prediction.get('class', '').lower()

        if filter_classes and pred_class not in filter_classes:
            continue

        counts[pred_class] = counts.get(pred_class, 0) + 1

        x = int(prediction['x'] - prediction['width'] / 2)
        y = int(prediction['y'] - prediction['height'] / 2)
        width = int(prediction['width'])
        height = int(prediction['height'])

        draw.rectangle([x, y, x + width, y + height], outline=color_choice, width=2)
        label = f"{pred_class}"
        draw.text((x, y - 10), label, fill=color_choice)

    return overlayed_img, counts

def process_floor_plan(image, zoom_factor, color_choice, selected_layers):
    try:
        if image is None:
            return None, None, json.dumps({"error": "Please upload a floor plan image."}, indent=2)
        selected_layers = selected_layers or []

        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
            image.save(temp_file.name)
            temp_file_path = temp_file.name

        zoomed_image = apply_zoom(Image.open(temp_file_path), zoom_factor)
        zoomed_image.save(temp_file_path)

        room_overlay = None
        dw_overlay = None

        if "Room Detection" in selected_layers:
            room_overlay, room_counts = detect_and_draw(
                temp_file_path,
                ROOM_MODEL,
                filter_classes=["room"],
                color_choice=color_choice
            )
        else:
            room_counts = {}

        if "Doors and Windows Detection" in selected_layers:
            dw_overlay, dw_counts = detect_and_draw(
                temp_file_path,
                DOOR_WINDOW_MODEL,
                filter_classes=["door", "window"],
                color_choice=color_choice
            )
        else:
            dw_counts = {}

        combined_counts = {}
        combined_counts.update(room_counts)
        combined_counts.update(dw_counts)

        return room_overlay, dw_overlay, json.dumps(combined_counts, indent=2)

    except Exception as e:
        return None, None, json.dumps({"error": str(e)}, indent=2)

# Gradio 5.x Interface
with gr.Blocks(title="Floor Plan Detection") as demo:
    gr.Markdown("# 🏗️ Floor Plan Detection")
    gr.Markdown("Upload a floor plan image to detect rooms, doors, and windows.")

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Upload Floor Plan")
            zoom_factor = gr.Slider(minimum=0.5, maximum=3.0, value=1.0, step=0.1, label="Zoom Factor")
            color_choice = gr.Dropdown(choices=color_options, value="Red", label="Annotation Color")
            selected_layers = gr.CheckboxGroup(choices=layer_options, value=layer_options, label="Detection Layers")
            process_btn = gr.Button("Process", variant="primary")

        with gr.Column():
            room_output = gr.Image(type="pil", label="Room Detection")
            door_window_output = gr.Image(type="pil", label="Doors and Windows Detection")
            output_json = gr.Textbox(label="Detection Results (JSON)", lines=10)

    process_btn.click(
        fn=process_floor_plan,
        inputs=[input_image, zoom_factor, color_choice, selected_layers],
        outputs=[room_output, door_window_output, output_json]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, show_api=False)