app.py

import io
import time
import sys
import tempfile
from pathlib import Path

import gradio as gr
from PIL import Image
import spaces
import numpy as np
import rasterio
from inference.geotiff import normalize_rgb, save_sr_geotiff

# ============================================================
# PATHS
# ============================================================

ROOT = Path(__file__).resolve().parent

if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))


# ============================================================
# ZERO-GPU MODEL LOADING (CPU initially)
# ============================================================

from models.hatsat.hatsat_inference import HATSATInference
from models.esrgan.esrgan_inference import ESRGANInference
from models.sen2sr.sen2sr_inference import Sen2SRInference

print("Loading models to CPU for ZeroGPU deployment...")
try:
    hatsat_model = HATSATInference(device="cpu")
    esrgan_model = ESRGANInference(
        checkpoint_path=str(ROOT / "weights" / "esrgan" / "RRDB_ESRGAN_x4.pth"),
        device="cpu"
    )
    sen2sr_model = Sen2SRInference(device="cpu")
    print("Models loaded successfully on CPU.")
except Exception as e:
    print(f"Error loading models: {e}")
    hatsat_model = None
    esrgan_model = None
    sen2sr_model = None

# ============================================================
# ZERO-GPU INFERENCE WRAPPERS
# ============================================================

@spaces.GPU(duration=120)
def infer_hatsat(image):
    if hatsat_model is None:
        raise RuntimeError("HATSAT model is not loaded.")
    return hatsat_model.predict(image)

@spaces.GPU(duration=120)
def infer_esrgan(image):
    if esrgan_model is None:
        raise RuntimeError("ESRGAN model is not loaded.")
    return esrgan_model.predict(image)

@spaces.GPU(duration=120)
def infer_sen2sr(image):
    if sen2sr_model is None:
        raise RuntimeError("Sen2SR model is not loaded.")
    return sen2sr_model.predict(image)


# ============================================================
# CUSTOM CSS
# ============================================================

