from functools import lru_cache

import numpy as np
import gradio as gr
from PIL import Image, ImageDraw
from paddleocr import PaddleOCR
import bbox_align

# Official PP-OCRv6 tiers. A single selector chooses the matching
# detection + recognition pair.
V6_MODELS = {
    "Tiny": ("PP-OCRv6_tiny_det", "PP-OCRv6_tiny_rec"),
    "Small": ("PP-OCRv6_small_det", "PP-OCRv6_small_rec"),
    "Medium": ("PP-OCRv6_medium_det", "PP-OCRv6_medium_rec"),
}


@lru_cache(maxsize=len(V6_MODELS))
def get_ocr(tier):
    det_model, rec_model = V6_MODELS[tier]
    return PaddleOCR(
        text_detection_model_name=det_model,
        text_recognition_model_name=rec_model,
        use_doc_orientation_classify=False,
        use_doc_unwarping=False,
        use_textline_orientation=False,
        device="cpu",
        enable_mkldnn=True,
        cpu_threads=2,
    )


def _result_dict(result_obj):
    """Normalize PaddleOCR's result object to its inner `res` dictionary."""
    if isinstance(result_obj, dict):
        data = result_obj
    elif hasattr(result_obj, "json"):
        data = result_obj.json
        if callable(data):
            data = data()
    elif hasattr(result_obj, "res"):
        return result_obj.res
    else:
        raise TypeError(f"Unsupported PaddleOCR result type: {type(result_obj)!r}")

    return data.get("res", data)


def _group_lines(polys, width, height):
    vertices = [
        [(float(x), float(y)) for x, y in np.asarray(poly).reshape(-1, 2)]
        for poly in polys
    ]
    page_boundary = [
        (0.0, 0.0),
        (float(width), 0.0),
        (float(width), float(height)),
        (0.0, float(height)),
    ]

    groups = bbox_align.process(vertices, page_boundary)

    groups = [
        sorted(group, key=lambda i: min(p[0] for p in vertices[i]))
        for group in groups
    ]
    groups.sort(
        key=lambda group: min(min(p[1] for p in vertices[i]) for i in group)
    )
    return groups, vertices


def _draw_groups(image, groups, vertices):
    out = image.copy().convert("RGB")
    draw = ImageDraw.Draw(out)

    # Thin boxes = OCR detections.
    for poly in vertices:
        draw.line(poly + [poly[0]], fill="red", width=2)

    # Thick boxes = reconstructed logical lines.
    for line_no, group in enumerate(groups, start=1):
        pts = [p for i in group for p in vertices[i]]
        xs = [p[0] for p in pts]
        ys = [p[1] for p in pts]
        x1, y1, x2, y2 = min(xs), min(ys), max(xs), max(ys)
        draw.rectangle((x1, y1, x2, y2), outline="blue", width=4)
        draw.text((x1 + 3, max(0, y1 - 14)), str(line_no), fill="blue")

    return out


def run_receipt(image, tier):
    if image is None:
        return None, "Upload a receipt image."

    image = image.convert("RGB")
    ocr = get_ocr(tier)
    results = list(ocr.predict(np.asarray(image)))
    if not results:
        return image, "No text found."

    res = _result_dict(results[0])
    texts = list(res.get("rec_texts", []))
    scores = list(res.get("rec_scores", []))
    polys = list(res.get("rec_polys", []))

    items = []
    for i, (text, poly) in enumerate(zip(texts, polys)):
        text = str(text).strip()
        if not text:
            continue
        score = float(scores[i]) if i < len(scores) else 0.0
        items.append((text, score, poly))

    if not items:
        return image, "No recognized text found."

    texts = [x[0] for x in items]
    polys = [x[2] for x in items]

    groups, vertices = _group_lines(polys, image.width, image.height)
    annotated = _draw_groups(image, groups, vertices)

    grouped_text = "\n".join(
        f"{line_no:02d}: " + " | ".join(texts[i] for i in group)
        for line_no, group in enumerate(groups, start=1)
    )

    return annotated, grouped_text


with gr.Blocks(title="PP-OCRv6 Receipt Line Grouping") as demo:
    gr.Markdown(
        "# Receipt OCR · PP-OCRv6 + bbox-align\n"
        "Upload a receipt, choose a PP-OCRv6 tier, and inspect the reconstructed logical lines."
    )

    model_tier = gr.Dropdown(
        choices=list(V6_MODELS.keys()),
        value="Small",
        label="PP-OCRv6 model",
    )

    with gr.Row():
        input_image = gr.Image(type="pil", label="Receipt")
        output_image = gr.Image(type="pil", label="Grouped lines")

    output_text = gr.Textbox(
        label="Grouped text",
        lines=16,
        buttons=["copy"],
    )

    run_button = gr.Button("Run OCR", variant="primary")
    run_button.click(
        run_receipt,
        inputs=[input_image, model_tier],
        outputs=[output_image, output_text],
    )


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