import gradio as gr
from pathlib import Path
from datetime import datetime
from PIL import Image

from ocrpkg.core import (
    ocr_file,                   # EasyOCR for PDFs/images by default
    run_ocr_on_image,           # EasyOCR for images
    run_ocr_on_image_paddle,    # PaddleOCR for images
    save_json,
)
from ocrpkg.pdf import pdf_to_images
from ocrpkg.schema import OCRBlock


def _make_preview_text(blocks):
    pages = sorted({b.page for b in blocks})
    chunks = []
    for pg in pages:
        page_text = "\n".join([b.text for b in blocks if b.page == pg])
        chunks.append(f"--- Page {pg} ---\n{page_text}")
    return "\n\n".join(chunks) if chunks else "No text detected."


def _save_json_artifact(blocks, base_stem: str) -> str:
    out_dir = Path("examples/output")
    out_dir.mkdir(parents=True, exist_ok=True)
    json_path = out_dir / f"{base_stem}.json"
    save_json(blocks, json_path)
    return str(json_path)


# -------- PDF first-page preview --------
def pdf_first_page_preview(file_obj):
    if file_obj is None:
        return None
    p = Path(file_obj.name)
    if p.suffix.lower() != ".pdf":
        return None
    try:
        imgs = pdf_to_images(p, dpi=150)
        return imgs[0] if imgs else None
    except Exception:
        return None


# -------- Handlers --------
def run_ocr_image(pil_image: Image.Image, langs: str, conf_thr: float, engine: str):
    if pil_image is None:
        return "No image uploaded", None

    conf = float(conf_thr)
    langs_list = [s.strip() for s in langs.split(",") if s.strip()]

    if engine == "PaddleOCR":
        primary_lang = langs_list[0] if langs_list else "en"
        blocks = run_ocr_on_image_paddle(pil_image.convert("RGB"), lang=primary_lang, conf_threshold=conf, page=1)
    else:
        blocks = run_ocr_on_image(pil_image.convert("RGB"), langs=langs_list or ["en"], conf_threshold=conf, page=1)

    preview_text = _make_preview_text(blocks)
    base = f"image_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}"
    json_path = _save_json_artifact(blocks, base)
    return preview_text, json_path


def run_ocr_pdf(file_obj, langs: str, dpi: int, conf_thr: float, engine: str):
    if file_obj is None:
        return "No PDF uploaded", None
    in_path = Path(file_obj.name)
    if in_path.suffix.lower() != ".pdf":
        return "Please upload a PDF file.", None

    conf = float(conf_thr)
    langs_list = [s.strip() for s in langs.split(",") if s.strip()] or ["en"]

    if engine == "PaddleOCR":
        # Rasterize PDF pages and run Paddle per page
        pages = pdf_to_images(in_path, dpi=int(dpi))
        all_blocks = []
        for i, img in enumerate(pages, start=1):
            # use first language only (Paddle supports one lang model at a time)
            blocks = run_ocr_on_image_paddle(img, lang=langs_list[0], conf_threshold=conf, page=i)
            all_blocks.extend(blocks)
        blocks = all_blocks
    else:
        # EasyOCR path (your default ocr_file)
        blocks = ocr_file(in_path, langs=langs_list, dpi=int(dpi), conf_threshold=conf)

    preview_text = _make_preview_text(blocks)
    json_path = _save_json_artifact(blocks, in_path.stem)
    return preview_text, json_path


# -------- UI --------
with gr.Blocks() as demo:
    gr.Markdown("# 📄 OCR & Document AI (EasyOCR + PaddleOCR)")
    gr.Markdown(
        "Images show a live preview in the upload area. PDFs show a first-page preview. "
        "Choose **PaddleOCR** for higher quality (CPU), or **EasyOCR** as fallback. "
        "Download structured JSON results."
    )

    with gr.Tabs():
        # ----- Image tab -----
        with gr.Tab("Image"):
            with gr.Row():
                with gr.Column():
                    img_in = gr.Image(type="pil", label="Upload image")
                    langs_img = gr.Textbox(value="en", label="Languages (comma-separated; EasyOCR uses this list)")
                    engine_img = gr.Dropdown(["PaddleOCR", "EasyOCR"], value="PaddleOCR", label="Engine")
                    conf_img = gr.Slider(0.0, 1.0, value=0.3, step=0.05, label="Confidence threshold")
                    btn_img = gr.Button("Run OCR on Image", variant="primary")
                with gr.Column():
                    txt_img = gr.Textbox(label="Recognized text (by page)", lines=16)
                    json_img = gr.File(label="Download JSON")

            btn_img.click(
                fn=run_ocr_image,
                inputs=[img_in, langs_img, conf_img, engine_img],
                outputs=[txt_img, json_img],
            )

        # ----- PDF tab -----
        with gr.Tab("PDF"):
            with gr.Row():
                with gr.Column():
                    pdf_in = gr.File(label="Upload PDF", file_types=[".pdf"])
                    pdf_preview = gr.Image(label="First Page Preview (auto)", interactive=False)
                    langs_pdf = gr.Textbox(value="en", label="Languages (comma-separated)")
                    engine_pdf = gr.Dropdown(["PaddleOCR", "EasyOCR"], value="PaddleOCR", label="Engine")
                    dpi_pdf = gr.Slider(100, 400, value=250, step=50, label="PDF render DPI")
                    conf_pdf = gr.Slider(0.0, 1.0, value=0.3, step=0.05, label="Confidence threshold")
                    btn_pdf = gr.Button("Run OCR on PDF", variant="primary")
                with gr.Column():
                    txt_pdf = gr.Textbox(label="Recognized text (by page)", lines=16)
                    json_pdf = gr.File(label="Download JSON")

            # Auto preview first page
            pdf_in.upload(fn=pdf_first_page_preview, inputs=pdf_in, outputs=pdf_preview)

            btn_pdf.click(
                fn=run_ocr_pdf,
                inputs=[pdf_in, langs_pdf, dpi_pdf, conf_pdf, engine_pdf],
                outputs=[txt_pdf, json_pdf],
            )

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