CSS = """
/* Google Font */
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');

body, .gradio-container {
    font-family: 'Inter', sans-serif !important;
    background-color: #f8fafc !important;
}

/* Hide default Gradio padding/footers */
footer { display: none !important; }
.gradio-container { max-width: 100% !important; padding: 0 !important; margin: 0 !important; }

/* Top Navbar */
.top-nav {
    background: linear-gradient(135deg, #03182b 0%, #07355c 100%);
    padding: 20px 40px;
    display: flex;
    justify-content: space-between;
    align-items: center;
    color: white;
}
.nav-left { display: flex; align-items: center; gap: 16px; }
.nav-logo-icon { width: 36px; height: 36px; fill: #38bdf8; }
.nav-title { display: flex; flex-direction: column; }
.nav-logo-text { font-size: 1.6rem; font-weight: 700; color: #38bdf8; margin: 0; line-height: 1.1; }
.nav-logo-text span { color: white; }
.nav-subtitle { font-size: 0.8rem; color: #cbd5e1; margin: 0; }
.nav-links { display: flex; gap: 30px; font-size: 0.95rem; align-items: center; }
.nav-links a { color: #cbd5e1; text-decoration: none; font-weight: 500; transition: color 0.2s; }
.nav-links a:hover { color: white; }
.nav-links a.active { color: white; border-bottom: 2px solid white; padding-bottom: 4px; }
.nav-btn { background: #34d399; color: #064e3b !important; padding: 10px 20px; border-radius: 6px; font-weight: 600; text-decoration: none; transition: background 0.2s; }
.nav-btn:hover { background: #10b981; }

/* Sidebar */
.sidebar-container { background: #f8fafc; padding: 30px 20px; height: 100%; border-right: 1px solid #e2e8f0; display: flex; flex-direction: column; justify-content: space-between; }
.sidebar-links { display: flex; flex-direction: column; gap: 8px; }
.sidebar-link { display: flex; align-items: center; padding: 14px 20px; border-radius: 10px; color: #475569; text-decoration: none; font-weight: 500; font-size: 0.95rem; gap: 16px; transition: background 0.2s; cursor: pointer; }
.sidebar-link:hover { background: #f1f5f9; }
.sidebar-link.active { background: #d1fae5; color: #065f46; font-weight: 600; }
.sidebar-link svg { width: 20px; height: 20px; }
.sidebar-banner { background: #e0f2fe; padding: 24px 20px; border-radius: 16px; text-align: center; color: #0369a1; display: flex; flex-direction: column; align-items: center; gap: 12px; }
.sidebar-banner svg { width: 40px; height: 40px; color: #0284c7; }
.sidebar-banner div { font-weight: 600; font-size: 0.95rem; line-height: 1.4; }

/* Main Content Area */
.main-wrapper { padding: 30px 40px; }
.page-title { font-size: 2.2rem; font-weight: 700; color: #0f172a; margin: 0 0 8px; letter-spacing: -0.5px; }
.page-title span { color: #0369a1; }
.page-subtitle { color: #64748b; font-size: 1.05rem; margin: 0 0 30px; }

/* Upload box */
.upload-card { background: white !important; border-radius: 16px !important; padding: 0 !important; border: 1px dashed #cbd5e1 !important; box-shadow: 0 4px 20px rgba(0,0,0,0.03) !important; overflow: hidden; }

/* Model Select & Enhance Button */
.model-card { background: white; border-radius: 12px; padding: 20px; border: 1px solid #e2e8f0; box-shadow: 0 4px 6px rgba(0,0,0,0.02); margin-bottom: 20px; }
.model-card-title { font-size: 1rem; font-weight: 600; color: #1e293b; margin-bottom: 12px; }
.model-info-box { background: #f0f9ff; color: #0369a1; padding: 14px; border-radius: 8px; font-size: 0.85rem; margin-top: 12px; display: flex; gap: 12px; align-items: flex-start; line-height: 1.5; }
#enhance-btn { background: #0f9d58 !important; color: white !important; font-weight: 600 !important; font-size: 1.1rem !important; padding: 16px !important; border-radius: 8px !important; border: none !important; width: 100%; transition: all 0.2s; }
#enhance-btn:hover { background: #0b8043 !important; transform: translateY(-1px); }

/* Features List */
.features-card { background: #f8fafc; border-radius: 12px; padding: 24px; border: 1px solid #e2e8f0; height: 100%; display: flex; flex-direction: column; justify-content: center; }
.features-title { font-size: 1.1rem; font-weight: 600; color: #0f172a; margin-bottom: 20px; }
.feature-item { display: flex; align-items: center; gap: 14px; margin-bottom: 16px; color: #334155; font-size: 0.95rem; font-weight: 500; }
.feature-icon { width: 24px; height: 24px; color: #0f9d58; flex-shrink: 0; }

/* Results */
.results-header { font-size: 1.4rem; font-weight: 700; color: #0f172a; margin: 40px 0 20px; }
.image-card { border-radius: 12px !important; overflow: hidden; border: 1px solid #e2e8f0; background: white; padding: 0 !important; box-shadow: 0 4px 12px rgba(0,0,0,0.04); }

/* Buttons */
.download-col { display: flex; flex-direction: column; gap: 12px; padding-top: 36px; }
.btn-primary { background: #0f9d58 !important; color: white !important; padding: 12px 16px !important; border-radius: 8px !important; font-weight: 600 !important; text-align: center; border: none !important; }
.btn-outline { background: white !important; color: #334155 !important; border: 1px solid #cbd5e1 !important; padding: 12px 16px !important; border-radius: 8px !important; font-weight: 500 !important; text-align: center; box-shadow: 0 1px 2px rgba(0,0,0,0.02) !important; }

/* Footer */
.footer-strip { background: #03182b; color: white; padding: 24px 40px; display: flex; justify-content: space-between; align-items: center; font-size: 0.9rem; margin-top: 40px; }
.footer-left { display: flex; gap: 10px; align-items: center; color: #cbd5e1; }
.footer-left strong { color: white; }
.footer-links { display: flex; gap: 24px; color: #94a3b8; align-items: center; }
.footer-links a { color: #94a3b8; text-decoration: none; transition: color 0.2s; }
.footer-links a:hover { color: white; }
"""

# ============================================================
# INFERENCE FUNCTION
# ============================================================

