import zipfile
from pathlib import Path

import gradio as gr

from model import extract, generate_template

templates_directory = Path(__file__).parent / "data" / "templates"
structured_json_templates = {
    "Basic receipt": (templates_directory / "basic-receipt.json").read_text(
        encoding="utf-8"
    ),
    "Invoice with line items": (
        templates_directory / "invoice-with-line-items.json"
    ).read_text(encoding="utf-8"),
    "Basic bank statement": (
        templates_directory / "basic-bank-statement.json"
    ).read_text(encoding="utf-8"),
    "Advanced bank statement": (
        templates_directory / "advanced-bank-statement.json"
    ).read_text(encoding="utf-8"),
    "Model task catalog": (
        templates_directory / "model-task-catalog.json"
    ).read_text(encoding="utf-8"),
    "Todo list": (templates_directory / "todo-list.json").read_text(
        encoding="utf-8"
    ),
}

default_structured_json_template = "Invoice with line items"


def download_agentskill(include_all_skills):
    skill_folders = (
        sorted(path for path in Path(".agents/skills").iterdir() if path.is_dir())
        if include_all_skills
        else [Path(".agents/skills/image-data-extractor")]
    )
    for skill_folder in skill_folders:
        with zipfile.ZipFile(
            f"{skill_folder.name}.zip", "w", compression=zipfile.ZIP_DEFLATED
        ) as archive:
            for path in skill_folder.rglob("*"):
                archive.write(path, path.relative_to(".agents/skills"))
    return [f"{skill_folder.name}.zip" for skill_folder in skill_folders]

with gr.Blocks(title="Image Data Extractor") as demo:
    gr.Markdown(
        "# Image Data Extractor\n"
        "Extract structured JSON from an image and optional text using "
        "[numind/NuExtract3](https://huggingface.co/numind/NuExtract3)."
    )
    with gr.Row():
        with gr.Column():
            image = gr.Image(label="Document image", type="pil", height=430)
            text = gr.Textbox(
                label="Document text (optional)",
                placeholder="Add text to process alongside the image",
                lines=3,
            )
            enable_thinking = gr.Checkbox(label="Enable thinking")
            template_preset = gr.Dropdown(
                choices=list(structured_json_templates),
                value=default_structured_json_template,
                label="JSON template preset",
                info="Choose a starting structure, then edit it below.",
            )
            generate_template_button = gr.Button("Generate template from Image")
            with gr.Accordion(
                "Structured JSON template", open=False
            ) as template_accordion:
                template = gr.Textbox(
                    label="Structured JSON template",
                    value=structured_json_templates[default_structured_json_template],
                    lines=18,
                    info="Used for structured extraction. Edit the example to match your document.",
                )
            run = gr.Button("Extract Image Data", variant="primary")
        with gr.Column():
            structured_output = gr.JSON(
                label="Structured JSON",
                open=True,
                show_indices=True,
            )
            gr.Markdown(
                "## Install the agent skill\n"
                "Download the ZIP, extract it, and place its skill folders in your "
                "project's `.agents/skills/` directory. Agents should read the "
                "`SKILL.md` and use it whenever a task requires structured extraction "
                "from an image. Keep **Include All Skills** checked to download every "
                "skill as its own ZIP file."
            )
            include_all_skills = gr.Checkbox(
                label="Include All Skills",
                value=True,
            )
            download_agentskill_button = gr.Button("Download agent skill")
            download_agentskill_file = gr.File(
                label="Agent skill ZIPs",
                file_count="multiple",
            )

    gr.Examples(
        examples=[
            [
                preset,
                str(Path(__file__).parent / "data" / "samples" / filename),
                "",
                False,
            ]
            for preset, filename in [
                ("Basic receipt", "receipt-ocr-original.webp"),
                ("Invoice with line items", "invoice-with-items.png"),
                ("Advanced bank statement", "BankStatementChequing.png"),
                ("Basic bank statement", "bank statement blog image.webp"),
                ("Model task catalog", "tags.png"),
                ("Todo list", "task-list.png"),
            ]
        ],
        inputs=[template_preset, image, text, enable_thinking],
    )

    template_preset.change(
        structured_json_templates.__getitem__,
        inputs=template_preset,
        outputs=template,
        api_name="select_structured_json_template",
    )

    download_agentskill_button.click(
        download_agentskill,
        inputs=include_all_skills,
        outputs=download_agentskill_file,
        api_name="download_agentskill",
    )

    generate_template_button.click(
        lambda: gr.Accordion(open=True),
        outputs=template_accordion,
        api_visibility="private",
        show_progress="hidden",
    ).then(
        generate_template,
        inputs=image,
        outputs=template,
        api_name="generate_template",
    )

    run.click(
        extract,
        inputs=[image, text, template, enable_thinking],
        outputs=structured_output,
        api_name="extract",
    )


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