#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""DeepSeek Search / Fetch 封装（自研搜索引擎）
从 Hermes 配置读取 key，不落盘、不打印。
"""
import json
import os
import re
import time

import requests

BASE = 'https://api.deepseek.com'
_CFG = [os.path.expanduser('~/.hermes/config.yaml'),
        os.path.expanduser('~/.hermes/profiles/zone-b/config.yaml')]


def _key():
    env = os.environ.get('DEEPSEEK_API_KEY')
    if env:
        return env.strip()
    for c in _CFG:
        if not os.path.exists(c):
            continue
        s = open(c, encoding='utf-8').read()
        m = re.search(r'deepseek:\s*\n(?:.*\n){0,12}?\s*api_key:\s*["\']?([A-Za-z0-9\-_]{16,})', s)
        if m:
            return m.group(1)
    raise RuntimeError('未找到 DeepSeek API Key')


def _headers():
    return {'Authorization': 'Bearer %s' % _key(), 'Content-Type': 'application/json'}


def search(query, count=10, need_content=False, ori_question=None, localization=None,
           retries=3):
    body = {'query': query, 'count': max(1, min(50, count)), 'need_content': need_content}
    if ori_question:
        body['ori_question'] = ori_question[:128]
    if localization:
        body['localization'] = localization
    last = None
    for k in range(retries):
        r = requests.post(BASE + '/v1/tools/search', headers=_headers(), json=body, timeout=60)
        if r.status_code == 200:
            return r.json().get('results', [])
        last = r
        if r.status_code in (408, 429) or r.status_code >= 500:
            time.sleep(2 * (k + 1))
            continue
        break
    raise RuntimeError('search %s: %s' % (last.status_code, last.text[:200]))


def fetch(url, retries=3):
    last = None
    for k in range(retries):
        r = requests.post(BASE + '/v1/tools/fetch', headers=_headers(),
                          json={'url': url}, timeout=90)
        if r.status_code == 200:
            return r.json()
        last = r
        if r.status_code in (408, 429) or r.status_code >= 500:
            time.sleep(2 * (k + 1))
            continue
        break
    raise RuntimeError('fetch %s: %s' % (last.status_code, last.text[:200]))


if __name__ == '__main__':
    import sys
    q = sys.argv[1] if len(sys.argv) > 1 else 'PCR 原料 酶 市场 2026'
    n = int(sys.argv[2]) if len(sys.argv) > 2 else 5
    t0 = time.time()
    rs = search(q, count=n, need_content=False)
    print('query: %s ｜%d 条 ｜%.1fs' % (q, len(rs), time.time() - t0))
    for x in rs:
        print('  [%.1f|auth %s] %s' % (x.get('rerank_score', 0), x.get('authority'),
                                       x.get('title', '')[:60]))
        print('     %s ｜%s' % (x.get('url', '')[:90], (x.get('date') or '')[:10]))
    if rs:
        t1 = time.time()
        f = fetch(rs[0]['url'])
        print('fetch: %d 字 ｜%.1fs' % (len(f.get('content', '')), time.time() - t1))
        print('  开头:', f.get('content', '')[:160].replace('\n', ' '))