def run_super_resolution(image_path, model_choice, progress=gr.Progress()):
    if image_path is None:
        raise gr.Error("⚠️ Please upload a satellite image before running.")

    model_map = {
        "HAT-SAT (Recommended)": "hatsat",
        "ESRGAN (Baseline)": "esrgan",
        "Sen2SR — WEO-SAS (Sentinel-2 CNN)": "sen2sr",
    }
    selected_model = model_map.get(model_choice)
    if selected_model is None:
        raise gr.Error("⚠️ Please select a valid model.")

    if not isinstance(image_path, list):
        image_path = [image_path]
        
    is_tiff = False
    
    if len(image_path) == 3:
        paths = [str(p) for p in image_path]
        if not all(p.lower().endswith(('.tif', '.tiff')) for p in paths):
            raise gr.Error("⚠️ When uploading multiple files, they must all be TIFF files.")
            
        red_path = next((p for p in paths if "B04" in p.upper()), None)
        green_path = next((p for p in paths if "B03" in p.upper()), None)
        blue_path = next((p for p in paths if "B02" in p.upper()), None)
        
        if not (red_path and green_path and blue_path):
            paths = sorted(paths)
            blue_path, green_path, red_path = paths[0], paths[1], paths[2]
            
        try:
            with rasterio.open(red_path) as src_r, rasterio.open(green_path) as src_g, rasterio.open(blue_path) as src_b:
                profile = src_r.profile.copy()
                profile.update(count=3)
                
                if not (src_r.width == src_g.width == src_b.width and src_r.height == src_g.height == src_b.height):
                    raise gr.Error("⚠️ The uploaded bands have mismatched dimensions.")
                
                r_data = src_r.read(1)
                g_data = src_g.read(1)
                b_data = src_b.read(1)
                
                data = np.stack([r_data, g_data, b_data], axis=-1)
                input_w, input_h = src_r.width, src_r.height
                is_tiff = True
                
        except Exception as e:
            raise gr.Error(f"⚠️ Error reading 3-band TIFFs: {str(e)}")
            
        if data.dtype != np.uint8:
            data = normalize_rgb(data)
            
        original_image = Image.fromarray(data, mode="RGB")
        
    elif len(image_path) == 1:
        file_path = str(image_path[0])
        is_tiff = file_path.lower().endswith(('.tif', '.tiff'))
        
        if is_tiff:
            try:
                with rasterio.open(file_path) as src:
                    profile = src.profile.copy()
                    count = src.count
                    if count not in [1, 3, 4]:
                        raise gr.Error(f"⚠️ Unsupported band count: {count}. Expected 1, 3, or 4 bands.")
                    
                    if count == 3:
                        data = src.read([1, 2, 3])
                        data = np.transpose(data, (1, 2, 0))
                    elif count == 4:
                        data = src.read([1, 2, 3])
                        data = np.transpose(data, (1, 2, 0))
                    else:
                        data = src.read(1)
                        data = np.stack([data, data, data], axis=-1)
                    
                    input_w, input_h = src.width, src.height
            except rasterio.errors.RasterioIOError:
                raise gr.Error("⚠️ Corrupted TIFF or unsupported format.")
            except Exception as e:
                raise gr.Error(f"⚠️ Error reading TIFF: {str(e)}")
    
            if data.dtype == np.uint8:
                pass
            else:
                data = normalize_rgb(data)
    
            original_image = Image.fromarray(data, mode="RGB")
        else:
            try:
                original_image = Image.open(file_path).convert("RGB")
                input_w, input_h = original_image.size
            except Exception:
                raise gr.Error("⚠️ Could not read image.")
    else:
        raise gr.Error("⚠️ Please upload either 1 file (multiband image) or 3 files (B04, B03, B02 GeoTIFFs).")

    if input_w < 16 or input_h < 16:
        raise gr.Error(f"⚠️ Image too small ({input_w}x{input_h}). Please upload an image at least 16x16 pixels.")
    if input_w > 1024 or input_h > 1024:
        raise gr.Error(f"⚠️ Image too large ({input_w}x{input_h}). Please use images up to 1024x1024 pixels.")

    start_time = time.time()
    try:
        progress(0.5, desc=f"Running {selected_model} 4x super-resolution (ZeroGPU) 🚀")
        if selected_model == "hatsat":
            result = infer_hatsat(original_image)
        elif selected_model == "esrgan":
            result = infer_esrgan(original_image)
        else:
            result = infer_sen2sr(original_image)
    except Exception as e:
        raise gr.Error(f"{selected_model} inference failed: {e}")

    progress(0.9, desc="Finalising output 🚀")

    # Generate PNG download
    tmp_png = tempfile.NamedTemporaryFile(suffix=f"_sr_{selected_model}_4x.png", delete=False)
    result.save(tmp_png.name, format="PNG")
    tmp_png.close()
    out_png = tmp_png.name

    # Generate TIFF download if input was TIFF
    if is_tiff:
        tmp_tiff = tempfile.NamedTemporaryFile(suffix=f"_sr_{selected_model}_4x_georef.tif", delete=False)
        tmp_tiff.close()
        save_sr_geotiff(result, profile, tmp_tiff.name, scale=4)
        out_tiff = tmp_tiff.name
    else:
        out_tiff = None

    progress(1.0, desc="Done.")
    
    # Return: original preview, enhanced preview, png download, tiff download
    return original_image, result, gr.update(value=out_png, interactive=True), gr.update(value=out_tiff, interactive=True)


