#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""本地视觉模型 qwen2.5vl:7b 票据识别 + 与台账对照

用法：
  python3 脚本/local_ocr_invoices.py --limit 20
输出：
  数据/local_ocr_results.json        逐张识别结果（可重跑续跑）
  销售分析/本地OCR票据识别_对照_<日期>.xlsx

说明：只做识别与对照，不写台账；金额以台账/票面人工核对为准。
"""
import argparse
import base64
import io
import json
import os
import re
import time

import openpyxl
import urllib.request
from openpyxl.styles import Alignment, Font, PatternFill
from openpyxl.utils import get_column_letter

ROOT = '/Users/liuxinyuan/Desktop/Hermes输出-工作类'
EXP = os.path.join(ROOT, '报销助手')
INBOX = os.path.join(EXP, '数据', '收件箱')
LEDGER = os.path.join(EXP, '数据', '报销台账.json')
RESULT = os.path.join(ROOT, '数据', 'local_ocr_results.json')
OUT = os.path.join(ROOT, '销售分析', '本地OCR票据识别_对照_2026-10-26.xlsx')
OLLAMA = 'http://127.0.0.1:11434/api/generate'
PROMPT = ('这是发票图片。请提取以下字段并用 JSON 输出，键名固定：'
          '发票号码、开票日期、销售方名称、价税合计、项目名称。'
          '日期格式 YYYY-MM-DD；金额只填数字；看不清就填 null，不要猜。')


def pdf_to_png(path, dpi=200):
    import fitz
    doc = fitz.open(path)
    pix = doc[0].get_pixmap(dpi=dpi)
    return pix.tobytes('png')


def load_image(path):
    ext = os.path.splitext(path)[1].lower()
    if ext == '.pdf':
        return pdf_to_png(path)
    return open(path, 'rb').read()


def parse_json(s):
    s = re.sub(r'^```(json)?|```$', '', (s or '').strip(), flags=re.M).strip()
    m = re.search(r'\{.*\}', s, re.S)
    if m:
        s = m.group(0)
    try:
        return json.loads(s)
    except Exception:
        return {}


def ocr(path):
    b64 = base64.b64encode(load_image(path)).decode()
    req = {'model': 'qwen2.5vl:7b', 'prompt': PROMPT, 'images': [b64], 'stream': False,
           'options': {'temperature': 0, 'num_predict': 260}}
    t = time.time()
    r = json.load(urllib.request.urlopen(urllib.request.Request(
        OLLAMA, data=json.dumps(req, ensure_ascii=False).encode(),
        headers={'Content-Type': 'application/json'}), timeout=600))
    return parse_json(r.get('response', '')), time.time() - t


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--limit', type=int, default=20)
    a = ap.parse_args()

    files = sorted(f for f in os.listdir(INBOX)
                   if os.path.splitext(f)[1].lower() in ('.pdf', '.jpg', '.jpeg', '.png'))
    done = json.load(open(RESULT, encoding='utf-8')) if os.path.exists(RESULT) else {}
    todo = [f for f in files if f not in done][:a.limit]
    print('待识别 %d / 共 %d（已完成 %d）' % (len(todo), len(files), len(done)))

    for i, f in enumerate(todo, 1):
        try:
            data, el = ocr(os.path.join(INBOX, f))
            done[f] = {'ocr': data, 'elapsed': round(el, 1)}
            print('  %-40s %.0fs %s' % (f[:40], el, json.dumps(data, ensure_ascii=False)[:110]), flush=True)
        except Exception as e:
            done[f] = {'ocr': {}, 'error': str(e)}
            print('  %-40s ERROR %s' % (f[:40], e), flush=True)
        json.dump(done, open(RESULT, 'w', encoding='utf-8'), ensure_ascii=False, indent=0)

    # 与台账对照
    led = json.load(open(LEDGER, encoding='utf-8'))['发票']
    by_no = {str(x.get('发票号码')): x for x in led if x.get('发票号码')}

    def num(v):
        try:
            return float(re.sub(r'[^\d.\-]', '', str(v)))
        except Exception:
            return None

    wb = openpyxl.Workbook()
    ws = wb.active
    ws.title = '识别对照'
    ws.append(['文件', 'OCR 发票号', 'OCR 日期', 'OCR 销售方', 'OCR 金额', 'OCR 项目',
               '台账 发票号', '台账 金额', '号码匹配', '金额差', '耗时(s)'])
    stat = {'n': 0, 'hit': 0, 'amt_ok': 0}
    for f, v in done.items():
        o = v.get('ocr') or {}
        no = str(o.get('发票号码') or '').strip()
        amt = num(o.get('价税合计'))
        inv = by_no.get(no)
        l_amt = num(inv.get('价税合计')) if inv else None
        diff = None if (amt is None or l_amt is None) else round(amt - l_amt, 2)
        stat['n'] += 1
        if inv:
            stat['hit'] += 1
            if diff is not None and abs(diff) <= 0.01:
                stat['amt_ok'] += 1
        ws.append([f, no, o.get('开票日期'), o.get('销售方名称'), amt, o.get('项目名称'),
                   (inv or {}).get('发票号码'), l_amt, '是' if inv else '否', diff, v.get('elapsed')])
    for c in ws[1]:
        c.font = Font(color='FFFFFF', bold=True)
        c.fill = PatternFill('solid', fgColor='1F4E79')
    for i, w in enumerate([34, 22, 12, 30, 10, 24, 20, 10, 9, 9, 9], 1):
        ws.column_dimensions[get_column_letter(i)].width = w
    ws.freeze_panes = 'A2'
    for row in ws.iter_rows(min_row=2):
        for c in row:
            c.alignment = Alignment(vertical='top', wrap_text=True)

    wb.save(OUT)
    print('✅ %s ｜ 已识别 %d ｜ 号码命中台账 %d ｜ 金额一致 %d' % (OUT, stat['n'], stat['hit'], stat['amt_ok']))


if __name__ == '__main__':
    main()