#!/usr/bin/env python3
"""Invoice field extraction with a vision LLM over an OpenAI-compatible API,
schema-enforced with response_format json_schema, scored against ground truth."""
import json, os, sys, time, base64, re, threading, queue, requests  # pip install requests
API="https://paraloncloud.com/v1/chat/completions"; MODEL="qwen3.8-27b"; KEY=os.environ["PRL_KEY"]
GT=json.load(open(os.environ.get("GT","img/ground_truth.json"))); LIMIT=int(os.environ.get("LIMIT","0")) or len(GT)
IN_FLIGHT=int(os.environ.get("IN_FLIGHT","4")); PRICE_IN,PRICE_OUT=0.12,1.70
SCHEMA={"type":"object","properties":{
 "vendor_name":{"type":"string"},"customer_name":{"type":"string"},"invoice_number":{"type":"string"},
 "invoice_date_as_printed":{"type":"string","description":"exactly as printed"},"due_date_as_printed":{"type":"string","description":"exactly as printed"},"invoice_date":{"type":"string","description":"YYYY-MM-DD"},"due_date":{"type":"string","description":"YYYY-MM-DD"},
 "currency":{"type":"string","enum":["EUR","USD","GBP","RON"]},
 "subtotal":{"type":"number"},"tax_rate_percent":{"type":"number"},"tax_amount":{"type":"number"},"total":{"type":"number"},
 "line_items":{"type":"array","items":{"type":"object","properties":{"description":{"type":"string"},"quantity":{"type":"number"},"unit_price":{"type":"number"},"amount":{"type":"number"}},"required":["description","quantity","unit_price","amount"]}}},
 "required":["vendor_name","customer_name","invoice_number","invoice_date_as_printed","due_date_as_printed","invoice_date","due_date","currency","subtotal","tax_rate_percent","tax_amount","total","line_items"]}
PROMPT=("Extract the fields of this invoice. Copy each date exactly as printed into the *_as_printed field, and also give it as YYYY-MM-DD. Numbers without thousands separators. "
        "Copy names and the invoice number exactly as printed. List every line item.")
def call(gt):
    path=gt["_file"]; mime="image/png" if path.endswith(".png") else "image/jpeg"
    b64=base64.b64encode(open(path,"rb").read()).decode()
    body={"model":MODEL,"messages":[{"role":"user","content":[{"type":"text","text":PROMPT},{"type":"image_url","image_url":{"url":f"data:{mime};base64,{b64}"}}]}],
          "response_format":{"type":"json_schema","json_schema":{"name":"invoice","schema":SCHEMA}},
          "chat_template_kwargs":{"enable_thinking":False},"temperature":0,"max_tokens":900}
    t0=time.time(); r=requests.post(API,headers={"Authorization":f"Bearer {KEY}"},json=body,timeout=240); lat=time.time()-t0
    if r.status_code!=200: return {"error":f"HTTP {r.status_code}: {r.text[:160]}","lat":lat}
    d=r.json(); u=d.get("usage",{})
    try: out=json.loads(d["choices"][0]["message"]["content"])
    except Exception as e: return {"error":f"parse: {e}","lat":lat,"usage":u}
    return {"out":out,"lat":lat,"usage":u}
norm=lambda s: re.sub(r"[^a-z0-9]","",str(s).lower())
import datetime
def parse_printed(txt, locale):
    txt=txt.strip()
    for fmt in (["%m/%d/%Y","%m/%d/%y"] if locale=="us" else ["%d.%m.%Y","%d/%m/%Y","%Y-%m-%d"]):
        try: return datetime.datetime.strptime(txt,fmt).strftime("%Y-%m-%d")
        except: pass
    return None
