#!/usr/bin/env python3
"""Local receipt OCR: given a receipt image, print an ordered list of rows of text.

Uses a local Qwen2.5-VL model served by ollama. No cloud calls, no API key.
Zero third-party dependencies — talks to the ollama HTTP API with the stdlib.

Usage:
    python3 receipt_ocr.py [image.jpg]            # pretty, numbered rows
    python3 receipt_ocr.py image.jpg --json       # raw JSON {"rows": [...]}

Env:
    OLLAMA_HOST   default http://127.0.0.1:11434
    OLLAMA_MODEL  default qwen2.5vl:7b
"""

import base64
import json
import os
import sys
import urllib.request

DEFAULT_IMAGE = os.path.join(os.path.dirname(__file__), "example", "data", "aldi.jpg")
HOST = os.environ.get("OLLAMA_HOST", "http://127.0.0.1:11434").rstrip("/")
MODEL = os.environ.get("OLLAMA_MODEL", "qwen2.5vl:7b")

PROMPT = (
    "You are an OCR engine for retail receipts. Transcribe EVERY line of text on "
    "this receipt, from top to bottom, exactly as printed. Each physical row of the "
    "receipt becomes one string in the output, with its columns joined left-to-right "
    "and separated by single spaces (e.g. an item code, its name, and its price all "
    "go in the same row). Do not merge or split rows, do not reorder, do not add "
    "commentary, do not correct spelling, do not invent text. Preserve numbers and "
    "punctuation verbatim. Return strict JSON only."
)

# Structured-output schema: forces the model to return {"rows": ["...", ...]}.
SCHEMA = {
    "type": "object",
    "properties": {"rows": {"type": "array", "items": {"type": "string"}}},
    "required": ["rows"],
}


def ocr(image_path):
    with open(image_path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode("ascii")

    payload = {
        "model": MODEL,
        "messages": [{"role": "user", "content": PROMPT, "images": [b64]}],
        "stream": False,
        "format": SCHEMA,
        "options": {"temperature": 0, "num_ctx": 8192},
    }

    req = urllib.request.Request(
        f"{HOST}/api/chat",
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=600) as resp:
        body = json.load(resp)

    content = body["message"]["content"]
    return json.loads(content)["rows"]


def main():
    args = [a for a in sys.argv[1:] if not a.startswith("--")]
    as_json = "--json" in sys.argv
    image_path = args[0] if args else DEFAULT_IMAGE

    if not os.path.exists(image_path):
        sys.exit(f"image not found: {image_path}")

    rows = ocr(image_path)

    if as_json:
        print(json.dumps({"rows": rows}, ensure_ascii=False, indent=2))
    else:
        width = len(str(len(rows)))
        for i, row in enumerate(rows, 1):
            print(f"{str(i).rjust(width)}  {row}")


if __name__ == "__main__":
    main()