app/ui.py

from __future__ import annotations

from pathlib import Path
from typing import Dict
from PIL import Image

import gradio as gr

from .config import DEFAULT_ANALYZER_SETTINGS, IMAGE_EXTS
from .data_sources import list_all_data_inputs
from .safety import AnalysisRequest, AnalysisSummary, SafetyAnalyzer
from .visualization import compose_view



def _make_request(
    footprint_m,
    std_thresh,
    grad_thresh,
    use_water_mask,
    use_road_mask,
    use_roof_mask,
    use_tree_mask,
    water_prompt,
    road_prompt,
    roof_prompt,
    tree_prompt,
    altitude_m,
    fov_deg,
    clearance_factor,
    process_res_cap,
    segmentation_max_side,
    segmentation_model_id,
    segmentation_score_thresh,
    segmentation_mask_thresh,
    coverage_strictness,
    model_id,
    openness_weight,
    texture_threshold,
) -> AnalysisRequest:
    return AnalysisRequest(
        footprint_m=footprint_m,
        std_thresh=std_thresh,
        grad_thresh=grad_thresh,
        use_water_mask=use_water_mask,
        use_road_mask=use_road_mask,
        use_roof_mask=use_roof_mask,
        use_tree_mask=use_tree_mask,
        water_prompt=water_prompt,
        road_prompt=road_prompt,
        roof_prompt=roof_prompt,
        tree_prompt=tree_prompt,
        altitude_m=altitude_m,
        fov_deg=fov_deg,
        clearance_factor=clearance_factor,
        process_res_cap=process_res_cap,
        depth_smoothing_base=0.0,
        segmentation_max_side=segmentation_max_side,
        segmentation_model_id=segmentation_model_id,
        segmentation_score_thresh=segmentation_score_thresh,
        segmentation_mask_thresh=segmentation_mask_thresh,
        coverage_strictness=coverage_strictness,
        model_id=model_id,
        openness_weight=openness_weight,
        texture_threshold=texture_threshold,
    )


def _format_status(summary: AnalysisSummary | None) -> str:
    if not summary:
        return "**Status**\nAwaiting analysis."
    masks_line = " / ".join(
        f"{label}:{'on' if enabled else 'off'}"
        for label, enabled in (
            ("Water", summary.water_mask_enabled),
            ("Road", summary.road_mask_enabled),
            ("Tree", summary.tree_mask_enabled),
            ("Roof", summary.roof_mask_enabled),
        )
    )
    warning_text = ""
    if summary.warnings:
        warning_text = "\nWarnings: " + " | ".join(summary.warnings)
    return (
        "**Run Status**\n"
        f"- Model: `{summary.model_id}`\n"
        f"- Process res: {summary.process_resolution}px; Runtime: {summary.runtime_ms:.0f} ms\n"
        f"- Footprint: {summary.footprint_m:.1f} m (~{summary.footprint_image_px}px)\n"
        f"- Safe: {summary.safe_area_pct:.1f}% | Hazard: {summary.hazard_pct:.1f}%"
        f"\n- Masks: {masks_line}"
        f"{warning_text}"
    )


def _format_metrics(summary: AnalysisSummary | None) -> str:
    if not summary:
        return "No metrics yet. Run the analyzer to populate this section."
    lines = []
    if not summary.used_valid_center:
        lines.append("Warning: No fully safe footprint; showing lowest-variance patch.")
    if summary.warnings:
        warn_lines = "<br/>".join(f"⚠️ {msg}" for msg in summary.warnings)
        lines.append(warn_lines)
    return "<br/>".join(lines)


