#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
lit_search_free.py — 免费文献"发现"通道（宝锐文献基础设施 · 杠杆1）

解决的是**只知题目/关键词 → 发现文献**（此前只会"已知 DOI 找全文"）。
三个免费源，全部无需 API key：
  pubmed    PubMed E-utilities（esearch + efetch）——生物医学题录、PMID/PMCID/DOI
  crossref  Crossref REST —— 全学科题录、刊名、被引、许可
  openalex  OpenAlex —— 题录 + 是否开放获取 + 引用数（补召回最好用）
  all       三源合并去重，标注 OA 状态与"是否已在本地区草稿清单里"

用法：
  /usr/bin/python3 脚本/lit_search_free.py all "LAMP lyophilized trehalose limit of detection" --n 15
  /usr/bin/python3 脚本/lit_search_free.py pubmed "recombinase polymerase amplification freeze-dried" --n 10
  /usr/bin/python3 脚本/lit_search_free.py all "duplex sequencing error rate" --n 20 --json
"""
import argparse
import html as _html
import json
import os
import re
import sys
import time
import urllib.parse
import urllib.request

HOME = os.path.expanduser('~')
BASE = os.path.join(HOME, 'Desktop/Hermes输出-工作类')
MAIL = 'research@biorisales.site'
UA = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124 Safari/537.36'
KDIR = os.path.join(BASE, '知识沉淀')
OA_DIR = os.path.join(KDIR, '_全文', 'OA')


def http(url, timeout=30):
    req = urllib.request.Request(url, headers={'User-Agent': UA, 'Accept': 'application/json'})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return r.read().decode('utf-8', 'ignore')


def jget(url, timeout=30):
    try:
        return json.loads(http(url, timeout=timeout))
    except Exception as e:
        return {'_error': str(e)}


def norm_id(s):
    return re.sub(r'[^a-z0-9]', '', str(s).lower())


def known_dois():
    """本地已有的 DOI（全文 + 四份草稿引用），用于标注"是否已覆盖"。"""
    s = set()
    for fn in (os.listdir(OA_DIR) if os.path.isdir(OA_DIR) else []):
        if fn.endswith('.txt'):
            s.add(norm_id(fn[:-4].replace('_', '/')))
    for fn in os.listdir(KDIR):
        if fn.startswith('_draft_') and fn.endswith('.md'):
            for d in re.findall(r'10\.\d{4,9}/[A-Za-z0-9._()/:;<>+-]+', open(os.path.join(KDIR, fn), encoding='utf-8').read()):
                s.add(norm_id(d.rstrip('.,;)')))
    return s


def src_pubmed(q, n):
    out = []
    try:
        ids = jget('https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi?db=pubmed&retmode=json&retmax=%d&term=%s'
                   % (n, urllib.parse.quote(q))).get('esearchresult', {}).get('idlist', [])
    except Exception:
        return out
    if not ids:
        return out
    xml = http('https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi?db=pubmed&retmode=xml&id=%s' % ','.join(ids))
    for art in re.findall(r'<PubmedArticle>.*?</PubmedArticle>', xml, re.S):
        t = re.search(r'<ArticleTitle>(.*?)</ArticleTitle>', art, re.S)
        j = re.search(r'<Title>(.*?)</Title>', art, re.S)
        y = re.search(r'<PubDate>.*?<Year>(\d{4})</Year>', art, re.S)
        doi = re.search(r'<ArticleId IdType="doi">(.*?)</ArticleId>', art)
        pmc = re.search(r'<ArticleId IdType="pmc">(.*?)</ArticleId>', art)
        pmid = re.search(r'<PMID[^>]*>(\d+)</PMID>', art)
        out.append({'src': 'pubmed', 'title': re.sub(r'<[^>]+>', '', t.group(1)).strip() if t else '',
                    'journal': re.sub(r'<[^>]+>', '', j.group(1)).strip() if j else '',
                    'year': y.group(1) if y else '', 'doi': doi.group(1) if doi else '',
                    'pmcid': pmc.group(1) if pmc else '', 'pmid': pmid.group(1) if pmid else '',
                    'oa': bool(pmc)})
    return out


def src_crossref(q, n):
    d = jget('https://api.crossref.org/works?rows=%d&mailto=%s&query.bibliographic=%s'
             % (n, MAIL, urllib.parse.quote(q)))
    out = []
    for it in (d.get('message') or {}).get('items', []) or []:
        out.append({'src': 'crossref', 'title': ((it.get('title') or [''])[0])[:220],
                    'journal': ((it.get('container-title') or [''])[0])[:90],
                    'year': str(((it.get('issued') or {}).get('date-parts') or [['']])[0][0] or ''),
                    'doi': it.get('DOI', ''), 'cited': it.get('is-referenced-by-count', 0),
                    'publisher': (it.get('publisher') or '')[:40], 'oa': False})
    return out


def src_openalex(q, n):
    d = jget('https://api.openalex.org/works?per-page=%d&mailto=%s&search=%s'
             % (n, MAIL, urllib.parse.quote(q)))
    out = []
    for it in (d.get('results') or []):
        out.append({'src': 'openalex', 'title': (it.get('title') or '')[:220],
                    'journal': (((it.get('primary_location') or {}).get('source') or {}) or {}).get('display_name', '')[:90],
                    'year': str(it.get('publication_year') or ''),
                    'doi': (it.get('doi') or '').replace('https://doi.org/', ''),
                    'cited': it.get('cited_by_count', 0),
                    'oa': bool((it.get('open_access') or {}).get('is_oa')),
                    'oa_url': (it.get('best_oa_location') or {}).get('pdf_url') or
                              (it.get('open_access') or {}).get('oa_url') or ''})
    return out


def merge(*lists):
    seen, rows = {}, []
    for lst in lists:
        for r in lst:
            k = norm_id(r['doi']) or norm_id(r['title'])[:60]
            if not k:
                continue
            if k in seen:
                cur = seen[k]
                cur['src'] = cur['src'] + '/' + r['src']
                for f in ('pmcid', 'pmid', 'oa_url', 'journal', 'year', 'cited'):
                    if not cur.get(f) and r.get(f):
                        cur[f] = r[f]
                cur['oa'] = cur.get('oa') or r.get('oa')
            else:
                seen[k] = dict(r)
                rows.append(seen[k])
    return rows


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('source', choices=['pubmed', 'crossref', 'openalex', 'all'])
    ap.add_argument('query')
    ap.add_argument('--n', type=int, default=15)
    ap.add_argument('--json', action='store_true')
    ap.add_argument('--only-new', action='store_true', help='只显示本地尚未覆盖的')
    a = ap.parse_args()

    if a.source == 'pubmed':
        rows = src_pubmed(a.query, a.n)
    elif a.source == 'crossref':
        rows = src_crossref(a.query, a.n)
    elif a.source == 'openalex':
        rows = src_openalex(a.query, a.n)
    else:
        rows = merge(src_pubmed(a.query, a.n), src_crossref(a.query, a.n), src_openalex(a.query, a.n))
        time.sleep(0.3)

    NOISE_D = re.compile(r'(\.s\d{3}$|/table-|/fig-|/suppl|/media/|/t\d+$|/review\d*$|/v\d+/review|/peer-review)')
    NOISE_T = re.compile(r'^(table|figure|fig\.|supplementary|supp\.|appendix|scheme|review for|author response|decision letter|peer review)\b', re.I)
    rows = [r for r in rows
            if not NOISE_T.match((r.get('title') or '').strip())
            and not NOISE_D.search(r.get('doi') or '')
            and len((r.get('title') or '').strip()) > 25]
    known = known_dois()
    for r in rows:
        r['local'] = norm_id(r.get('doi', '')) in known if r.get('doi') else False
        r['title'] = _html.unescape(re.sub(r'\s+', ' ', (r.get('title') or '')).strip())
        r['journal'] = _html.unescape((r.get('journal') or '')).strip()
    if a.only_new:
        rows = [r for r in rows if not r['local']]
    if a.json:
        print(json.dumps(rows, ensure_ascii=False, indent=1))
        return
    new = [r for r in rows if not r['local']]
    oa = [r for r in rows if r.get('oa')]
    print('查询「%s」　源=%s　召回 %d 条｜开放获取 %d 条｜本地未覆盖 %d 条\n' % (a.query, a.source, len(rows), len(oa), len(new)))
    for i, r in enumerate(rows, 1):
        flag = '★新' if not r['local'] else ' 有'
        oaf = 'OA' if r.get('oa') else '  '
        print('%2d. [%s][%s] %s' % (i, flag, oaf, r['title'][:96]))
        print('     %s %s %s | %s' % (r.get('journal', '')[:38], r.get('year', ''), r.get('doi', ''),
                                      r.get('src', '')))
    if new:
        print('\n★ 本地未覆盖的 %d 条（可进下一轮深调研候选）：' % len(new))
        for r in new:
            print('   - %s (%s %s) %s' % (r['title'][:80], r.get('journal', '')[:30], r.get('year', ''), r.get('doi', '')))


if __name__ == '__main__':
    main()