"""PP-OCRv6 medium recognition model demo.

Upload a cropped text-line image and get recognized text with a confidence score.
PP-OCRv6_medium_rec is a 19M-parameter text recognition model from PaddlePaddle.
"""

import json
import os
import tempfile

# On ZeroGPU, `import spaces` must come before any torch/CUDA import.
# PaddlePaddle doesn't use torch, but the ZeroGPU runtime still requires
# this import and a @spaces.GPU decorator on the inference function.
import spaces

import gradio as gr
from paddleocr import TextRecognition

# Load model at module scope — PaddlePaddle runs on CPU natively.
MODEL_NAME = "PP-OCRv6_medium_rec"
rec_model = TextRecognition(model_name=MODEL_NAME)


@spaces.GPU(duration=30)
def recognize_text(image):
    """Recognize text from a cropped text-line image.

    Args:
        image: A file path to an image containing a single line of text.

    Returns:
        A tuple of (recognized_text, confidence_score, raw_json).
    """
    if image is None:
        return "No image provided", 0.0, ""

    img_path = image if isinstance(image, str) else None
    if img_path is None:
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
            image.save(f, format="PNG")
            img_path = f.name

    try:
        output = rec_model.predict(input=img_path, batch_size=1)
        results = []
        for res in output:
            raw_json = getattr(res, "json", None)
            if isinstance(raw_json, dict):
                data = raw_json
            elif isinstance(raw_json, str):
                data = json.loads(raw_json)
            elif hasattr(res, "to_dict"):
                data = res.to_dict()
            else:
                data = {
                    "rec_text": getattr(res, "rec_text", str(res)),
                    "rec_score": getattr(res, "rec_score", 0.0),
                }
            results.append(data)

        if not results:
            return "No text recognized", 0.0, ""

        # Results may be nested under a "res" key
        first = results[0]
        if "res" in first:
            first = first["res"]
        text = first.get("rec_text", "")
        score = first.get("rec_score", 0.0)
        raw = json.dumps(results, ensure_ascii=False, indent=2)
        return text, float(score), raw
    except Exception as e:
        return f"Error: {e}", 0.0, ""
    finally:
        if not isinstance(image, str) and os.path.exists(img_path):
            os.unlink(img_path)


# --- UI ---

CSS = """
#col-container { max-width: 900px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# PP-OCRv6 Medium Text Recognition")
        gr.Markdown(
            "Upload a cropped image of a single text line to recognize text. "
            "This is the [PP-OCRv6_medium_rec](https://huggingface.co/PaddlePaddle/PP-OCRv6_medium_rec) "
            "model — a 19M-parameter OCR recognition model from PaddlePaddle that "
            "supports 50 languages."
        )

        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(
                    label="Text-line image",
                    type="filepath",
                    height=200,
                )
                run_btn = gr.Button("Recognize Text", variant="primary")

            with gr.Column(scale=1):
                text_output = gr.Textbox(label="Recognized text", interactive=False)
                score_output = gr.Number(label="Confidence score", precision=4, interactive=False)
                raw_output = gr.Code(label="Raw JSON output", language="json")

        gr.Examples(
            examples=[
                ["examples/sample_readme.jpeg"],
                ["examples/sample_text.png"],
                ["examples/sample_text2.png"],
                ["examples/sample_text3.png"],
                ["examples/sample_text4.png"],
            ],
            inputs=input_image,
            outputs=[text_output, score_output, raw_output],
            fn=recognize_text,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        fn=recognize_text,
        inputs=input_image,
        outputs=[text_output, score_output, raw_output],
        api_name="recognize",
    )


if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)