app.py

import gradio as gr
import arxiv
import re
import os
import requests # 用于进行 HTTP 请求，获取 DOI 和通用 URL 元数据
import datetime # 用于获取当前日期作为访问日期
import hashlib # 用于为通用 URL 生成唯一的 BibTeX 键

# --- 核心逻辑函数 ---

def extract_identifiers(input_string):
    """
    从输入的字符串中提取 arXiv ID、DOI 或通用 URL。
    支持新旧格式的 arXiv 链接/ID、DOI 链接/ID 和常见的 URL 格式。
    返回一个包含 (type, id) 元组的列表，例如 [('arxiv', '2006.01234'), ('doi', '10.1007/s11227-020-03299-8'), ('web', 'https://example.com/blog')]
    """
    input_string = input_string.strip()
    identifiers = []

    # 尝试提取 arXiv ID
    arxiv_new_format_match = re.search(r'(?P<id>\d{4}\.\d{5})(v\d+)?', input_string)
    if arxiv_new_format_match:
        identifiers.append(('arxiv', arxiv_new_format_match.group('id')))
        return identifiers

    arxiv_old_format_match = re.search(r'(?P<id>[a-z\-]+/\d{7})(v\d+)?', input_string)
    if arxiv_old_format_match:
        identifiers.append(('arxiv', arxiv_old_format_match.group('id')))
        return identifiers

    if re.match(r'^\d{4}\.\d{5}(v\d+)?$', input_string):
        identifiers.append(('arxiv', input_string.split('v')[0]))
        return identifiers

    if re.match(r'^[a-z\-]+/\d{7}(v\d+)?$', input_string):
        identifiers.append(('arxiv', input_string.split('v')[0]))
        return identifiers

    # 尝试提取 DOI
    doi_match = re.search(r'(10\.\d{4,9}/[-._;()/:A-Z0-9]+)', input_string, re.IGNORECASE)
    if doi_match:
        identifiers.append(('doi', doi_match.group(1)))
        return identifiers

    # 如果以上都不是，尝试作为通用 URL
    # 简单的 URL 匹配，确保它看起来像一个 URL (http/https开头或www.开头)
    if re.match(r'^(https?://|www\.)[^\s/$.?#].[^\s]*$', input_string, re.IGNORECASE):
        identifiers.append(('web', input_string))
        return identifiers

    return identifiers

def get_bibtex_from_arxiv(arxiv_id):
    """
    根据 arXiv ID 获取 BibTeX 条目。
    """
    try:
        search_results = arxiv.Search(id_list=[arxiv_id]).results()
        paper = next(search_results)

        authors = " and ".join([author.name for author in paper.authors])
        title = paper.title.replace("\n", " ").strip()
        year = paper.published.year
        arxiv_category = paper.primary_category
        doi_field = f",\n  doi = {{{paper.doi}}}" if paper.doi else ""

        bib_key = arxiv_id.replace('.', '_').replace('/', '_').replace('-', '_')

        bibtex_entry = f"""
@article{{{bib_key},
  title={{ {title} }},
  author={{ {authors} }},
  journal={{ arXiv preprint arXiv:{paper.entry_id.split('/')[-1]} [{arxiv_category}] }},
  year={{ {year} }}{doi_field}
}}
"""
        return bibtex_entry, None # 返回 BibTeX 和 None (无错误)
    except StopIteration:
        return None, f"错误: 未找到 arXiv ID '{arxiv_id}' 的论文。请检查 ID 是否正确或已发布。"
    except Exception as e:
        return None, f"无法获取 arXiv ID '{arxiv_id}' 的 BibTeX 条目: {e}"

def get_bibtex_from_doi(doi):
    """
    根据 DOI 获取 BibTeX 条目，通过 CrossRef API。
    """
    try:
        # CrossRef API 直接提供 BibTeX 格式的转换
        url = f"https://api.crossref.org/v1/works/{doi}/transform/application/x-bibtex"
        headers = {'Accept': 'application/x-bibtex'} # 请求 BibTeX 格式
        response = requests.get(url, headers=headers, timeout=10) # 增加超时

        if response.status_code == 200:
            return response.text, None # 返回 BibTeX 文本和 None (无错误)
        elif response.status_code == 404:
            return None, f"错误: 未找到 DOI '{doi}' 的论文。请检查 DOI 是否正确。"
        else:
            return None, f"错误: 获取 DOI '{doi}' 的 BibTeX 失败，状态码: {response.status_code}，响应: {response.text[:100]}..."
    except requests.exceptions.RequestException as e:
        return None, f"网络请求错误，无法获取 DOI '{doi}' 的 BibTeX: {e}"
    except Exception as e:
        return None, f"处理 DOI '{doi}' 时发生未知错误: {e}"

