import os

import spaces  # Initialize ZeroGPU before importing torch or the model.
import gradio as gr
import torch
from PIL import Image

from openhtr.inference import Recognizer


# ZeroGPU exposes a CUDA device at startup to register the model for GPU workers.
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cpu":
    torch.set_num_threads(min(2, os.cpu_count() or 1))
recognizer = Recognizer(device)


@spaces.GPU(duration=15)
def recognize(image: Image.Image | None) -> str:
    if image is None:
        raise gr.Error("Upload or paste a handwritten line image first.")
    if not isinstance(image, Image.Image):
        raise gr.Error("Upload an image in a supported format such as PNG, JPEG, or WebP.")
    text = recognizer.recognize(image)
    if not text:
        gr.Warning("No text was recognized. Try a clearer crop of a single line.")
    return text


with gr.Blocks(title="OpenHTR Chinese Handwriting Recognition", analytics_enabled=False) as demo:
    gr.Markdown("# OpenHTR-Rec\nUpload or paste a cropped image of a single line of Chinese handwriting.")
    with gr.Row():
        image = gr.Image(
            type="pil", image_mode=None, sources=["upload", "clipboard"], format="png",
            label="Handwritten Line Image", height=280,
        )
        output = gr.Textbox(label="Transcription", interactive=False, lines=9, buttons=["copy"])
    with gr.Row():
        submit = gr.Button("Recognize", variant="primary")
        gr.ClearButton([image, output], value="Clear")

    submit.click(recognize, inputs=image, outputs=output, api_name="recognize")

demo.queue(default_concurrency_limit=1, max_size=16)

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