# ============================================================
# HTML COMPONENTS
# ============================================================

TOP_NAV = """
<div class="top-nav">
    <div class="nav-left">
        <svg class="nav-logo-icon" viewBox="0 0 24 24"><path d="M12 2L2 22h20L12 2zm0 4.2l6.8 13.8H5.2L12 6.2z"/><path d="M12 10.5l-2.5 5h5z"/></svg>
        <div class="nav-title">
            <h1 class="nav-logo-text">Senti<span>Enhance</span></h1>
            <p class="nav-subtitle">AI-Powered Super-Resolution for Sentinel-2 Imagery</p>
        </div>
    </div>
    <div class="nav-links">
        <a href="#" class="active">Home</a>
        <a href="#">About</a>
        <a href="#">How It Works</a>
        <a href="#">Use Cases</a>
        <a href="#">Contact</a>
        <a href="#" class="nav-btn">Get Started</a>
    </div>
</div>
"""

SIDEBAR = """
<div class="sidebar-container">
    <div class="sidebar-links">
        <a href="#" class="sidebar-link active">
            <svg fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M3 12l2-2m0 0l7-7 7 7M5 10v10a1 1 0 001 1h3m10-11l2 2m-2-2v10a1 1 0 01-1 1h-3m-6 0a1 1 0 001-1v-4a1 1 0 011-1h2a1 1 0 011 1v4a1 1 0 001 1m-6 0h6"></path></svg>
            Super-Resolution
        </a>
        <a href="#" class="sidebar-link">
            <svg fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M9 19v-6a2 2 0 00-2-2H5a2 2 0 00-2 2v6a2 2 0 002 2h2a2 2 0 002-2zm0 0V9a2 2 0 012-2h2a2 2 0 012 2v10m-6 0a2 2 0 002 2h2a2 2 0 002-2m0 0V5a2 2 0 012-2h2a2 2 0 012 2v14a2 2 0 01-2 2h-2a2 2 0 01-2-2z"></path></svg>
            Compare
        </a>
        <a href="#" class="sidebar-link">
            <svg fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M4 16l4.586-4.586a2 2 0 012.828 0L16 16m-2-2l1.586-1.586a2 2 0 012.828 0L20 14m-6-6h.01M6 20h12a2 2 0 002-2V6a2 2 0 00-2-2H6a2 2 0 00-2 2v12a2 2 0 002 2z"></path></svg>
            Image Gallery
        </a>
        <a href="#" class="sidebar-link">
            <svg fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M9 12h6m-6 4h6m2 5H7a2 2 0 01-2-2V5a2 2 0 012-2h5.586a1 1 0 01.707.293l5.414 5.414a1 1 0 01.293.707V19a2 2 0 01-2 2z"></path></svg>
            Documentation
        </a>
        <a href="#" class="sidebar-link">
            <svg fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M13 16h-1v-4h-1m1-4h.01M21 12a9 9 0 11-18 0 9 9 0 0118 0z"></path></svg>
            About Project
        </a>
    </div>
    <div class="sidebar-banner">
        <svg fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M3.055 11H5a2 2 0 012 2v1a2 2 0 002 2 2 2 0 012 2v2.945M8 3.935V5.5A2.5 2.5 0 0010.5 8h.5a2 2 0 012 2 2 2 0 104 0 2 2 0 012-2h1.064M15 20.488V18a2 2 0 012-2h3.064M21 12a9 9 0 11-18 0 9 9 0 0118 0z"></path></svg>
        <div>Satellite data<br>for a sustainable<br>future</div>
    </div>
</div>
"""

