import gradio as gr
import os
import sys
import json
import shutil
import gdown
import time
from PIL import Image
from io import BytesIO

# ==================================
# SETUP
# ==================================

print("🚀 Gradio App Starting...")

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# Paths
UPLOAD_DIR = "/tmp/uploads/"
JSON_DIR = "/tmp/results/"
OUTPUT_DIR = "/tmp/output/"
MODEL_DIR = os.path.join(BASE_DIR, "rcnn_model", "scripts")
logo_path = os.path.join(BASE_DIR, "public", "logo.png")
model_path = os.path.join(OUTPUT_DIR, "model_final.pth")

# Google Drive model
GOOGLE_DRIVE_FILE_ID = "1yr64AOgaYZPTcQzG6cxG6lWBENHR9qjW"
GDRIVE_URL = f"https://drive.google.com/uc?id={GOOGLE_DRIVE_FILE_ID}"

# Create folders
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(JSON_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)

# Download model if missing
if not os.path.exists(model_path):
    print("🚀 Model file not found! Downloading...")
    try:
        # gdown.download(GDRIVE_URL, model_path, quiet=False)
        gdown.download(GDRIVE_URL, model_path, quiet=False, use_cookies=False)
        print("✅ Model downloaded successfully.")
    except Exception as e:
        print(f"❌ Failed to download model: {e}")

# Import model
sys.path.append(MODEL_DIR)
from rcnn_model.scripts.rcnn_run import main, write_config

cfg = write_config()

# ==================================
# MAIN PREDICTION FUNCTION
# ==================================

# def predict(uploaded_file):
#     if uploaded_file is None:
#         return None, None, "No file uploaded."
#     # uploaded_path = uploaded_file
#     # input_filename = os.path.basename(uploaded_path)
#     # print(f"✅ Image received at {uploaded_path}")

#     uploaded_path = os.path.join(UPLOAD_DIR, uploaded_file.name)
#     with open(uploaded_path, "wb") as f:
#         f.write(uploaded_file.read())
#     print(f"✅ Image saved to {uploaded_path}")


#     # Save uploaded image
#     # input_bytes = uploaded_file.read()
#     # img = Image.open(BytesIO(input_bytes)).convert("RGB")
#     # input_filename = uploaded_file.name
#     # uploaded_path = os.path.join(UPLOAD_DIR, input_filename)
#     # img.save(uploaded_path)
#     # print(f"✅ Image saved to {uploaded_path}")

#     # Prepare output paths
#     input_filename = uploaded_file.name
#     output_json_name = input_filename.replace(".png", "_result.json").replace(".jpg", "_result.json").replace(".jpeg", "_result.json")
#     output_image_name = input_filename.replace(".png", "_result.png").replace(".jpg", "_result.png").replace(".jpeg", "_result.png")

#     output_json_path = os.path.join(JSON_DIR, output_json_name)
#     output_image_path = os.path.join(JSON_DIR, output_image_name)

#     # Run model
#     main(cfg, uploaded_path, output_json_path, output_image_path)

#     # Read outputs
#     result_img = Image.open(output_image_path) if os.path.exists(output_image_path) else None
#     result_json = {}
#     if os.path.exists(output_json_path):
#         with open(output_json_path, "r") as jf:
#             result_json = json.load(jf)

#     return result_img, json.dumps(result_json, indent=2), None

def predict(uploaded_file_path):
    if uploaded_file_path is None:
        return None, None, "No file uploaded."

    input_filename = os.path.basename(uploaded_file_path)
    uploaded_path = os.path.join(UPLOAD_DIR, input_filename)
    shutil.copy(uploaded_file_path, uploaded_path)
    print(f"✅ Image saved to {uploaded_path}")

    # Prepare output paths
    output_json_name = input_filename.replace(".png", "_result.json").replace(".jpg", "_result.json").replace(".jpeg", "_result.json")
    output_image_name = input_filename.replace(".png", "_result.png").replace(".jpg", "_result.png").replace(".jpeg", "_result.png")

    output_json_path = os.path.join(JSON_DIR, output_json_name)
    output_image_path = os.path.join(JSON_DIR, output_image_name)

    # Run model
    main(cfg, uploaded_path, output_json_path, output_image_path)

    # Read outputs
    result_img = Image.open(output_image_path) if os.path.exists(output_image_path) else None
    result_json = {}
    if os.path.exists(output_json_path):
        with open(output_json_path, "r") as jf:
            result_json = json.load(jf)

    return result_img, json.dumps(result_json, indent=2), None

# ==================================
# GRADIO UI
# ==================================

with gr.Blocks() as demo:
    gr.Markdown("<h1 style='text-align: center;'>🏠 Inovonics 2D Floorplan Vectorizer</h1>")

    with gr.Row():
        with gr.Column():
            # uploaded_file = gr.File(label="Upload your Floorplan Image", type="filepath")
            uploaded_file = gr.File(label="Upload your Floorplan Image", type="file")
            run_button = gr.Button("Run Vectorizer 🔥")

        with gr.Column():
            output_image = gr.Image(label="🖼 Output Vectorized Image")
            output_json = gr.JSON(label="🧾 Output JSON")

    error_output = gr.Textbox(label="Error Message", visible=False)

    run_button.click(
        predict,
        inputs=[uploaded_file],
        outputs=[output_image, output_json, error_output]
    )

demo.launch(share=True)