"""YouRunAI: local image-to-JSON receipt reader. No hosted inference."""
import argparse
import json
from datetime import datetime
from pathlib import Path
import gradio as gr
from ollama import Client
from PIL import Image, ImageDraw, ImageFont

parser = argparse.ArgumentParser()
parser.add_argument("--model", default="qwen3.5:4b")
args = parser.parse_args()
out = Path("outputs")
out.mkdir(exist_ok=True)
sample = Path("sample-receipt.png")
if not sample.exists():
    picture = Image.new("RGB", (900, 600), "white")
    pen = ImageDraw.Draw(picture)
    font = ImageFont.load_default(size=34)
    pen.multiline_text((50, 45), "CEDAR CAFE\nDate: 2026-09-28\nCoffee   4.00 USD\nSandwich 8.00 USD\nTOTAL   12.00 USD", fill="black", font=font, spacing=20)
    picture.save(sample)

schema = {"type": "object", "properties": {
    "merchant": {"type": ["string", "null"]},
    "date": {"type": ["string", "null"]},
    "total": {"type": ["number", "null"]},
    "currency": {"type": ["string", "null"]}},
    "required": ["merchant", "date", "total", "currency"], "additionalProperties": False}
client = Client(host="http://127.0.0.1:11434", timeout=600)

def extract(image):
    if image is None:
        raise gr.Error("Choose a receipt image first.")
    try:
        result = client.chat(model=args.model, messages=[{
            "role": "user", "content": "Read this receipt. Extract merchant, date, total and currency. Use null for anything unreadable; never invent values.",
            "images": [image]}], format=schema, think=False,
            options={"num_ctx": 4096, "num_predict": 1024, "temperature": 0}, keep_alive="10m")
        data = json.loads(result.message.content)
        path = out / (datetime.now().strftime("receipt-%Y%m%d-%H%M%S-%f") + ".json")
        path.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
        return data, str(path)
    except Exception as error:
        raise gr.Error(f"Extraction failed: {error}. Check that Ollama is running and the selected model was pulled.") from error

with gr.Blocks(title="Private receipt reader") as app:
    gr.Markdown(f"# Private receipt reader\nModel: **{args.model}** · 4K context. Compare every result with the image before using it.")
    image = gr.Image(value=str(sample), type="filepath", label="Receipt image (PNG or JPEG)")
    button = gr.Button("Extract receipt", variant="primary")
    result = gr.JSON(label="Extracted fields")
    download = gr.File(label="Download JSON")
    button.click(extract, inputs=image, outputs=[result, download])
app.launch(server_name="127.0.0.1", server_port=7860, share=False)