"""OCR a scanned invoice, print its rows as CSV, and flag bad math.

Usage: python invoice_ocr.py invoice.jpg
"""

import csv
import re
import sys

import pytesseract
from PIL import Image

MONEY = r"\$?(\d[\d,]*\.\d\d)"  # 1234.56, 1,234.56 or $1,234.56
# One item row: SKU, description, qty, unit price, amount
ROW = re.compile(rf"^\W*(\w+)\W+(.+?)\W+(\d+)\W+{MONEY}\W+{MONEY}\W*$")
FOOTER = re.compile(rf"^\W*(Subtotal|Tax|Total)\b.*?{MONEY}\W*$", re.IGNORECASE)
MARK = {True: "ok", False: "CHECK"}


def ocr(path, psm=4, cut=160):
    image = Image.open(path)
    dpi = round(image.info.get("dpi", (300, 300))[0])  # without it, Tesseract guesses
    # Pure black and white drops the scan noise
    ink = image.convert("L").point(lambda v: 255 if v > cut else 0)
    return pytesseract.image_to_string(ink, config=f"--psm {psm} --dpi {dpi}")


def cents(money):  # whole cents, so the math has no float rounding
    return int(money.replace(",", "").replace(".", "") or 0)


def parse(text):
    number = re.search(r"Invoice\s*#?\s*(\w+)", text)
    date = re.search(r"Date:?\s*([A-Z][a-z]+ \d{1,2}, \d{4})", text)
    rows, foot, checks = [], {}, []
    for line in text.splitlines():
        if m := ROW.match(line):
            sku, desc, qty, unit, amount = m.groups()
            ok = int(qty) * cents(unit) == cents(amount)
            ok = ok and not re.search(rf"-{MONEY}", line)  # a credit loses its minus
            unit, amount = unit.replace(",", ""), amount.replace(",", "")
            rows.append([sku, desc, qty, unit, amount, MARK[ok]])
        elif m := FOOTER.match(line):
            foot[m[1].lower()] = m[2].replace(",", "")
        elif re.search(r"\d[.,]\d\d\b", line):  # a price, but not a row
            first = re.findall(r"\w+", line)[0]
            rows.append([first, line.strip(), "", "", "", "CHECK"])
    if "subtotal" in foot:  # the amounts must add up to the subtotal
        ok = sum(cents(r[4]) for r in rows) == cents(foot["subtotal"])
        checks.append(["subtotal", foot["subtotal"], MARK[ok]])
    if {"subtotal", "tax", "total"} <= foot.keys():  # subtotal + tax = total
        ok = cents(foot["subtotal"]) + cents(foot["tax"]) == cents(foot["total"])
        checks += [["tax", foot["tax"]], ["total", foot["total"], MARK[ok]]]
    return number and number[1], date and date[1], rows, checks


def report(text, out=sys.stdout):
    number, date, rows, checks = parse(text)
    head = [["invoice", number], ["date", date]]
    head.append(["sku", "description", "qty", "unit_price", "amount", "check"])
    csv.writer(out, lineterminator="\n").writerows(head + rows + checks)


if __name__ == "__main__":
    report(ocr(sys.argv[1]))