from pathlib import Path

import gradio as gr

from model import extract, generate_template

structured_json_templates = {
    "Basic receipt": """{
  "merchant": "verbatim-string",
  "purchase_date": "date",
  "purchase_time": "time",
  "total": "number",
  "currency": "currency",
  "item_count": "integer",
  "paid": "boolean",
  "payment_method": ["cash", "credit-card", "debit-card", "other"],
  "items": ["verbatim-string"]
}""",
    "Advanced receipt": """{
  "merchant": {
    "name": "verbatim-string",
    "address": "verbatim-string",
    "phone": "phone-number"
  },
  "purchase_date": "date",
  "purchase_time": "time",
  "store_number": "verbatim-string",
  "operator_number": "verbatim-string",
  "terminal_number": "verbatim-string",
  "transaction_number": "verbatim-string",
  "items": [
    {
      "description": "verbatim-string",
      "sku": "verbatim-string",
      "quantity": "number",
      "unit": "unit-code",
      "unit_price": "number",
      "discount": "number",
      "line_total": "number",
      "tax_code": "verbatim-string"
    }
  ],
  "item_count": "integer",
  "subtotal": "number",
  "tax": "number",
  "total": "number",
  "currency": "currency",
  "payment_method": ["cash", "credit-card", "debit-card", "other"],
  "card_last_four": "verbatim-string",
  "amount_tendered": "number",
  "change_due": "number"
}""",
    "Basic bank statement": """{
  "financial_institution": "verbatim-string",
  "account_holder": "verbatim-string",
  "account_number": "verbatim-string",
  "statement_period": {
    "start_date": "date",
    "end_date": "date"
  },
  "opening_balance": "number",
  "closing_balance": "number",
  "currency": "currency",
  "transaction_count": "integer",
  "transactions": [
    {
      "date": "date",
      "description": "verbatim-string",
      "transaction_type": ["debit", "credit", "other"],
      "amount": "number",
      "balance": "number"
    }
  ]
}""",
    "Advanced bank statement": """{
  "financial_institution": {
    "name": "verbatim-string",
    "address": "verbatim-string",
    "phone": "phone-number",
    "website": "url"
  },
  "account_holder": {
    "name": "verbatim-string",
    "address": "verbatim-string"
  },
  "account": {
    "type": "verbatim-string",
    "number": "verbatim-string",
    "branch_number": "verbatim-string",
    "transit_number": "verbatim-string",
    "routing_number": "verbatim-string",
    "iban": "iban",
    "bic": "bic"
  },
  "statement_period": {
    "start_date": "date",
    "end_date": "date"
  },
  "summary": {
    "opening_balance": "number",
    "total_credits": "number",
    "total_debits": "number",
    "total_fees": "number",
    "total_interest": "number",
    "closing_balance": "number",
    "currency": "currency"
  },
  "transaction_count": "integer",
  "transactions": [
    {
      "transaction_date": "date",
      "posting_date": "date",
      "description": "verbatim-string",
      "reference_number": "verbatim-string",
      "category": [
        "deposit",
        "withdrawal",
        "transfer",
        "payment",
        "purchase",
        "fee",
        "interest",
        "refund",
        "other"
      ],
      "debit": "number",
      "credit": "number",
      "balance": "number",
      "currency": "currency"
    }
  ]
}""",
    "Invoice with line items": """{
  "document_type": "verbatim-string",
  "vendor": {
    "name": "verbatim-string",
    "address": "verbatim-string",
    "email": "email-address",
    "phone": "phone-number"
  },
  "invoice_number": "verbatim-string",
  "invoice_date": "verbatim-string",
  "line_items": [
    {
      "description": "verbatim-string",
      "category": ["product", "service", "shipping", "tax", "other"],
      "quantity": "integer",
      "unit_price": "number",
      "unit": "unit-code",
      "amount": "number"
    }
  ],
  "subtotal": "number",
  "tax": "number",
  "total": "number",
  "currency": "currency"
}""",
    "Advanced global record": """{
  "summary": "string",
  "reference_text": "verbatim-string",
  "sequence_number": "integer",
  "amount": "number",
  "effective_date": "date",
  "effective_time": "time",
  "signed_at": "date-time",
  "term": "duration",
  "active": "boolean",
  "country": "country",
  "currency": "currency",
  "language": "language",
  "language_tag": "language-tag",
  "script": "script",
  "website": "url",
  "email": "email-address",
  "phone": "phone-number",
  "iban": "iban",
  "bic": "bic",
  "measurement_unit": "unit-code",
  "regions": {
    "us": "region:US",
    "fr": "region:FR",
    "ie": "region:IE",
    "gb": "region:GB",
    "it": "region:IT",
    "es": "region:ES",
    "de": "region:DE",
    "pt": "region:PT",
    "ca": "region:CA",
    "mx": "region:MX",
    "br": "region:BR",
    "au": "region:AU",
    "jp": "region:JP",
    "kr": "region:KR"
  },
  "contacts": [
    {
      "name": "verbatim-string",
      "role": "string",
      "email": "email-address"
    }
  ],
  "status": ["draft", "active", "expired", "unknown"],
  "applicable_labels": [["legal", "financial", "technical", "personal"]],
  "keywords": ["string"]
}""",
}

default_structured_json_template = "Invoice with line items"


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):
                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.Examples(
        examples=[
            [
                "Basic receipt",
                str(
                    Path(__file__).parent
                    / "data"
                    / "samples"
                    / "receipt-ocr-original.webp"
                ),
                "",
                False,
            ],
            [
                "Invoice with line items",
                str(
                    Path(__file__).parent
                    / "data"
                    / "samples"
                    / "invoice-with-items.png"
                ),
                "",
                False,
            ],
            [
                "Advanced bank statement",
                str(
                    Path(__file__).parent
                    / "data"
                    / "samples"
                    / "BankStatementChequing.png"
                ),
                "",
                False,
            ],
            [
                "Basic bank statement",
                str(
                    Path(__file__).parent
                    / "data"
                    / "samples"
                    / "bank statement blog image.webp"
                ),
                "",
                False,
            ],
            *[
                [
                    default_structured_json_template,
                    str(path),
                    "",
                    False,
                ]
                for path in sorted(
                    path
                    for path in (Path(__file__).parent / "data" / "samples").glob(
                        "*.png"
                    )
                    if path.name
                    not in {"BankStatementChequing.png", "invoice-with-items.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",
    )

    generate_template_button.click(
        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()