#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""提炼失败条目补跑：缩短正文 + 提高额度，合并回 enriched 文件。

用法：intel_enrich_retry.py [first|daily]
"""
import json
import os
import re
import sys
import time
from concurrent.futures import ThreadPoolExecutor

import requests

W = '/Users/liuxinyuan/Desktop/Hermes输出-工作类'
RUNS = os.path.join(W, '数据/market_intel/runs')
MODE = sys.argv[1] if len(sys.argv) > 1 else 'daily'
F = os.path.join(RUNS, '20261013_' + ('enriched_first.jsonl' if MODE == 'first' else 'enriched.jsonl'))
CFG = os.path.expanduser('~/.hermes/profiles/zone-b/config.yaml')
KEY = re.search(r'base_url:\s*https://api\.deepseek\.com[\s\S]{0,200}?api_key:\s*(\S+)',
                open(CFG, encoding='utf-8').read()).group(1)
UA = ('Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 '
      '(KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36')
S = requests.Session()
S.headers.update({'User-Agent': UA})

PROMPT = """你是产业情报分析员。只依据正文提炼要点，输出 JSON（无多余文字）：
{"时间":"","地点":"","人物":"","主体":"","事件":"","关键事项":"","与我方关联":""}
规则：正文没有的写「未提及」；「事件」≤40字；「关键事项」用「；」分隔 2-4 条（每条≤25字）；
「与我方关联」站在分子诊断原料供应商（酶/预混料/冻干体系/引物探针）角度写一句判断，无则「无直接关联」。

标题：{t}
正文：{b}
"""


def fetch(u):
    try:
        if u.lower().endswith('.pdf'):
            r = S.get(u, timeout=40)
            p = os.path.join(os.environ.get('TMPDIR', '/tmp'),
                             '_retry_%08d.pdf' % (abs(hash(u)) % 10 ** 8))
            open(p, 'wb').write(r.content)
            import fitz
            d = fitz.open(p)
            return re.sub(r'\s+', ' ', d[0].get_text())[:1400]
        r = S.get(u, timeout=25)
        r.encoding = r.apparent_encoding or r.encoding
        t = re.sub(r'(?s)<[^>]+>', ' ', re.sub(r'(?is)<(script|style|nav|footer|head)[^>]*>.*?</\1>', ' ', r.text))
        return re.sub(r'\s+', ' ', t)[:1400]
    except Exception:
        return ''


def run(i):
    b = fetch(i['链接'])
    prompt = PROMPT.replace('{t}', i['标题']).replace('{b}', b or (i.get('摘要') or '')[:600])
    for mt in (6000, 8000):
        try:
            r = requests.post('https://api.deepseek.com/chat/completions', timeout=150,
                              headers={'Authorization': 'Bearer ' + KEY, 'Content-Type': 'application/json'},
                              json={'model': 'deepseek-flash', 'temperature': 0, 'max_tokens': mt,
                                    'messages': [{'role': 'user', 'content': prompt}]})
            ch = (r.json().get('choices') or [{}])[0]
            c = (ch.get('message', {}).get('content') or '').strip()
            if not c:
                continue
            m = re.search(r'\{.*\}', re.sub(r'^```(?:json)?|```$', '', c, flags=re.M), re.S)
            if m:
                return json.loads(m.group(0))
        except Exception:
            time.sleep(2)
    return None


data = json.load(open(F, encoding='utf-8'))
ONLY = [a for a in sys.argv[2:] if a]
if ONLY:
    bad = [i for i in data if i['来源'] in ONLY]
    print('定向重跑来源:', ONLY)
else:
    bad = [i for i in data if not (isinstance(i.get('要点'), dict) and '错误' not in i['要点'])]
print('失败条目: %d 条 → 补跑中' % len(bad))
with ThreadPoolExecutor(max_workers=4) as ex:
    outs = list(ex.map(run, bad))
n = 0
for it, o in zip(bad, outs):
    if o:
        it['要点'] = o
        n += 1
json.dump(data, open(F, 'w', encoding='utf-8'), ensure_ascii=False, indent=1)
still = sum(1 for i in data if not (isinstance(i.get('要点'), dict) and '错误' not in i['要点']))
print('✅ 补回 %d 条 ｜ 仍失败 %d 条 ｜ 文件 %s' % (n, still, os.path.basename(F)))