WHY_USE_US = """
<div class="features-title">Why Use Our Platform?</div>
<div class="feature-item">
    <svg class="feature-icon" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M13 10V3L4 14h7v7l9-11h-7z"></path></svg>
    AI-powered super-resolution
</div>
<div class="feature-item">
    <svg class="feature-icon" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M4 7v10c0 2.21 3.582 4 8 4s8-1.79 8-4V7M4 7c0 2.21 3.582 4 8 4s8-1.79 8-4M4 7c0-2.21 3.582-4 8-4s8 1.79 8 4m0 5c0 2.21-3.582 4-8 4s-8-1.79-8-4"></path></svg>
    Supports Sentinel-2 imagery
</div>
<div class="feature-item">
    <svg class="feature-icon" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M17.657 16.657L13.414 20.9a1.998 1.998 0 01-2.827 0l-4.244-4.243a8 8 0 1111.314 0z"></path><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M15 11a3 3 0 11-6 0 3 3 0 016 0z"></path></svg>
    Preserves geospatial metadata
</div>
<div class="feature-item">
    <svg class="feature-icon" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M9.75 17L9 20l-1 1h8l-1-1-.75-3M3 13h18M5 17h14a2 2 0 002-2V5a2 2 0 00-2-2H5a2 2 0 00-2 2v10a2 2 0 002 2z"></path></svg>
    Easy to use web interface
</div>
<div class="feature-item">
    <svg class="feature-icon" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M3.055 11H5a2 2 0 012 2v1a2 2 0 002 2 2 2 0 012 2v2.945M8 3.935V5.5A2.5 2.5 0 0010.5 8h.5a2 2 0 012 2 2 2 0 104 0 2 2 0 012-2h1.064M15 20.488V18a2 2 0 012-2h3.064M21 12a9 9 0 11-18 0 9 9 0 0118 0z"></path></svg>
    Free for research & education
</div>
"""

FOOTER = """
<div class="footer-strip">
    <div class="footer-left">
        <strong>SentiEnhance</strong> | AI for a Sustainable Planet
    </div>
    <div class="footer-links">
        <a href="#">Privacy Policy</a> |
        <a href="#">Terms of Use</a> |
        <a href="#">Contact</a>
        <a href="#"><svg width="18" height="18" fill="currentColor" viewBox="0 0 24 24"><path d="M19 0h-14c-2.761 0-5 2.239-5 5v14c0 2.761 2.239 5 5 5h14c2.762 0 5-2.239 5-5v-14c0-2.761-2.238-5-5-5zm-11 19h-3v-11h3v11zm-1.5-12.268c-.966 0-1.75-.79-1.75-1.764s.784-1.764 1.75-1.764 1.75.79 1.75 1.764-.783 1.764-1.75 1.764zm13.5 12.268h-3v-5.604c0-3.368-4-3.113-4 0v5.604h-3v-11h3v1.765c1.396-2.586 7-2.777 7 2.476v6.759z"/></svg></a>
        <a href="#"><svg width="18" height="18" fill="currentColor" viewBox="0 0 24 24"><path d="M12 0c-6.626 0-12 5.373-12 12 0 5.302 3.438 9.8 8.207 11.387.599.111.793-.261.793-.577v-2.234c-3.338.726-4.033-1.416-4.033-1.416-.546-1.387-1.333-1.756-1.333-1.756-1.089-.745.083-.729.083-.729 1.205.084 1.839 1.237 1.839 1.237 1.07 1.834 2.807 1.304 3.492.997.107-.775.418-1.305.762-1.604-2.665-.305-5.467-1.334-5.467-5.931 0-1.311.469-2.381 1.236-3.221-.124-.303-.535-1.524.117-3.176 0 0 1.008-.322 3.301 1.23.957-.266 1.983-.399 3.003-.404 1.02.005 2.047.138 3.006.404 2.291-1.552 3.297-1.23 3.297-1.23.653 1.653.242 2.874.118 3.176.77.84 1.235 1.911 1.235 3.221 0 4.609-2.807 5.624-5.479 5.921.43.372.823 1.102.823 2.222v3.293c0 .319.192.694.801.576 4.765-1.589 8.199-6.086 8.199-11.386 0-6.627-5.373-12-12-12z"/></svg></a>
    </div>
</div>
"""

# ============================================================
# APP LAYOUT
# ============================================================

