"""PP-OCRv6 medium text detection demo.

Loads the PaddlePaddle/PP-OCRv6_medium_det model and runs text-line
bounding-box detection on an uploaded image.  Returns the image with
detected boxes drawn on top, plus a JSON summary of the boxes.
"""

import json
import os
import tempfile

# Disable PIR + oneDNN to avoid a known PaddlePaddle 3.3.x bug:
#   "ConvertPirAttribute2RuntimeAttribute not support [pir::ArrayAttribute<pir::DoubleAttribute>]"
os.environ.setdefault("FLAGS_enable_pir_api", "0")
os.environ.setdefault("FLAGS_use_mkldnn", "0")

import gradio as gr
from PIL import Image

from paddleocr import TextDetection

MODEL_ID = "PP-OCRv6_medium_det"

# Load the model at module scope so it's ready for the first request.
print(f"Loading {MODEL_ID} ...")
detector = TextDetection(model_name=MODEL_ID, enable_mkldnn=False)
print(f"{MODEL_ID} loaded.")


def detect_text(image: str) -> tuple:
    """Detect text regions in an image and return annotated output.

    Args:
        image: Path to the input image (filepath from Gradio).

    Returns:
        A tuple of (annotated_image, json_summary) where annotated_image
        is the input image with detected bounding boxes drawn on it, and
        json_summary is a JSON string listing each detected region's
        polygon and confidence score.
    """
    if image is None:
        return None, "Please provide an image."

    # Run detection
    results = detector.predict(input=image, batch_size=1)

    all_boxes = []
    annotated_path = None

    for res in results:
        # Save the visualization image
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            res.save_to_img(save_path=tmp.name)
            annotated_path = tmp.name

        # Extract bounding polygons and confidence scores
        if hasattr(res, "json") and res.json:
            data = res.json
            if isinstance(data, dict) and "res" in data:
                det_data = data["res"]
                if isinstance(det_data, dict) and "dt_polys" in det_data:
                    polys = det_data["dt_polys"]
                    scores = det_data.get("dt_scores", [])
                    for i, poly in enumerate(polys):
                        score = float(scores[i]) if i < len(scores) else 0.0
                        all_boxes.append({
                            "polygon": [list(p) for p in poly],
                            "confidence": round(score, 4),
                        })

    summary = {
        "num_text_regions": len(all_boxes),
        "regions": all_boxes,
    }

    annotated = Image.open(annotated_path) if annotated_path else None
    if annotated_path and os.path.exists(annotated_path):
        os.unlink(annotated_path)

    return annotated, json.dumps(summary, indent=2, ensure_ascii=False)


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

with gr.Blocks() as demo:
    gr.Markdown("# PP-OCRv6 Medium Text Detection")
    gr.Markdown(
        "Detect text regions in images using "
        "[PaddlePaddle/PP-OCRv6_medium_det](https://huggingface.co/PaddlePaddle/PP-OCRv6_medium_det) — "
        "a lightweight 15.5M-parameter OCR detection model from the PaddleOCR team."
    )

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="Input image",
                type="filepath",
                height=400,
            )
            run_btn = gr.Button("Detect text", variant="primary")
        with gr.Column():
            output_image = gr.Image(
                label="Detected text regions",
                type="pil",
                height=400,
            )

    output_json = gr.Code(
        label="Detection summary (JSON)",
        language="json",
        lines=12,
    )

    run_btn.click(
        fn=detect_text,
        inputs=[input_image],
        outputs=[output_image, output_json],
        api_name="detect_text",
    )

    with gr.Accordion("Advanced settings", open=False):
        gr.Markdown(
            "This demo uses the detection-only model, which identifies bounding "
            "boxes of text regions but does not perform character recognition. "
            "For full OCR (detection + recognition), see the "
            "[PaddleOCR pipeline](https://github.com/PaddlePaddle/PaddleOCR)."
        )

    gr.Examples(
        examples=[
            ["example1.png"],
            ["example_sign.jpg"],
            ["example_noodles.jpg"],
        ],
        inputs=[input_image],
        outputs=[output_image, output_json],
        fn=detect_text,
        cache_examples=True,
        cache_mode="lazy",
    )

    gr.Markdown(
        "---\n"
        "**Model**: [PaddlePaddle/PP-OCRv6_medium_det](https://huggingface.co/PaddlePaddle/PP-OCRv6_medium_det) "
        "| **License**: Apache-2.0 | **Framework**: PaddlePaddle"
    )


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