app.py

import streamlit as st
import re
import requests
import xml.etree.ElementTree as ET
import os
import tempfile
from win32com.client import Dispatch
import pythoncom  # 新增：用于初始化 COM
import time

# ===== 初始化 COM 库（必须放在所有 Dispatch 调用之前）=====
pythoncom.CoInitialize()

# 设置网页标题和图标
st.set_page_config(page_title="PMID 自动化处理工具", page_icon="📚")
st.title("📚 拖拽 Word 一键处理 PMID")
st.write("上传你的 Word 文档，自动提取 PMID、生成 EndNote 导入文件，并自动添加批注。")

# 创建文件上传区域（支持拖拽）
uploaded_file = st.file_uploader("将 Word 文档拖拽到这里", type=['docx'])

if uploaded_file is not None:
    # 1. 保存用户上传的文件到临时目录
    with tempfile.NamedTemporaryFile(delete=False, suffix='.docx') as tmp_input:
        tmp_input.write(uploaded_file.getvalue())
        input_path = tmp_input.name

    # 2. 提取 PMID
    with st.spinner('正在提取 PMID...'):
        word = Dispatch('Word.Application')
        word.Visible = False
        doc = word.Documents.Open(input_path)
        full_text = doc.Range().Text
        doc.Close(False)
        word.Quit()
        
        pattern = r'[Pp][Mm][Ii][Dd]\s*[:：]\s*(\d+)'
        pmid_list = list(dict.fromkeys(re.findall(pattern, full_text)))
        st.success(f"✅ 提取成功！共找到 {len(pmid_list)} 个 PMID")

    if pmid_list:
        # 3. 获取 PubMed 数据并生成 RIS
        with st.spinner('正在从 PubMed 抓取文献信息...'):
            base_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi"
            params = {"db": "pubmed", "id": ",".join(pmid_list), "retmode": "xml"}
            response = requests.get(base_url, params=params, timeout=60)
            root = ET.fromstring(response.text)
            
            ris_lines = []
            for article in root.findall(".//PubmedArticle"):
                pmid = article.find(".//PMID").text
                title = article.find(".//ArticleTitle")
                title = title.text if title is not None else "No Title"
                journal = article.find(".//Title")
                journal = journal.text if journal is not None else "No Journal"
                year = article.find(".//PubDate/Year")
                year = year.text if year is not None else ""
                authors = []
                for author in article.findall(".//Author")[:3]:
                    last = author.find("LastName")
                    fore = author.find("ForeName")
                    if last is not None and fore is not None:
                        authors.append(f"{last.text} {fore.text}")
                author_str = ", ".join(authors)
                
                ris_lines.append("TY  - JOUR")
                ris_lines.append(f"TI  - {title}")
                ris_lines.append(f"T2  - {journal}")
                ris_lines.append(f"PY  - {year}")
                ris_lines.append(f"AU  - {author_str}")
                ris_lines.append(f"ID  - {pmid}")
                ris_lines.append(f"UR  - https://pubmed.ncbi.nlm.nih.gov/{pmid}/")
                ris_lines.append("ER  -")
                ris_lines.append("")
            
            ris_content = "\n".join(ris_lines)

        # 4. 提供 RIS 下载
        st.download_button(
            label="📥 下载 RIS 文件（导入 EndNote）",
            data=ris_content,
            file_name="my_references.ris",
            mime="text/plain"
        )

        # 5. 可选：自动添加批注（用复选框控制）
        if st.checkbox("同时自动添加批注并删除原文（谨慎操作，建议先备份）"):
            with st.spinner('正在处理批注...'):
                word = Dispatch('Word.Application')
                word.Visible = True  # 显示出来让你看到变化
                doc = word.Documents.Open(input_path)
                
                find_obj = word.Selection.Find
                find_obj.ClearFormatting()
                find_obj.Text = "\(PMID:*\)"
                find_obj.MatchWildcards = True
                find_obj.Forward = True
                find_obj.Wrap = 1
                
                count = 0
                while find_obj.Execute():
                    count += 1
                    matched_text = word.Selection.Text
                    pmid_number = matched_text[6:-1]
                    word.Selection.Text = ""
                    word.ActiveDocument.Comments.Add(word.Selection.Range, pmid_number)
                
                doc.Save()
                # 保存修改后的文件到临时目录
                output_path = input_path.replace('.docx', '_modified.docx')
                doc.SaveAs2(output_path)
                doc.Close()
                word.Quit()
                
                st.success(f"✅ 已添加 {count} 条批注！")
                
                # 提供修改后的 Word 下载
                with open(output_path, 'rb') as f:
                    st.download_button(
                        label="📥 下载修改后的 Word 文档（含批注）",
                        data=f,
                        file_name="modified_文档.docx",
                        mime="application/vnd.openxmlformats-officedocument.wordprocessingml"
                    )
                
                # 清理临时文件
                os.unlink(output_path)
        
        # 清理上传的临时文件
        os.unlink(input_path)
    else:
        st.warning("⚠️ 文档中未找到任何 PMID 格式的内容。")