with gr.Blocks(theme=gr.themes.Base(), css=CSS) as demo:
    
    gr.HTML(TOP_NAV)
    
    with gr.Row():
        with gr.Column(scale=1, min_width=240, elem_classes="sidebar"):
            gr.HTML(SIDEBAR)
            
        with gr.Column(scale=5, elem_classes="main-wrapper"):
            gr.HTML("""
                <h2 class="page-title">Enhance <span>Sentinel-2</span> Satellite Images</h2>
                <p class="page-subtitle">Upload your Sentinel-2 image (GeoTIFF, TIFF, JP2, or PNG) and use AI to generate a high-resolution version.</p>
            """)
            
            with gr.Row():
                with gr.Column(scale=4, elem_classes="upload-card"):
                    input_image = gr.File(
                        label="Drag & drop your satellite image here",
                        file_count="multiple",
                        file_types=["image", ".tif", ".tiff", ".jp2"]
                    )
                    
                with gr.Column(scale=3):
                    with gr.Group(elem_classes="model-card"):
                        gr.HTML("<div class='model-card-title'>Model Selection</div>")
                        model_selector = gr.Dropdown(
                            choices=[
                                "HAT-SAT (Recommended)",
                                "ESRGAN (Baseline)",
                                "Sen2SR — WEO-SAS (Sentinel-2 CNN)",
                            ],
                            value="HAT-SAT (Recommended)",
                            show_label=False,
                            container=False
                        )
                        gr.HTML("""
                            <div class="model-info-box">
                                <svg width="20" height="20" fill="currentColor" viewBox="0 0 24 24"><path d="M12 2C6.48 2 2 6.48 2 12s4.48 10 10 10 10-4.48 10-10S17.52 2 12 2zm1 15h-2v-6h2v6zm0-8h-2V7h2v2z"/></svg>
                                <span><strong>HAT-SAT:</strong> State-of-the-art transformer model. &nbsp;<strong>ESRGAN:</strong> General-purpose RRDB baseline. &nbsp;<strong>Sen2SR:</strong> WEO-SAS Sentinel-2 CNN — <a href="https://huggingface.co/WEO-SAS/sen2sr" target="_blank" style="color:#0369a1;">HuggingFace ↗</a>.</span>
                            </div>
                        """)
                    enhance_btn = gr.Button("✨ Enhance Image", elem_id="enhance-btn")
                    gr.HTML("<div style='text-align:center; font-size:0.85rem; color:#94a3b8; margin-top:8px;'>Processing may take a few moments...</div>")
                    
                with gr.Column(scale=3, elem_classes="features-card"):
                    gr.HTML(WHY_USE_US)
            
            # Results Section
            gr.HTML("<div class='results-header'>Results</div>")
            
            with gr.Row():
                with gr.Column(scale=4, elem_classes="image-card"):
                    original_image = gr.Image(label="Original (Sentinel-2, 10m)", type="pil", interactive=False)
                    
                with gr.Column(scale=4, elem_classes="image-card"):
                    enhanced_image = gr.Image(label="Enhanced (HAT-SAT, 2x)", type="pil", interactive=False)
                    
                with gr.Column(scale=3, elem_classes="download-col"):
                    gr.HTML("<div style='font-weight:700; font-size:1.05rem; color:#0f172a; margin-bottom:4px;'>Download & Share</div>")
                    
                    download_png = gr.DownloadButton("📥 Download Enhanced Image", elem_classes="btn-primary")
                    download_tiff = gr.DownloadButton("📄 Download as GeoTIFF", elem_classes="btn-outline")
                    compare_btn = gr.Button("🔍 Compare Images", elem_classes="btn-outline")
                    share_btn = gr.Button("🔗 Share Result", elem_classes="btn-outline")

    gr.HTML(FOOTER)

    # ============================================================
    # EVENTS
    # ============================================================
    
    enhance_btn.click(
        fn=run_super_resolution,
        inputs=[input_image, model_selector],
        outputs=[original_image, enhanced_image, download_png, download_tiff],
        show_progress="full",
    )

    # Interactive elements that don't do anything functionally but can show alerts
    def show_alert():
        gr.Info("This feature is coming soon!")
    compare_btn.click(fn=show_alert)
    share_btn.click(fn=show_alert)


# ============================================================
# LAUNCH
# ============================================================

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