def score(gt,out):
    for k in ("invoice_date","due_date"):
        p=parse_printed(out.get(k+"_as_printed",""), gt["_datefmt"])
        if p: out[k]=p
    s={}
    for k in ("vendor_name","customer_name","invoice_number","invoice_date","due_date","currency"): s[k]=norm(out.get(k,""))==norm(gt[k])
    for k in ("subtotal","tax_rate_percent","tax_amount","total"):
        try: s[k]=abs(float(out.get(k,-1))-float(gt[k]))<0.011
        except: s[k]=False
    li=out.get("line_items") or []; s["line_item_count"]=len(li)==len(gt["line_items"])
    ok=0
    for a,b in zip(li,gt["line_items"]):
        try: ok+= (norm(a.get("description",""))==norm(b["description"]) and abs(float(a.get("quantity",-1))-b["quantity"])<0.01 and abs(float(a.get("amount",-1))-b["amount"])<0.011)
        except: pass
    s["line_items_exact"]=ok==len(gt["line_items"]) and len(li)==len(gt["line_items"])
    s["_li_ok"]=ok; s["_li_n"]=len(gt["line_items"])
    return s
lock=threading.Lock(); results=[]; q=queue.Queue()
for g in GT[:LIMIT]: q.put(g)
def worker():
    while True:
        try: g=q.get_nowait()
        except queue.Empty: return
        r=call(g); 
        with lock: results.append((g,r))
th=[threading.Thread(target=worker) for _ in range(IN_FLIGHT)]; t0=time.time(); [t.start() for t in th]; [t.join() for t in th]; wall=time.time()-t0
fields=["vendor_name","customer_name","invoice_number","invoice_date","due_date","currency","subtotal","tax_rate_percent","tax_amount","total","line_item_count","line_items_exact"]
agg={"docs":len(results),"errors":0,"tokens_in":0,"tokens_out":0,"lat":[],"by_field":{f:0 for f in fields},"all_fields_correct":0,"clean":{"n":0,"all":0},"degraded":{"n":0,"all":0},"li_ok":0,"li_n":0,"rows":[]}
for g,r in results:
    if "error" in r: agg["errors"]+=1; agg["rows"].append({"file":g["_file"],"error":r["error"]}); continue
    u=r["usage"]; agg["tokens_in"]+=u.get("prompt_tokens",0); agg["tokens_out"]+=u.get("completion_tokens",0); agg["lat"].append(r["lat"])
    s=score(g,r["out"]); allok=all(s[f] for f in fields)
    for f in fields: agg["by_field"][f]+=int(s[f])
    agg["all_fields_correct"]+=int(allok); agg["li_ok"]+=s["_li_ok"]; agg["li_n"]+=s["_li_n"]
    grp="degraded" if g["_degraded"] else "clean"; agg[grp]["n"]+=1; agg[grp]["all"]+=int(allok)
    agg["rows"].append({"file":g["_file"],"degraded":g["_degraded"],"layout":g["_layout"],"font":g["_font"],"datefmt":g["_datefmt"],"all_ok":allok,"wrong":[f for f in fields if not s[f]],"out":r["out"],"lat":round(r["lat"],1)})
n=max(1,agg["docs"]-agg["errors"]); lat=sorted(agg["lat"])
summary={"docs":agg["docs"],"errors":agg["errors"],"all_fields_correct":agg["all_fields_correct"],"all_fields_correct_pct":round(100*agg["all_fields_correct"]/n,1),
 "by_field_pct":{f:round(100*v/n,1) for f,v in agg["by_field"].items()},"line_items_pct":round(100*agg["li_ok"]/max(1,agg["li_n"]),1),
 "clean_all_ok":f'{agg["clean"]["all"]}/{agg["clean"]["n"]}',"degraded_all_ok":f'{agg["degraded"]["all"]}/{agg["degraded"]["n"]}',
 "tokens_in_per_doc":round(agg["tokens_in"]/n),"tokens_out_per_doc":round(agg["tokens_out"]/n),
 "latency_p50_s":round(lat[len(lat)//2],1) if lat else None,"latency_p90_s":round(lat[int(len(lat)*0.9)],1) if lat else None,"wall_s":round(wall,1),
 "docs_per_hour":round(n/wall*3600),"cost_per_doc_usd":round((agg["tokens_in"]*PRICE_IN+agg["tokens_out"]*PRICE_OUT)/1e6/n,5),
 "cost_per_1k_docs_usd":round((agg["tokens_in"]*PRICE_IN+agg["tokens_out"]*PRICE_OUT)/1e6/n*1000,3)}
print(json.dumps(summary,indent=1)); json.dump({"summary":summary,"rows":agg["rows"]},open(os.environ.get("OUT","results.json"),"w"),indent=1)