# -*- coding: utf-8 -*-
"""对识别为空/失败的页重跑（空内容也重试）"""
import base64, json, os, glob, time
import urllib.request
from concurrent.futures import ThreadPoolExecutor

API="https://api.deepseek.com/v1/chat/completions"; MODEL="deepseek-v4-flash-vision-exp"
KEY="sk-2af0093dfbac45359276d8aa85d83618"
SRC="/home/zyw/Downloads/dl-hub/26-合同审查/originals"
OUT="/home/zyw/Downloads/dl-hub/26-合同审查/pages"
PROMPT=open("/home/zyw/Downloads/dl-hub/26-合同审查/prompt.txt",encoding="utf-8").read()

def ocr(path):
    b64=base64.b64encode(open(path,"rb").read()).decode()
    payload={"model":MODEL,"temperature":0,"max_tokens":8000,"messages":[{"role":"user","content":[
      {"type":"image_url","image_url":{"url":"data:image/jpeg;base64,"+b64}},
      {"type":"text","text":PROMPT}]}]}
    req=urllib.request.Request(API,data=json.dumps(payload).encode(),
      headers={"Authorization":"Bearer "+KEY,"Content-Type":"application/json"})
    with urllib.request.urlopen(req,timeout=600) as r:
        d=json.loads(r.read())
    return d["choices"][0]["message"].get("content") or ""

def work(f):
    name=os.path.splitext(os.path.basename(f))[0]
    dst=os.path.join(OUT,name+".md")
    if os.path.exists(dst) and os.path.getsize(dst)>50: return "skip "+name
    for a in range(1,6):
        try:
            t=ocr(f)
            if len(t.strip())>50:
                open(dst,"w",encoding="utf-8").write(t)
                return f"OK {name} ({len(t)}字) 第{a}次"
        except Exception as e:
            t=""
        time.sleep(4*a)
    return "FAIL "+name

files=sorted(glob.glob(SRC+"/*.jpg"))
with ThreadPoolExecutor(max_workers=3) as ex:
    for r in ex.map(work,files): print(r,flush=True)
print("RETRY DONE")