"""MathSnap deblur Space (Gradio on ZeroGPU). Only callers holding DEBLUR_KEY (the Supabase function) may use it."""
import base64
import binascii
import hmac
import os

import gradio as gr
import spaces

import deblur

MAX_BASE64 = 11_000_000  # ~8 MB JPEG


@spaces.GPU(duration=30)
def restore(rgb):
    return deblur.deblur_array(rgb)


def run(image_base64, key):
    expected = os.environ.get("DEBLUR_KEY", "")
    # The key is checked before any GPU time is requested, so strangers cannot spend the quota.
    if len(expected) < 32 or not hmac.compare_digest((key or "").encode(), expected.encode()):
        raise gr.Error("Forbidden")
    if not image_base64 or len(image_base64) > MAX_BASE64:
        raise gr.Error("Invalid image")
    try:
        rgb = deblur.decode(base64.b64decode(image_base64, validate=True))
    except (binascii.Error, ValueError, OSError):
        raise gr.Error("Invalid image")
    return base64.b64encode(deblur.encode(restore(rgb))).decode()


with gr.Blocks(title="MathSnap Deblur") as demo:
    gr.Markdown("Private endpoint for the MathSnap app. Requests without the app's key are refused.")
    image = gr.Textbox(label="JPEG (base64)")
    key = gr.Textbox(label="Key", type="password")
    result = gr.Textbox(label="Result (base64)")
    gr.Button("Deblur").click(run, [image, key], result, api_name="deblur")

demo.queue(default_concurrency_limit=1).launch()