"""
KYC OCR — Gradio app for Hugging Face Spaces.

Upload a PAN card, cheque, Aadhaar, or passport (PDF or image) and get
structured field extraction. Runs on CPU using RapidOCR (ONNX).
"""
# ── Patch a known gradio_client bug ──────────────────────────────────────
# gradio_client.utils.get_type assumes schema is always a dict, but JSON
# Schema's `additionalProperties` field can legitimately be a bool. When
# Gradio's /info endpoint runs (HF Spaces health checks ping it), the
# function crashes with: "TypeError: argument of type 'bool' is not iterable".
# Fixed upstream in newer releases but still present in many 5.x. Patch it.
import gradio_client.utils as _gcu
_orig_get_type = _gcu.get_type
def _patched_get_type(schema):
    if not isinstance(schema, dict):
        return "Any"
    return _orig_get_type(schema)
_gcu.get_type = _patched_get_type
# ─────────────────────────────────────────────────────────────────────────

import gradio as gr
import time
import traceback
from pathlib import Path

from kyc_extract import process, RapidOCR

# ── Initialise OCR engine once at startup ─────────────────────────────
print("Loading RapidOCR engine (ONNX models)…")
_t0 = time.time()
ENGINE = RapidOCR()
print(f"Ready. Engine init took {time.time() - _t0:.1f}s")


# ── Pretty-format the extracted fields as Markdown ────────────────────
def _fmt_fields(fields: dict) -> str:
    if not fields:
        return "_No fields extracted._"
    lines = []
    for k, v in fields.items():
        label = k.replace("_", " ").title()
        if isinstance(v, bool):
            v = "✅ Yes" if v else "❌ No"
        elif v is None:
            v = "_(not found)_"
        lines.append(f"**{label}:** {v}")
    return "\n\n".join(lines)


def _fmt_info(result: dict) -> str:
    bits = [
        f"**Detected type:** `{result.get('detected_type', 'unknown').upper()}`",
        f"**OCR time:** {result.get('ocr_time_s')}s",
        f"**Rotation applied:** {result.get('rotation_applied', 0)}°",
        f"**OCR quality score:** {result.get('ocr_score')}",
    ]
    return "  •  ".join(bits)


# ── Main handler ──────────────────────────────────────────────────────
def extract(file):
    if file is None:
        return "_Please upload a file._", "", ""

    file_path = file if isinstance(file, str) else file.name

    try:
        pages = process(file_path, ENGINE)
    except Exception as exc:
        return (
            f"❌ **Error:** {exc}\n\n```\n{traceback.format_exc()}\n```",
            "",
            "",
        )

    if not pages:
        return "_No pages processed._", "", ""

    # Render all pages (cheques and PANs are normally single-page; multi-page
    # PDFs of compound KYC packets get each page rendered separately).
    info_blocks, field_blocks, raw_blocks = [], [], []
    for p in pages:
        if "error" in p:
            info_blocks.append(f"### Page {p['page']}\n_{p['error']}_")
            continue
        info_blocks.append(f"### Page {p['page']}\n{_fmt_info(p)}")
        field_blocks.append(f"### Page {p['page']}\n{_fmt_fields(p['fields'])}")
        raw_blocks.append(
            f"━━━ Page {p['page']} ━━━\n"
            + "\n".join(f"[{c:.2f}] {t}" for c, t in p["_raw_lines"])
        )

    return (
        "\n\n".join(info_blocks),
        "\n\n---\n\n".join(field_blocks) if field_blocks else "_No fields extracted._",
        "\n\n".join(raw_blocks),
    )


# ── Gradio UI ─────────────────────────────────────────────────────────
DESCRIPTION = """
# 🇮🇳 KYC Document OCR

Extract structured fields from Indian KYC documents — **PAN card**, **cheque**,
**Aadhaar**, and **passport** — from images or PDFs. Runs entirely on CPU using
[RapidOCR](https://github.com/RapidAI/RapidOCR) (PaddleOCR PP-OCRv4 models in ONNX).

**Supported documents and fields:**
- **PAN:** PAN number, name, father's name, DOB / date of incorporation, entity type
- **Cheque:** bank name, address, IFSC, account number, MICR code, cheque number, transaction code
- **Aadhaar / Passport:** _coming soon — Aadhaar reads best from QR code (`pyzbar`), passport from MRZ (`passporteye`)_

**Privacy note:** this is a public Space. Don't upload real personal documents.
For production use, fork it and run on private infrastructure.
"""

with gr.Blocks(title="KYC OCR") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(
                label="Upload PAN / Cheque (PDF or image)",
                file_types=[".pdf", ".png", ".jpg", ".jpeg", ".webp"],
                type="filepath",
            )
            submit_btn = gr.Button("Extract fields", variant="primary", size="lg")
            # Use an absolute path; skip the Examples block if the example file
            # wasn't uploaded (won't crash the whole app).
            _example_path = Path(__file__).parent / "examples" / "synthetic_pan.png"
            if _example_path.exists():
                gr.Examples(
                    examples=[[str(_example_path)]],
                    inputs=file_input,
                    label="Example (synthetic PAN with fake data)",
                )

        with gr.Column(scale=2):
            info_out = gr.Markdown()
            fields_out = gr.Markdown()

    with gr.Accordion("Raw OCR output (for debugging)", open=False):
        raw_out = gr.Textbox(label="Raw OCR lines with confidence scores",
                             lines=20)

    gr.Markdown(
        "---\n"
        "Pipeline source: `kyc_extract.py`. "
        "All extraction is rule-based — no hardcoded names/banks. "
        "Doc-type detection → optional rotation correction (cheques) → "
        "format-spec regex + heuristics + checksum validation."
    )

    submit_btn.click(
        extract,
        inputs=file_input,
        outputs=[info_out, fields_out, raw_out],
    )
    file_input.change(
        extract,
        inputs=file_input,
        outputs=[info_out, fields_out, raw_out],
    )


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