def get_bibtex_from_url(url):
    """
    根据通用 URL 获取 BibTeX 条目。
    尝试从网页标题中提取信息，并生成 @misc 类型的 BibTeX。
    注意：从通用网页提取准确的元数据（作者、日期等）非常困难且不可靠。
    此函数主要提取标题和 URL，并记录访问日期。
    """
    try:
        # 确保 URL 包含协议，否则 requests 可能无法处理
        if not url.startswith('http://') and not url.startswith('https://'):
            url = 'http://' + url # 默认使用 http，requests 会自动重定向到 https

        response = requests.get(url, timeout=10)
        response.raise_for_status() # 对 4xx/5xx 响应抛出 HTTPError
        html_content = response.text

        # 尝试从 <title> 标签中提取标题
        title_match = re.search(r'<title>(.*?)</title>', html_content, re.IGNORECASE | re.DOTALL)
        title = title_match.group(1).strip() if title_match else url # 如果未找到标题，则使用 URL 作为标题

        # 尝试从 Open Graph 协议中提取标题 (优先级更高)
        og_title_match = re.search(r'<meta\s+property="og:title"\s+content="([^"]*)"', html_content, re.IGNORECASE)
        if og_title_match:
            title = og_title_match.group(1).strip()
        
        # 简单的 bib_key 生成，基于 URL 的 MD5 哈希值，取前8位
        bib_key = "web_" + hashlib.md5(url.encode('utf-8')).hexdigest()[:8]

        access_date = datetime.date.today().strftime("%Y-%m-%d")

        bibtex_entry = f"""
@misc{{{bib_key},
  title={{ {title} }},
  howpublished={{ \\url{{{url}}} }},
  note={{ Accessed: {access_date} }}
}}
"""
        return bibtex_entry, None
    except requests.exceptions.RequestException as e:
        return None, f"网络请求错误，无法获取 URL '{url}' 的内容: {e}"
    except Exception as e:
        return None, f"处理 URL '{url}' 时发生未知错误: {e}"


def get_bibtex_entries_and_display(input_string):
    """
    根据输入的字符串（每行一个 arXiv ID/链接、DOI/链接或通用 URL）获取 BibTeX 条目，
    返回日志信息和所有BibTeX条目的字符串形式。
    """
    input_lines = input_string.split('\n')
    
    processed_identifiers = [] # 存储 (type, id) 元组
    log_messages = []
    all_bibtex_content = [] # 存储所有BibTeX条目

    log_messages.append("正在从输入中提取 arXiv ID、DOI 或通用 URL...")
    for line in input_lines:
        line = line.strip()
        if not line:
            continue
        extracted = extract_identifiers(line)
        if extracted:
            for id_type, identifier in extracted:
                processed_identifiers.append((id_type, identifier))
        else:
            log_messages.append(f"警告: 无法从 '{line}' 中提取有效的 arXiv ID、DOI 或 URL。将跳过此项。")

    # 去重处理，确保每个唯一的 (type, id) 对只处理一次
    unique_identifiers = list(set(processed_identifiers))
    
    if not unique_identifiers:
        log_messages.append("没有检测到有效的 arXiv ID、DOI 或 URL。请检查输入。")
        return "\n".join(log_messages), "", None

    log_messages.append(f"\n成功提取到 {len(unique_identifiers)} 个唯一的标识符。正在获取 BibTeX 条目...")

    for id_type, identifier in unique_identifiers:
        bibtex_entry = None
        error_message = None

        if id_type == 'arxiv':
            log_messages.append(f"正在获取 arXiv ID '{identifier}' 的 BibTeX 条目...")
            bibtex_entry, error_message = get_bibtex_from_arxiv(identifier)
        elif id_type == 'doi':
            log_messages.append(f"正在获取 DOI '{identifier}' 的 BibTeX 条目...")
            bibtex_entry, error_message = get_bibtex_from_doi(identifier)
        elif id_type == 'web': # 新增的通用 URL 处理
            log_messages.append(f"正在获取通用 URL '{identifier}' 的 BibTeX 条目...")
            bibtex_entry, error_message = get_bibtex_from_url(identifier)
        else:
            error_message = f"未知标识符类型: {id_type} (ID: {identifier})"

        if bibtex_entry:
            all_bibtex_content.append(bibtex_entry)
            log_messages.append(f"成功获取 {id_type.upper()} '{identifier}' 的 BibTeX 条目。")
        else:
            log_messages.append(error_message)
            
    final_bibtex_string = "\n".join(all_bibtex_content)

    temp_bib_filepath = None
    if final_bibtex_string.strip():
        temp_bib_filepath = "temp_references.bib" # 修改文件名以反映更广泛的来源
        try:
            with open(temp_bib_filepath, 'w', encoding='utf-8') as f:
                f.write(final_bibtex_string)
            log_messages.append(f"\nBibTeX 内容已保存到临时文件 '{temp_bib_filepath}' 以供下载。")
        except Exception as e:
            log_messages.append(f"错误: 无法将BibTeX内容写入临时文件: {e}")
            temp_bib_filepath = None

    if not final_bibtex_string.strip():
        log_messages.append("\n没有成功获取任何 BibTeX 条目。")

    return "\n".join(log_messages), final_bibtex_string, temp_bib_filepath

# --- Gradio 界面部分 ---

# 定义 Gradio 界面
iface = gr.Interface(
    fn=get_bibtex_entries_and_display,
    inputs=[
        gr.Textbox(
            lines=10, 
            label="粘贴 arXiv 链接/ID、DOI 或通用 URL (每行一个)", 
            placeholder="例如:\nhttps://arxiv.org/abs/2006.01234\n1904.08133v1\n10.1109/TNNLS.2020.3006789\nhttps://doi.org/10.1038/s41586-023-06670-w\nhttps://www.example.com/blog-post\nwww.anotherblog.org/article"
        )
    ],
    outputs=[
        gr.Textbox(label="处理日志", lines=8, interactive=False),
        gr.Textbox(label="在此复制 BibTeX 内容", lines=10, interactive=True),
        gr.File(label="下载 BibTeX 文件")
    ],
    title="Anything2Bib: 通用 BibTeX 获取工具",
    description="在此粘贴 arXiv 论文的链接/ID、DOI 或通用网页链接 (每行一个)，获取其 BibTeX 条目。",
    allow_flagging="never"
)

# 启动 Gradio 应用
if __name__ == "__main__":
    print("Gradio 应用正在启动...")
    print("你可以在浏览器中访问以下地址:")
    iface.launch(share=False)
    print("Gradio 应用已关闭。")