"""RoofSight — Roof Plane Segmentation

Gradio web app using NVIDIA C-RADIOv4-H (631M params) for zero-shot
roof type segmentation combined with RANSAC 3D plane fitting on
Google Solar API DSM data.

Deployed on HuggingFace Spaces with ZeroGPU.
"""

import json
import os
import tempfile

import gradio as gr
import numpy as np
import torch

# Pre-load model weights to CPU at startup (moved to GPU per-request by ZeroGPU)
from radio_backbone import load_model
print("Initializing RoofSight...")
device = "cuda" if torch.cuda.is_available() else "cpu"
load_model(device="cpu")

# Try importing spaces for ZeroGPU decorator
try:
    import spaces
    GPU_DECORATOR = spaces.GPU
except ImportError:
    # Local dev without ZeroGPU
    def GPU_DECORATOR(fn):
        return fn

GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY", "")


@GPU_DECORATOR
def process(address, api_key, radius, ransac_threshold, max_planes, min_area):
    """Run the full pipeline and return results for Gradio."""
    from pipeline import run as run_pipeline

    api_key = api_key.strip() if api_key.strip() else GOOGLE_API_KEY
    if not api_key:
        return None, None, "", "No API key provided. Enter your Google Solar API key or set GOOGLE_API_KEY."

    compute_device = "cuda" if torch.cuda.is_available() else "cpu"

    try:
        result = run_pipeline(
            address=address,
            api_key=api_key,
            radius_meters=int(radius),
            ransac_threshold=ransac_threshold,
            max_planes=int(max_planes),
            min_area_sqft=min_area,
            device=compute_device,
        )

        geojson_str = json.dumps(result.geojson, indent=2)
        status = "\n".join(result.status)
        return result.original_image, result.overlay, geojson_str, status

    except Exception as e:
        return None, None, "", f"Error: {str(e)}"


def save_geojson(geojson_str):
    """Save GeoJSON to a downloadable temp file."""
    if not geojson_str or geojson_str.strip() == "":
        return None
    tmp = tempfile.NamedTemporaryFile(suffix=".geojson", delete=False, mode="w")
    tmp.write(geojson_str)
    tmp.close()
    return tmp.name


# --- Gradio UI ---
with gr.Blocks(title="RoofSight", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# RoofSight")
    gr.Markdown(
        "Roof plane segmentation from satellite imagery.  \n"
        "**C-RADIOv4-H** (DINOv3 + SAM3 + SigLIP2) + **RANSAC** 3D plane fitting."
    )

    with gr.Row():
        # --- Left column: inputs ---
        with gr.Column(scale=1):
            address = gr.Textbox(
                label="Property Address",
                placeholder="123 Main St, City, ST 12345",
            )
            api_key = gr.Textbox(
                label="Google API Key",
                type="password",
                placeholder="Leave blank to use server key",
            )

            with gr.Accordion("Settings", open=False):
                radius = gr.Slider(
                    25, 100, value=50, step=5,
                    label="Search Radius (m)",
                )
                ransac_thresh = gr.Slider(
                    0.05, 0.50, value=0.15, step=0.01,
                    label="RANSAC Threshold (m)",
                    info="Distance tolerance for plane inliers. Lower = tighter fit.",
                )
                max_planes = gr.Slider(
                    2, 15, value=8, step=1,
                    label="Max Planes",
                )
                min_area = gr.Slider(
                    10, 500, value=50, step=10,
                    label="Min Area (sqft)",
                    info="Polygons smaller than this are excluded.",
                )

            run_btn = gr.Button("Segment Roof", variant="primary", size="lg")

        # --- Right column: outputs ---
        with gr.Column(scale=2):
            with gr.Row():
                orig_img = gr.Image(label="Satellite Image", show_download_button=False)
                overlay_img = gr.Image(label="Roof Planes", show_download_button=True)

            status_md = gr.Markdown(label="Status")

            with gr.Accordion("GeoJSON Output", open=True):
                geojson_out = gr.Code(language="json", lines=15, label="GeoJSON")
                with gr.Row():
                    dl_btn = gr.Button("Download GeoJSON", size="sm")
                    dl_file = gr.File(label="Download", visible=True)

    # --- Wiring ---
    run_btn.click(
        fn=process,
        inputs=[address, api_key, radius, ransac_thresh, max_planes, min_area],
        outputs=[orig_img, overlay_img, geojson_out, status_md],
    )
    dl_btn.click(fn=save_geojson, inputs=[geojson_out], outputs=[dl_file])


if __name__ == "__main__":
    demo.launch()