import json
import re

import gradio as gr
import spaces
import torch
from PIL import Image
from transformers import AutoProcessor, Gemma3ForConditionalGeneration

# Gemma 3 1B has NO vision encoder (Google only added the SigLIP tower to the
# 4B/12B/27B sizes), so it can't read images. 4B is the smallest size that can.
MODEL_ID = "unsloth/gemma-3-4b-it-qat"

# The vision encoder always resizes to a fixed 896x896 square and emits a fixed
# 256 tokens per pass (when pan_and_scan is off), so anything above this edge
# length buys zero extra fidelity — it only costs upload time, RAM, and resize
# latency. We downscale before it ever reaches the processor.
MAX_EDGE = 1568

# Itemized receipts (groceries, long invoices) can easily need more than a
# few hundred output tokens. Truncation mid-JSON is the #1 cause of "invalid
# JSON" failures on small models, so we give this generous headroom.
MAX_NEW_TOKENS = 1536

processor = AutoProcessor.from_pretrained(MODEL_ID)
model = Gemma3ForConditionalGeneration.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
).to("cuda")


def preprocess_image(image: Image.Image) -> Image.Image:
    """Downscale + normalize before it hits the processor.

    Caps the long edge at MAX_EDGE (well above the 896px the encoder will use
    internally, so no quality is lost) and converts to RGB so odd modes
    (RGBA, CMYK, palette) don't trip up the processor.
    """
    if image.mode != "RGB":
        image = image.convert("RGB")

    long_edge = max(image.size)
    if long_edge > MAX_EDGE:
        scale = MAX_EDGE / long_edge
        new_size = (round(image.width * scale), round(image.height * scale))
        image = image.resize(new_size, Image.LANCZOS)

    return image


SCHEMA_PROMPT = """You are a document-extraction engine. You are given a photo or \
scan of an invoice or receipt. Read it carefully and return ONLY a single valid \
JSON object — no prose, no markdown code fences, nothing before or after it — \
matching exactly this schema:

{
  "document_type": "invoice" | "receipt" | "unknown",
  "vendor_name": string | null,
  "vendor_address": string | null,
  "date": string | null,
  "time": string | null,
  "currency": string | null,
  "tax_inclusive": boolean | null,
  "subtotal": number | null,
  "discount": number | null,
  "taxes": [
    {"label": string, "rate_percent": number | null, "amount": number | null}
  ],
  "tax": number | null,
  "service_charge": number | null,
  "tip": number | null,
  "total": number | null,
  "amount_paid": number | null,
  "change_due": number | null,
  "payment_method": string | null,
  "line_items": [
    {
      "description": string,
      "quantity": number | null,
      "unit": string | null,
      "unit_price": number | null,
      "discount": number | null,
      "tax_rate_percent": number | null,
      "amount": number | null
    }
  ],
  "notes": string | null
}

Rules:
- Use null for any field you cannot read. Never guess or invent values.
- Numbers must be plain numbers (no currency symbols, no thousands commas).
- "date" should be ISO 8601 (YYYY-MM-DD) if determinable; "time" as 24h HH:MM.
- "currency" should be an ISO 4217 code (e.g. "USD", "EUR") if determinable.
- If the receipt shows a SINGLE combined tax figure, put it in "tax" and leave
  "taxes" as an empty array. If it breaks tax out by type or rate (e.g. GST +
  PST, or a VAT rate table), list each one in "taxes" AND put their sum in
  "tax" for convenience.
- "tax_inclusive": true if listed prices already include tax (common outside
  the US), false if tax is added on top, null if you can't tell.
- Discounts, coupons, and refunds are negative in effect: put the discount
  amount as a positive number in the "discount" field (it will be subtracted),
  and for a line item that IS a discount/refund, use a negative "amount".
- "service_charge" is a mandatory/automatic charge (e.g. restaurant service
  fee); "tip" is a voluntary gratuity. Don't conflate them.
- If the image is not an invoice or receipt (or is unreadable), set
  "document_type" to "unknown", leave other fields null, "line_items" as an
  empty array, and briefly say why in "notes".
- If you notice the totals don't add up, don't silently correct them — record
  what the document actually states and mention the discrepancy in "notes".
- Output valid JSON and nothing else.
"""


def isolate_json(text: str) -> str:
    """Best-effort isolation of the JSON object from the model's raw output."""
    text = text.strip()
    if text.startswith("```"):
        text = text.strip("`")
        if text.lower().startswith("json"):
            text = text[4:]
        text = text.strip()

    start = text.find("{")
    end = text.rfind("}")
    if start != -1 and end != -1 and end > start:
        text = text[start : end + 1]
    return text