def build_ui(analyzer: SafetyAnalyzer | None = None) -> gr.Blocks:
    analyzer = analyzer or SafetyAnalyzer()
    defaults = DEFAULT_ANALYZER_SETTINGS
    data_inputs = list_all_data_inputs()

    with gr.Blocks(title="Drone Landing Site Safety Analyzer") as demo:
        gr.Markdown(
            "## Drone Landing Site Safety Analyzer\n"
            "Evaluate aerial imagery to spot flat, obstacle-free landing sites.",
            elem_classes="tight-title",
        )
        gr.HTML(
                    """
                    <style>
                    #preview-wrap { position: relative; }
                    #preview-wrap .hover-legend {
                        position: absolute;
                        right: 12px;
                        bottom: 12px;
                        background: rgba(0, 0, 0, 0.65);
                color: #fff;
                padding: 8px 10px;
                border-radius: 10px;
                font-size: 12px;
                line-height: 1.4;
                opacity: 0;
                transition: opacity 0.2s ease;
                pointer-events: none;
            }
            #preview-wrap:hover .hover-legend { opacity: 1; }
                    #preview-wrap .hover-legend .row { display: flex; align-items: center; gap: 6px; margin-bottom: 4px; }
                    #preview-wrap .hover-legend .swatch { width: 12px; height: 12px; border-radius: 3px; display: inline-block; }
                    </style>
                    """
                )
        images_state = gr.State({})

        with gr.Row(equal_height=False):
            with gr.Column(scale=1, min_width=280):
                input_path = gr.Dropdown(
                    label="Input file",
                    choices=data_inputs,
                    value=data_inputs[0] if data_inputs else None,
                    allow_custom_value=True,
                    info="Pick any VISLOC image under data/Image/VISLOC/.",
                )
                model_id = defaults.model_id
                process_res_cap = gr.Slider(
                    label="Processing max side (px)",
                    value=defaults.process_res_cap,
                    minimum=512,
                    maximum=2048,
                    step=32,
                    info="Global resolution cap for depth/analysis steps.",
                )
                footprint_m = gr.Slider(
                    label="Landing footprint (meters)",
                    value=defaults.footprint_m,
                    minimum=1,
                    maximum=150,
                    step=1,
                    info="Side length required for a clear landing zone.",
                )
                std_thresh = gr.Slider(
                    label="Flatness threshold",
                    value=defaults.std_thresh,
                    minimum=0.001,
                    maximum=0.08,
                    step=0.001,
                    info="Lower favors flatter regions but may reduce candidates.",
                )
                grad_thresh = gr.Slider(
                    label="Gradient threshold",
                    value=defaults.grad_thresh,
                    minimum=0.02,
                    maximum=1.0,
                    step=0.01,
                    info="Lower suppresses slopes/edges; higher tolerates tilt.",
                )
                with gr.Accordion("Hazard segmentation", open=False):
                    gr.Markdown("Select which hazards to mask using SAM3 segmentation.")
                    with gr.Row():
                        use_water_mask = gr.Checkbox(label="Exclude water", value=True, info="Mask water regions.")
                        use_road_mask = gr.Checkbox(label="Exclude roads", value=True, info="Mask road surfaces.")
                    with gr.Row():
                        use_tree_mask = gr.Checkbox(label="Exclude trees", value=True, info="Mask trees/foliage.")
                        use_roof_mask = gr.Checkbox(label="Exclude rooftops", value=True, info="Mask rooftop areas.")
                with gr.Accordion("Segmentation settings", open=False):
                    gr.Markdown("Control SAM3 and prompts.")
                    segmentation_model_id = gr.Dropdown(
                        label="Segmentation model",
                        value=defaults.segmentation_model_id,
                        choices=[
                            ("SAM3", "facebook/sam3"),
                        ],
                        info="Choose segmentation backbone for water/road masks (SAM3 only).",
                    )
                    segmentation_max_side = gr.Slider(
                        label="Segmentation max side (px)",
                        value=defaults.segmentation_max_side,
                        minimum=256,
                        maximum=2048,
                        step=32,
                        info="Largest long-side resolution for running the segmentation model.",
                    )
                    water_prompt = gr.Textbox(
                        label="Water prompt",
                        value=defaults.water_prompt,
                        placeholder="e.g., water",
                    )
                    road_prompt = gr.Textbox(
                        label="Road prompt",
                        value=defaults.road_prompt,
                        placeholder="e.g., road",
                    )
                    roof_prompt = gr.Textbox(
                        label="Roof prompt",
                        value=defaults.roof_prompt,
                        placeholder="e.g., roof",
                    )
                    tree_prompt = gr.Textbox(
                        label="Tree prompt",
                        value=defaults.tree_prompt,
                        placeholder="e.g., tree",
                    )
                    segmentation_score_thresh = gr.Slider(
                        label="Segmentation score threshold",
                        value=defaults.segmentation_score_thresh,
                        minimum=0.1,
                        maximum=0.9,
                        step=0.05,
                        info="Minimum instance confidence from SAM3 to keep a mask.",
                    )
                    segmentation_mask_thresh = gr.Slider(
                        label="Segmentation mask threshold",
                        value=defaults.segmentation_mask_thresh,
                        minimum=0.1,
                        maximum=0.9,
                        step=0.05,
                        info="Pixel probability threshold when binarizing SAM3 masks.",
                    )
                with gr.Accordion("Advanced settings", open=False):
                    gr.Markdown("Adjust detail levels and scoring.")
                    clearance_factor = gr.Slider(
                        label="Clearance factor",
                        value=defaults.clearance_factor,
                        minimum=0.0,
                        maximum=2.0,
                        step=0.05,
                        info="Dilate unsafe regions relative to footprint size.",
                    )
                    coverage_strictness = gr.Slider(
                        label="Coverage strictness",
                        value=defaults.coverage_strictness,
                        minimum=0.8,
                        maximum=1.0,
                        step=0.001,
                        info="Minimum fraction of the footprint that must be safe.",
                    )
                    openness_weight = gr.Slider(
                        label="Open-area preference",
                        value=defaults.openness_weight,
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        info="Higher favors distance from hazards over absolute flatness when picking the landing spot.",
                    )
                    texture_threshold = gr.Slider(
                        label="Clutter tolerance",
                        value=defaults.texture_threshold,
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        info="Lower values avoid visually textured (high-contrast) regions like tracks or debris.",
                    )
                    # Plane removal fixed to least squares; toggle removed.
                with gr.Accordion("Camera settings", open=False):
                    gr.Markdown("Configure capture assumptions for footprint sizing.")
                    altitude_m = gr.Slider(
                        label="Camera altitude (m)",
                        value=defaults.altitude_m,
                        minimum=10,
                        maximum=1500,
                        step=5,
                    )
                    fov_deg = gr.Slider(
                        label="Camera FOV (deg)",
                        value=defaults.fov_deg,
                        minimum=30,
                        maximum=150,
                        step=1,
                    )
            with gr.Column(scale=2, min_width=520, elem_id="preview-wrap", elem_classes="preview-col"):
                main_view = gr.Image(
                    label="Analyzed",
                    height=485,
                    elem_id="main-preview",
                    buttons=["download"],
                )
                orig_view = gr.Image(
                    label="Original",
                    height=485,
                    buttons=["download"],
                )
                gr.HTML(
                    """
                    <style>
                    .preview-col { gap: 0 !important; }
                    .tight-status { margin-top: 0 !important; padding-top: 0 !important; }
                    .tight-title { margin-bottom: 0 !important; }
                    </style>
                    """
                )
            with gr.Column(scale=1, min_width=220):
                base_view = gr.Dropdown(
                    label="Base view",
                    value="RGB",
                    choices=[
                        "RGB",
                        "Depth",
                        "Flatness map (std)",
                        "Depth gradient",
                        "Gradient mask",
                        "Water mask",
                        "Road mask",
                        "Tree mask",
                        "Safety score",
                        "Safety heatmap overlay",
                    ],
                )
                spot_on = gr.Checkbox(label="Show optimal landing spot", value=True, info="Show recommended landing box/marker.")
                hazard_on = gr.Checkbox(label="Risk highlight", value=True, info="Show depth-based risk map.")
                hazards_on = gr.Checkbox(label="Segmented hazard highlight", value=True, info="Show segmentation + roof hazards.")
                grad_on = gr.Checkbox(label="Depth gradient", value=False, info="Gradient magnitude overlay.")
                with gr.Row():
                    run_btn = gr.Button("Run", variant="primary")
                    stop_btn = gr.Button("Stop", variant="stop")
                status_card = gr.Markdown("**Status:** Awaiting analysis.", elem_classes="tight-status")
                metrics_card = gr.Markdown("")

        def process_any(
            input_path,
            footprint_m,
            std_thresh,
            grad_thresh,
            use_water_mask,
            use_road_mask,
            use_roof_mask,
            use_tree_mask,
            water_prompt,
            road_prompt,
            roof_prompt,
            tree_prompt,
            altitude_m,
            fov_deg,
            clearance_factor,
            process_res_cap,
            segmentation_max_side,
            segmentation_model_id,
            segmentation_score_thresh,
            segmentation_mask_thresh,
            coverage_strictness,
            openness_weight,
            texture_threshold,
            base_view,
            hazard_on,
            hazards_on,
            grad_on,
            spot_on,
        ):
            if not input_path:
                raise gr.Error("Select an input image first.")
            path = Path(input_path)
            if not path.exists():
                raise gr.Error(f"Input not found: {path}")
            if path.suffix.lower() not in IMAGE_EXTS:
                raise gr.Error(f"Unsupported input type for path: {path} (images only)")
            request = _make_request(
                footprint_m,
                std_thresh,
                grad_thresh,
                use_water_mask,
                use_road_mask,
                use_roof_mask,
                use_tree_mask,
                water_prompt,
                road_prompt,
                roof_prompt,
                tree_prompt,
                altitude_m,
                fov_deg,
                clearance_factor,
                process_res_cap,
                segmentation_max_side,
                segmentation_model_id,
                segmentation_score_thresh,
                segmentation_mask_thresh,
                coverage_strictness,
                model_id,
                openness_weight,
                texture_threshold,
            )
            try:
                result = analyzer.process_path(path, request)
            except ValueError as exc:
                raise gr.Error(str(exc)) from exc
            imgs = result.images
            summary = result.summary
            composed = compose_view(
                imgs,
                base_view,
                True,
                0.2,
                hazard_on,
                0.2,
                hazards_on,
                grad_on,
                False,
                False,
                spot_on=spot_on,
            )
            orig = imgs.get("RGB")
            orig_with_spot = orig
            if spot_on and imgs.get("Landing spot overlay") is not None and orig is not None:
                orig_with_spot = Image.alpha_composite(orig.convert("RGBA"), imgs["Landing spot overlay"]).convert("RGB")
            view = composed
            return imgs, view, orig_with_spot or orig, _format_status(summary), _format_metrics(summary)

        run_inputs = [
            input_path,
            footprint_m,
            std_thresh,
            grad_thresh,
            use_water_mask,
            use_road_mask,
            use_roof_mask,
            use_tree_mask,
            water_prompt,
            road_prompt,
            roof_prompt,
            tree_prompt,
            altitude_m,
            fov_deg,
            clearance_factor,
            process_res_cap,
            segmentation_max_side,
            segmentation_model_id,
            segmentation_score_thresh,
            segmentation_mask_thresh,
            coverage_strictness,
            openness_weight,
            texture_threshold,
            base_view,
            hazard_on,
            hazards_on,
            grad_on,
            spot_on,
        ]

        run_event = run_btn.click(
            fn=process_any,
            inputs=run_inputs,
            outputs=[images_state, main_view, orig_view, status_card, metrics_card],
        )
        stop_btn.click(fn=None, inputs=None, outputs=None, cancels=[run_event])

        overlay_inputs = [
            images_state,
            base_view,
            hazard_on,
            hazards_on,
            grad_on,
            spot_on,
        ]

        def update_overlays_only(
            images_state_val,
            base_view_val,
            hazard_on_val,
            hazards_on_val,
            grad_on_val,
            spot_on_val,
        ):
            if not images_state_val:
                return images_state_val, gr.update(), gr.update(), gr.update(), gr.update()
            composed = compose_view(
                images_state_val,
                base_view_val,
                True,
                0.2,
                hazard_on_val,
                0.2,
                hazards_on_val,
                grad_on_val,
                False,
                False,
                spot_on_val,
            )
            return images_state_val, composed, gr.update(), gr.update(), gr.update()

        base_view.change(
            fn=update_overlays_only,
            inputs=overlay_inputs,
            outputs=[images_state, main_view, orig_view, status_card, metrics_card],
        )
        overlay_toggle_controls = (
            hazard_on,
            hazards_on,
            grad_on,
            spot_on,
        )
        for control in overlay_toggle_controls:
            control.change(
                fn=update_overlays_only,
                inputs=overlay_inputs,
                outputs=[images_state, main_view, orig_view, status_card, metrics_card],
            )
        # Opacity sliders removed; overlays now use fixed alpha.

    return demo


__all__ = ["build_ui"]