from __future__ import annotations

import gradio as gr
import spaces

from geochat.inference.engine import GeoChatEngine
from geochat.inference.sar import load_image_input, sar1_to_rgb


# ------------------------------------------------------------
# Load GeoChat once at Space startup.
# ------------------------------------------------------------

print("[Akasha] Loading GeoChat...")

engine = GeoChatEngine()

print("[Akasha] GeoChat ready.")


# ------------------------------------------------------------
# Gradio inference function
# ------------------------------------------------------------

@spaces.GPU(duration=60)
def geochat(
    image: str,
    prompt: str,
    max_new_tokens: int,
    sar_vh: str | None = None,
) -> str:

    if image is None:
        raise ValueError("Image is required.")

    try:
        image = sar1_to_rgb(image, sar_vh) if sar_vh else load_image_input(image)
    except ValueError:
        raise
    except Exception as exc:
        raise ValueError("Unable to process the supplied image.") from exc

    if not prompt or not prompt.strip():
        raise ValueError("Prompt must not be empty.")

    max_new_tokens = int(max_new_tokens)

    if not 1 <= max_new_tokens <= 512:
        raise ValueError(
            "max_new_tokens must be between 1 and 512."
        )

    return engine.generate(
        image=image,
        prompt=prompt,
        max_new_tokens=max_new_tokens,
    )


# ------------------------------------------------------------
# Gradio application
# ------------------------------------------------------------

demo = gr.Interface(
    fn=geochat,
    inputs=[
        gr.File(
            type="filepath",
            label="Satellite Image / TIFF / GeoTIFF",
        ),
        gr.Textbox(
            label="Prompt",
            value="Describe the visible features in this satellite image.",
            lines=3,
        ),
        gr.Slider(
            minimum=1,
            maximum=512,
            value=128,
            step=1,
            label="Max New Tokens",
        ),
        gr.File(
            type="filepath",
            label="Optional SAR VH (pair with VV in first input)",
        ),
    ],
    outputs=gr.Textbox(
        label="GeoChat Response",
        lines=8,
    ),
    title="Akasha GeoChat",
    description="GeoChat specialist for satellite image analysis.",
    api_name="geochat",
)


# ------------------------------------------------------------
# Launch
# ------------------------------------------------------------

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