def parse_json_loose(text: str):
    """Try strict parsing, then one repair pass for the most common small-model
    glitch (a trailing comma before a closing brace/bracket)."""
    try:
        return json.loads(text), None
    except json.JSONDecodeError as e:
        repaired = re.sub(r",\s*([}\]])", r"\1", text)
        try:
            return json.loads(repaired), None
        except json.JSONDecodeError:
            return None, str(e)


def validate_totals(data: dict) -> dict:
    """Lightweight arithmetic cross-check. Never raises — this is a hint for
    downstream review, not a hard guarantee, since it can only catch
    internal inconsistency, not a value that was misread but still adds up."""
    try:
        line_items = data.get("line_items") or []
        item_amounts = [
            i["amount"] for i in line_items if isinstance(i.get("amount"), (int, float))
        ]
        computed_subtotal = round(sum(item_amounts), 2) if item_amounts else None

        taxes = data.get("taxes") or []
        tax_amounts = [
            t["amount"] for t in taxes if isinstance(t.get("amount"), (int, float))
        ]
        computed_tax = round(sum(tax_amounts), 2) if tax_amounts else data.get("tax")

        subtotal = data.get("subtotal") if data.get("subtotal") is not None else computed_subtotal
        discount = data.get("discount")
        service_charge = data.get("service_charge")
        tip = data.get("tip")

        parts = [subtotal, computed_tax, service_charge, tip]
        numeric_parts = [p for p in parts if isinstance(p, (int, float))]
        computed_total = None
        if numeric_parts:
            computed_total = sum(numeric_parts)
            if isinstance(discount, (int, float)):
                computed_total -= discount
            computed_total = round(computed_total, 2)

        stated_total = data.get("total")
        matches = None
        if isinstance(stated_total, (int, float)) and computed_total is not None:
            matches = abs(computed_total - stated_total) <= 0.02

        return {
            "computed_subtotal_from_line_items": computed_subtotal,
            "computed_total": computed_total,
            "stated_total": stated_total,
            "matches_stated_total": matches,
        }
    except Exception:
        return {"error": "validation check failed"}


@spaces.GPU(duration=90)
def extract(image: Image.Image, high_accuracy: bool = False):
    if image is None:
        return json.dumps({"error": "no image provided"}, indent=2)

    image = preprocess_image(image)

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": SCHEMA_PROMPT},
            ],
        }
    ]

    # pan_and_scan crops the image into extra tiles so long/narrow receipts
    # (thermal-paper strips) don't get squashed by the 896x896 square resize.
    # It costs ~256 more tokens per extra tile, so it's opt-in, not default.
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
        add_generation_prompt=True,
        do_pan_and_scan=high_accuracy,
    ).to(model.device, dtype=torch.bfloat16)

    input_len = inputs["input_ids"].shape[-1]

    with torch.inference_mode():
        generated = model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
        )

    output_ids = generated[0][input_len:]
    likely_truncated = output_ids.shape[-1] >= MAX_NEW_TOKENS

    text = processor.decode(output_ids, skip_special_tokens=True)
    cleaned = isolate_json(text)

    parsed, err = parse_json_loose(cleaned)

    if parsed is None:
        return json.dumps(
            {
                "error": "model did not return valid JSON",
                "likely_truncated": likely_truncated,
                "parse_error": err,
                "raw_output": text.strip(),
            },
            indent=2,
        )

    parsed["validation"] = validate_totals(parsed)
    if likely_truncated:
        parsed["validation"]["warning"] = (
            "Output may have been truncated before completion — treat with caution, "
            "especially for line_items near the end of the list."
        )

    return json.dumps(parsed, indent=2)


demo = gr.Interface(
    fn=extract,
    inputs=[
        gr.Image(type="pil", label="Invoice / Receipt image"),
        gr.Checkbox(
            label="High accuracy (long/narrow receipts)",
            value=False,
            info="Costs more tokens/time — only needed for thermal-paper strips or tiny print.",
        ),
    ],
    outputs=gr.Textbox(label="Structured JSON", lines=32),
    title="Invoice & Receipt Extractor (Gemma 3 4B)",
    description=(
        "Upload a photo or scan of an invoice or receipt. Runs on ZeroGPU. "
        "Callable via the Gradio API (see 'Use via API' link below, or gradio_client)."
    ),
    api_name="extract",
)

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