import io
import rasterio
from rasterio.mask import mask
from shapely import wkt as load_wkt
import matplotlib.pyplot as plt
import gradio as gr
import leafmap.foliumap as leafmap
import tempfile
import os


# -------------------------------
# Helper: Get raster bounds
# -------------------------------
def get_raster_bounds(file):
    try:
        with rasterio.open(file.name) as src:
            bounds = src.bounds
        return f"{bounds.left},{bounds.bottom},{bounds.right},{bounds.top}"
    except Exception as e:
        return f"❌ Error: {str(e)}"


# -------------------------------
# Clip TIFF (from upload)
# -------------------------------
def clip_tiff_from_upload(file, bbox):
    try:
        minx, miny, maxx, maxy = map(float, bbox.split(","))
        wkt = f"POLYGON(({minx} {miny}, {minx} {maxy}, {maxx} {maxy}, {maxx} {miny}, {minx} {miny}))"
        geom = load_wkt.loads(wkt)

        with rasterio.open(file.name) as src:
            clipped, transform = mask(src, [geom], crop=True)
            meta = src.meta.copy()
            meta.update({
                "driver": "GTiff",
                "height": clipped.shape[1],
                "width": clipped.shape[2],
                "transform": transform
            })

            # Save clipped file temporarily
            out_path = tempfile.mktemp(suffix=".tif")
            with rasterio.open(out_path, "w", **meta) as dest:
                dest.write(clipped)

        # Preview image
        fig, ax = plt.subplots(figsize=(5, 5))
        ax.imshow(clipped[0], cmap="terrain")
        ax.axis("off")
        buf = io.BytesIO()
        plt.savefig(buf, format="png", bbox_inches="tight")
        buf.seek(0)
        plt.close(fig)

        return buf, out_path, "✅ Clipped image generated successfully"
    except Exception as e:
        return None, None, f"❌ Error: {str(e)}"


# -------------------------------
# Visualize clipped TIFF with Leafmap
# -------------------------------
def visualize_leafmap(tif_path):
    try:
        m = leafmap.Map(center=[0, 0], zoom=2)
        m.add_raster(tif_path, layer_name="Clipped TIFF", colormap="terrain")
        return m.to_html()
    except Exception as e:
        return f"<p>❌ Leafmap Error: {str(e)}</p>"


# -------------------------------
# Build Gradio UI
# -------------------------------
with gr.Blocks() as demo:
    gr.Markdown("## 🌍 GeoTIFF Clipping Tool (Upload + Leafmap)")

    with gr.Tab("From Upload"):
        raster_file = gr.File(label="Upload GeoTIFF", file_types=[".tif"])
        bbox_file = gr.Textbox(label="Bounding Box (minx,miny,maxx,maxy)", value="")
        
        get_bounds_btn = gr.Button("📐 Get Raster Bounds")
        get_bounds_btn.click(get_raster_bounds, [raster_file], [bbox_file])

        output_image = gr.Image(label="Clipped TIFF Preview")
        download_file = gr.File(label="Download Clipped TIFF")
        status = gr.Textbox(label="Status")
        leafmap_view = gr.HTML(label="Leafmap Visualization")

        submit = gr.Button("✂️ Clip & Visualize")
        submit.click(
            clip_tiff_from_upload,
            [raster_file, bbox_file],
            [output_image, download_file, status]
        ).then(
            visualize_leafmap,
            [download_file],
            [leafmap_view]
        )


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