import marimo


__generated_with = "0.17.7"
app = marimo.App(width="medium")


@app.cell
def _():
    """
    标题：使用 Pydantic AI 批量分析 PDF 文档

    描述：
    本 notebook 演示如何使用 Pydantic AI 批量分析 PDF 文档内容。
    使用 pypdfium2 进行 PDF 文本提取，配合 Pydantic AI 进行结构化内容分析。

    学习目标：
    1. 如何使用 pypdfium2 提取 PDF 文本
    2. 如何配置 AnalysisConfig 进行 PDF 内容分析
    3. 如何批量异步处理多个 PDF 文件
    """
    import json
    import os
    from pathlib import Path

    import pypdfium2 as pdfium  # PDF 文本提取库
    import marimo as mo
    import polars as pl

    from information_composer.pubmed.analyzer import (
        AnalysisConfig,
        PaperAnalyzer,
        PaperInput,
    )

    # 从环境变量读取配置
    openai_base_url = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
    openai_api_key = os.getenv("OPENAI_API_KEY", "")

    # 文件浏览器 - 支持多选
    file_browser = mo.ui.file_browser(
        initial_path="/Users/forrest/Downloads/uORF文章汇总/",
        multiple=True,
        filetypes=[".pdf"],
    )

    mo.md(f"""
    ## 📚 批量 PDF 文档分析工具（Pydantic AI 版本）

    使用 Pydantic AI 进行智能 PDF 内容分析。

    **技术栈**:
    - PDF 解析: pypdfium2 (轻量、高效的 PDF 处理库)
    - LLM 调用: Pydantic AI (类型安全、结构化输出)

    {file_browser}
    """)
    return (file_browser, mo, pl, pdfium, json, openai_api_key, openai_base_url)


@app.cell
def _(file_browser, mo):
    if file_browser.value and len(file_browser.value) > 0:
        file_list = "\n".join([f"- `{f.name}`" for f in file_browser.value])
        mo.md(f"""
        ### ✅ 已选择的 PDF 文件:

        {file_list}
        """)
    else:
        mo.md("📂 请选择至少一个 PDF 文件")
    return


@app.cell
def _(pdfium):
    def extract_pdf_text(pdf_path: str, max_chars: int = 80000) -> str:
        """使用 pypdfium2 提取 PDF 文本"""
        try:
            pdf = pdfium.PdfDocument(pdf_path)
            text_parts = []

            for page_num in range(len(pdf)):
                page = pdf[page_num]
                textpage = page.get_textpage()
                text_parts.append(textpage.get_text())

            full_text = "\n\n".join(text_parts)

            # 截断长文本
            if len(full_text) > max_chars:
                full_text = full_text[:max_chars] + "\n\n... (文本已截断)"

            return full_text
        except Exception as e:
            return f"PDF 解析错误: {str(e)}"

    return (extract_pdf_text,)


@app.cell
def _(extract_pdf_text, json, mo):
    async def analyze_pdf_with_pydantic_ai(
        pdf_file, question: str, config: dict, semaphore, progress_list
    ):
        """使用 Pydantic AI 分析单个 PDF"""
        filename = pdf_file.name

        async with semaphore:
            try:
                # 提取 PDF 文本
                pdf_text = extract_pdf_text(str(pdf_file.path))

                if pdf_text.startswith("PDF 解析错误"):
                    progress_list.append(f"❌ {filename}")
                    return {"文件名": filename, "错误": pdf_text}

                # 使用 Pydantic AI 分析
                from information_composer.pubmed.analyzer import (
                    AnalysisConfig,
                    PaperAnalyzer,
                    PaperInput,
                )

                analysis_config = AnalysisConfig(
                    analysis_types=["summary"],
                    llm_provider=config.get("provider", "openai"),
                    llm_model=config.get("model", "claude-sonnet-4.5"),
                    use_pydantic_ai=True,
                )

                analyzer = PaperAnalyzer(analysis_config)

                # 构造分析输入（将 PDF 内容作为 abstract）
                paper = PaperInput(
                    pmid=f"pdf_{filename}",
                    title=f"PDF: {filename}",
                    abstract=f"用户问题: {question}\n\nPDF 内容:\n{pdf_text[:10000]}",
                )

                result = await analyzer.analyze_paper(paper)

                # 提取结构化结果
                parsed_result = {
                    "文件名": filename,
                    "是否与uORF有关": "是" if "uORF" in pdf_text.lower() else "待判定",
                    "核心发现": "; ".join(result.summary.main_findings[:3])
                    if result.summary
                    else "",
                    "token使用量": result.processing_metadata.tokens_used or 0,
                }

                progress_list.append(filename)
                return parsed_result

            except Exception as e:
                progress_list.append(f"❌ {filename}")
                return {"文件名": filename, "错误": str(e)}

    async def batch_analyze_pdfs(pdf_files, question, max_concurrent=3, config=None):
        """批量异步处理 PDF"""
        import asyncio

        if config is None:
            config = {"provider": "openai", "model": "claude-sonnet-4.5"}

        semaphore = asyncio.Semaphore(max_concurrent)
        progress_list = []

        tasks = [
            analyze_pdf_with_pydantic_ai(f, question, config, semaphore, progress_list)
            for f in pdf_files
        ]

        return await asyncio.gather(*tasks)

    return (batch_analyze_pdfs,)


@app.cell
async def _(batch_analyze_pdfs, file_browser, mo):
    # 分析问题
    question = "这篇论文是否与 uORF 相关？如果是请给出相关信息。"

    if file_browser.value and len(file_browser.value) > 0:
        with mo.status.spinner(
            title=f"🤔 正在使用 Pydantic AI 分析 {len(file_browser.value)} 个文件..."
        ):
            results = await batch_analyze_pdfs(
                file_browser.value, question, max_concurrent=3
            )
    else:
        results = []

    return (results,)


@app.cell
def _(mo):
    mo.md("""
    ### Pydantic AI 优势说明

    1. **结构化输出**: 分析结果自动转换为 Pydantic 模型，字段类型安全
    2. **Token 追踪**: 自动记录每个 PDF 的 token 使用量，便于成本控制
    3. **错误处理**: 验证错误自动捕获，不会出现 JSON 解析失败
    4. **性能**: pypdfium2 是高效的 PDF 处理库，支持快速文本提取
    """)
    return


@app.cell
def _(mo, pl, results):
    # 展示结果表格
    df_results = None
    if results:
        df_results = pl.DataFrame(results)

        mo.vstack(
            [
                mo.md("### 📈 分析结果表格"),
                mo.ui.table(df_results),
            ]
        )
    else:
        mo.md("")
    return (df_results,)


@app.cell
def _(df_results):
    df_results
    return


@app.cell
def _(df_results, mo):
    # 导出功能
    try:
        if df_results is not None and len(df_results) > 0:
            export_button = mo.ui.button(label="📥 导出为 CSV")

            mo.vstack(
                [
                    mo.md("### 💾 导出数据"),
                    export_button,
                ]
            )

            if export_button.value:
                try:
                    # 保存 CSV
                    output_path = "pdf_analysis_results.csv"
                    df_results.write_csv(output_path)

                    mo.md(f"""
                    ✅ **导出成功！**

                    文件已保存至: `{output_path}`

                    你可以在文件系统中找到这个 CSV 文件。
                    """)
                except Exception as e:
                    mo.md(f"❌ 导出失败: {str(e)}")
        else:
            mo.md("")
    except:
        mo.md("")
    return


if __name__ == "__main__":
    app.run()
