# -*- coding: utf-8 -*-
"""接口测试

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1sTqCE6RFjGND2PA3c52cQjx4EV1SXJsr
"""

!pip install --no-deps unsloth_zoo
!pip install --no-deps unsloth
!pip install trl
!pip install -U bitsandbytes
!pip install pymupdf

import os
os.kill(os.getpid(), 9)

from huggingface_hub import login
login()  # 或设置环境变量 HF_TOKEN

from transformers import pipeline
from unsloth import FastLanguageModel

# 2. 导入必要包
import gradio as gr
from transformers import pipeline
from unsloth import FastLanguageModel  # ✅ 插入这里

import torch
import torch
from transformers import AutoTokenizer, pipeline
from unsloth import FastLanguageModel

model_name = "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"

# 全局模型与tokenizer（只加载一次）
tokenizer = AutoTokenizer.from_pretrained(model_name)

model = None
generator = None

def load_model_dynamic_maxlen(example_text):
    global model, generator
    tokenized_len = len(tokenizer.encode(example_text))
    buffer_len = 512
    max_limit = 4096
    max_seq_length = min(tokenized_len + buffer_len, max_limit)

    print(f"[✓] 动态设置 max_seq_length 为: {max_seq_length}")

    model, _ = FastLanguageModel.from_pretrained(
        model_name=model_name,
        max_seq_length=max_seq_length,
        dtype=torch.float16,
        load_in_4bit=True
    )

    generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")
    return generator

# 全局模型与tokenizer（只加载一次）
model_name = "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"

tokenizer = AutoTokenizer.from_pretrained(model_name)

# 设置最大长度一次性加载
max_seq_length = 4096

model, _ = FastLanguageModel.from_pretrained(
    model_name=model_name,
    max_seq_length=max_seq_length,
    dtype=torch.float16,
    load_in_4bit=True
)

generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")

def generate_suggestion_llama(prompt_text):
    inputs = tokenizer(prompt_text, return_tensors="pt").to(model.device)
    generation_output = model.generate(
        **inputs,
        max_new_tokens=512,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        eos_token_id=tokenizer.eos_token_id
    )
    output_text = tokenizer.decode(generation_output[0], skip_special_tokens=True)
    # 通常模型会把prompt也返回，切掉prompt部分，只保留生成
    return output_text[len(prompt_text):].strip()

# 1. 安装依赖
!pip install gradio transformers accelerate

# 2. 导入必要包
import gradio as gr
from transformers import pipeline

# 先写个PDF转文本的简单函数
def pdf_to_text(pdf_file):
    try:
        # 这里写你的PDF解析代码，比如用 PyMuPDF / pdfplumber / PyPDF2 等
        # demo示例：
        import fitz  # PyMuPDF
        doc = fitz.open(pdf_file.name)
        text = ""
        for page in doc:
            text += page.get_text()
        return text
    except Exception as e:
        print(f"PDF解析错误: {e}")
        return ""  # 或者返回提示文本


# 先写个PDF转文本的简单函数
def extract_title_author_from_pdf(pdf_file):
    import fitz  # PyMuPDF
    try:
        doc = fitz.open(pdf_file.name)

        # 先尝试读取PDF元数据
        metadata = doc.metadata
        title = metadata.get('title', '').strip()
        author = metadata.get('author', '').strip()

        # 如果元数据里没有，尝试从第一页内容提取
        if not title or not author:
            first_page_text = doc[0].get_text("text")
            lines = [line.strip() for line in first_page_text.split('\n') if line.strip()]
            if not title and len(lines) > 0:
                title = lines[0]  # 假设第一页第一行是标题
            if not author and len(lines) > 1:
                author = lines[1]  # 假设第二行是作者信息

        # 返回标题、作者和全文文本
        full_text = ""
        for page in doc:
            full_text += page.get_text("text") + "\n"

        return title, author, full_text
    except Exception as e:
        print(f"PDF解析错误: {e}")
        return "", "", ""

# 在import部分新增
from typing import List, Dict, Tuple
import re

# 新增高亮处理函数
def highlight_severity(text, suggestions):
    severity_colors = {
        "High": "red",
        "Medium": "orange",
        "Low": "blue"
    }

    color_map = {}
    marked_text = text

    # Handle both string (single suggestion) and list (multiple suggestions) inputs
    suggestions_iter = [suggestions] if isinstance(suggestions, dict) else suggestions

    for suggestion in suggestions_iter:
        excerpt = suggestion.get("excerpt", "")
        severity = suggestion.get("severity", "Low")
        color = severity_colors.get(severity, "black")
        if excerpt in marked_text and excerpt not in color_map:
            span = f"<span style='background-color:{color};'>{excerpt}</span>"
            marked_text = marked_text.replace(excerpt, span, 1)
            color_map[excerpt] = color

    if not color_map:
        marked_text = f"<div style='padding:10px;'>{text}</div>"

    return marked_text, color_map

from google.colab import drive
drive.mount('/content/drive')

import os
base_path = "/content/drive/MyDrive/Colab Notebooks"
print(os.listdir(base_path))

import os
model_path = "/content/drive/MyDrive/Colab Notebooks/longformer_dmp_final_model"
print("目录是否存在:", os.path.exists(model_path))

import os
model_path = "/content/drive/MyDrive/Colab Notebooks/longformer_dmp_final_model"

print("目录是否存在:", os.path.exists(model_path))
print("目录内容:", os.listdir(model_path))

from transformers import AutoTokenizer, AutoModelForSequenceClassification

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)

print("模型和分词器加载成功！")

!ls /content/drive/MyDrive/Colab\ Notebooks/

# 2. 导入必要包
import gradio as gr
from transformers import pipeline

# 先写个PDF转文本的简单函数
def pdf_to_text(pdf_file):
    try:
        # 这里写你的PDF解析代码，比如用 PyMuPDF / pdfplumber / PyPDF2 等
        # demo示例：
        import fitz  # PyMuPDF
        doc = fitz.open(pdf_file.name)
        text = ""
        for page in doc:
            text += page.get_text()
        return text
    except Exception as e:
        print(f"PDF解析错误: {e}")
        return ""  # 或者返回提示文本

# 定义统一的标签映射，分类名称 <-> 数字标签
num_to_category = {
    0: "Good",
    1: "Not Good",
    2: "Intermediate DMP Needed",
    3: "Good with Suggestions"
}

category_map = {
    "Good": 0,
    "Not Good": 1,
    "Intermediate DMP Needed": 2,
    "Good with Suggestions": 3,
    "Other": 1,  # 仍然可以保留“其他”映射到1，即 Not Good
}

# === 新增：Longformer 分类模型 ===
from transformers import AutoModelForSequenceClassification, AutoTokenizer

longformer_model_name = "allenai/longformer-base-4096"  # 你的模型路径
longformer_tokenizer = AutoTokenizer.from_pretrained(longformer_model_name)
longformer_model = AutoModelForSequenceClassification.from_pretrained(longformer_model_name)

from transformers import pipeline
longformer_classifier = pipeline(
    "text-classification",
    model=longformer_model,
    tokenizer=longformer_tokenizer,
    device_map="auto",
    truncation=True,
    max_length=4096
)

# 3. 模拟分类模型（用现成pipeline替代，后续换你的Longformer）
# === 修改：用Longformer替代关键词判断 ===
def classify_text(text):
    # 将长文本切成大约 3500 tokens 的块
    tokens = longformer_tokenizer.encode(text, truncation=False)
    chunk_size = 3500  # 预留余量，避免超出 max_length=4096
    chunks = [tokens[i:i + chunk_size] for i in range(0, len(tokens), chunk_size)]

    # 初始化分数字典
    label_scores = {0: 0, 1: 0, 2: 0, 3: 0}

    # 标签映射：将模型输出的 label 转为对应的数字分类
    label_map = {
        "good": 0,
        "not good": 1,
        "intermediate dmp needed": 2,
        "good with suggestions": 3,
    }

    for chunk in chunks:
        chunk_text = longformer_tokenizer.decode(chunk)
        result = longformer_classifier(chunk_text, top_k=None)[0]

        label = result["label"].lower().strip()  # 转小写并去除可能的空格
        score = result["score"]

        # 如果标签未被识别，则默认归为“not good”（1）
        label_idx = label_map.get(label, 1)
        label_scores[label_idx] += score

    # 返回得分最高的标签编号
    final_label = max(label_scores, key=label_scores.get)
    return final_label

# 4. 模拟结构化建议生成（用Meta-Llama调用示范）
def generate_suggestion(text, category_label):
    prompt = f"""
I want you to act as a Research Data Management expert.
You will be provided with:

1. An excerpt from a Data Management Plan (DMP) document:
{text}

2. A quality label for the DMP:
{category_label}  # 这里是数字标签，确保和分类模型一致

Your task is to analyze the given DMP excerpt and explain the reasons behind the provided quality label. Provide clear, constructive, and actionable feedback to help improve the DMP where needed.

Please organize your output in the following format:

# Summary Diagnosis
- Briefly interpret the overall quality based on the given label, naming the main strengths and/or weaknesses regardless of label value.
- Highlight the top 2–3 aspects that most influenced the categorization.

# Structured Analysis by Dimension
Analyze the DMP excerpt along these dimensions, presented in descending order of importance:
- Data Description
- Metadata and Documentation
- Ethics and Privacy
- Storage and Backup
- Legal and Security
- Data Sharing
- Preservation
- Responsibilities

For **each dimension where relevant issues or observations exist**, provide:
- **Severity Level:** (e.g., Low, Medium, High) — the importance or urgency of addressing this point.
- **Core Keywords:** 3 to 5 keywords that summarize the main focus of the issue or strength.
- **Core Issue Summary:** A concise explanation of the strengths or weaknesses seen in this area.
- **Supporting Excerpts:** Relevant quotes or sentences taken directly from the *DMP text itself* (not reviewer comments), which illustrate your point.
- **Improvement Suggestions:** Specific, practical, and encouraging recommendations for enhancement.

# Other Notes (if applicable)
If there are supplementary comments, minor issues, or helpful suggestions that do not fit neatly into the above dimensions, briefly list them here.

---

**Important:** Please maintain a positive and encouraging tone throughout your feedback to motivate and guide improvement.
**Note:** Do not explicitly mention or describe your tone or style in the output. Just apply a constructive and supportive wording naturally.
**Note:** Do NOT repeat disclaimers or polite closings. Only provide concise, actionable feedback.
---

Please treat each DMP and label you receive individually, applying this structured approach consistently.
"""

    # 这里调用你模型接口或transformers pipeline的代码，传入prompt，获得生成结果
    # 例如：
    # response = llama_model.generate(prompt)
    # return response

    outputs = generator(prompt, max_new_tokens=2000, do_sample=True, temperature=0.7, top_p=0.9)
    result = outputs[0]['generated_text']
    if result.startswith(prompt):
        result = result[len(prompt):].strip()
    elif prompt.strip() in result:
        result = result.split(prompt.strip(), 1)[-1].strip()
    return result

# 在import后，紧接着加入：
reviewed_docs = []  # 存放已审核结果的列表
docs_to_review = []  # 批量待审核的文档列表
current_index = 0  # 当前审核文档索引

# 主处理函数，输入PDF文件或文本，输出分类和建议

# 修改process_document函数，添加高亮处理

# Modify the process_document function to return proper results
def process_document(pdf_file, text_input, selected_category):
    if pdf_file is not None:
        title, author, text = extract_title_author_from_pdf(pdf_file)
        if not text.strip():
            return "❌ PDF解析失败或内容为空", "", "", "请上传有效的PDF文件或输入文本。", ""
    else:
        title, author = "", ""
        text = text_input
        if not text.strip():
            return "❌ 无输入文本", "", "", "请输入文本或上传PDF文件。", ""

    # Classification
    if selected_category == "Auto Classification":
        category_num = classify_text(text)
    else:
        category_num = category_map.get(selected_category, 1)

    category_str = num_to_category.get(category_num, str(category_num))

    # Generate suggestions with proper highlighting
    try:
        # GPT 生成建议
        suggestion = generate_suggestion(text, category_num)  # 返回的是字符串，真实建议文本

        # 解析建议（可以不改）
        # ...

        # 生成带高亮的文本
        highlighted_text, _ = highlight_severity(text, suggestions_list)
    except Exception as e:
        suggestion = f"❌ 生成建议时出错: {e}"
        highlighted_text = text

    # 返回给前端：分类标签, 建议文本, 原文, 文档信息, 高亮文本
    # 其中 category_str 是类似 "Good" 的标签
    return category_str, suggestion, text, f"标题: {title}\n作者: {author}", highlighted_text

# 2.1 新增批量加载文档函数，替换你主函数里对文件和文本读取逻辑：
# 改进PDF元数据提取
def load_documents(pdf_files):
    """Load all PDF files and extract text with better metadata handling"""
    import fitz
    texts = []
    doc_infos = []

    for pdf_file in pdf_files:
        try:
            doc = fitz.open(pdf_file.name)
            text = ""
            for page in doc:
                text += page.get_text() + "\n"

            # 改进的元数据提取
            metadata = doc.metadata
            title = metadata.get("title", "").strip() or pdf_file.name.split('/')[-1]
            author = metadata.get("author", "").strip() or "Unknown"

            # 尝试从第一页提取标题
            if not title or title == pdf_file.name:
                first_page = doc[0].get_text("text")
                possible_title = first_page.split('\n')[0].strip()
                if len(possible_title) > 3 and len(possible_title) < 100:
                    title = possible_title

            doc_infos.append({
                "title": title,
                "author": author,
                "text": text
            })
            doc.close()
        except Exception as e:
            print(f"Error processing {pdf_file.name}: {str(e)}")

    return doc_infos

# 单条文档分类+生成建议（调整了参数，不再传selected_category）：

def classify_and_suggest(doc):
    category_num = classify_text(doc["text"])
    category_str = num_to_category.get(category_num, "Unknown")

    suggestion = generate_suggestion(doc["text"], category_num)

    return category_str, suggestion

# 2.3 保存当前审核结果：
def save_review(doc, manual_label_value, edited_suggestion):
    reviewed_docs.append({
        "title": doc.get("title", "Untitled"),
        "author": doc.get("author", "Unknown"),
        "text": doc["text"],
        "final_label": manual_label_value,
        "suggestion": edited_suggestion
    })
    return f"Document '{doc.get('title', 'Untitled')}' saved"

# ===== 新增全局变量和导入 =====
import json
import pandas as pd
from datetime import datetime

saved_data = []  # 存储所有保存的文档数据
export_format = gr.Radio(  # 导出格式选择器
    choices=["CSV", "JSON", "Excel"],
    value="CSV",
    label="Export Format"
)

# # 6. Gradio界面搭建

# 额外：保存结果按钮，导出JSON示范（简易）
# 导出JSON
# Export JSON
def export_json(text, category, suggestion):
    import json
    result = {
        "text": text,
        "category": category,
        "suggestion": suggestion
    }
    return json.dumps(result, ensure_ascii=False, indent=2)

# 替换当前文档来源为 Title & Author
# 启动批量审查
# 修改start_review函数，添加高亮处理
def start_review(pdf_files):
    if not pdf_files:
        return [gr.update()] * 10  # 匹配10个output

    import fitz
    docs = []
    titles = []
    title_author_list = []

    for pdf_file in pdf_files:
        doc = fitz.open(pdf_file.name)
        text = "\n".join(page.get_text() for page in doc)

        # 改进标题/作者提取
        metadata = doc.metadata
        title = metadata.get("title", "").strip()
        author = metadata.get("author", "").strip()
        first_page = doc[0].get_text("text")
        lines = [line.strip() for line in first_page.split('\n') if line.strip()]
        if not title and lines:
            title = lines[0]
        if not author and len(lines) > 1:
            author = lines[1]

        title = title or pdf_file.name.split("/")[-1]
        author = author or "Unknown"
        full_title_author = f"{title} - {author}"

        docs.append(text)
        titles.append(full_title_author)
        title_author_list.append(full_title_author)

    # 处理第一个文档
    doc_text = docs[0]
    category_num = classify_text(doc_text)
    category_str = num_to_category.get(category_num, "Unknown")
    suggestion = generate_suggestion(doc_text, category_num)

    # 解析高亮建议
    suggestions_list = []
    severity_pattern = r"\*\*Severity Level:\*\* (High|Medium|Low)"
    excerpt_pattern = r"\*\*Supporting Excerpts:\*\* (.*?)(?=\n\*\*|\n\n|$)"
    severities = re.findall(severity_pattern, suggestion)
    excerpts = re.findall(excerpt_pattern, suggestion, re.DOTALL)

    for sev, exc in zip(severities, excerpts):
        suggestions_list.append({
            "excerpt": exc.strip(),
            "severity": sev
        })

    highlighted_text, _ = highlight_severity(doc_text, suggestions_list)

    page_info_str = f"Document 1/{len(docs)}"

    return [
        titles[0],         # Title & Author
        category_str,      # Auto category
        suggestion,        # Suggestion text
        category_str,      # Manual label default
        page_info_str,     # Page info
        highlighted_text,  # Left side content
        docs,              # all texts
        titles,            # titles
        0,                 # current index
        1                  # page num (placeholder)
    ]

# 保存并切换到下一个

# 修改save_and_next函数，添加高亮处理
def save_and_next(manual_label_value, edited_suggestion):
    global current_index

    if current_index >= len(docs_to_review):
        return gr.update(), gr.update(), gr.update(), gr.update(), "No more documents", gr.update()

    # Save current
    doc = docs_to_review[current_index]
    save_msg = save_review(doc, manual_label_value, edited_suggestion)
    current_index += 1

    if current_index >= len(docs_to_review):
        return gr.update(), gr.update(), gr.update(), gr.update(), "All documents reviewed", gr.update()

    # Process next
    doc = docs_to_review[current_index]
    category_num = classify_text(doc["text"])
    category_str = num_to_category.get(category_num, "Unknown")
    suggestion = generate_suggestion(doc["text"], category_num)

    # Highlight
    suggestions_list = []
    severity_pattern = r"\*\*Severity Level:\*\* (High|Medium|Low)"
    excerpt_pattern = r"\*\*Supporting Excerpts:\*\* (.*?)(?=\n\*\*|\n\n|$)"
    severities = re.findall(severity_pattern, suggestion)
    excerpts = re.findall(excerpt_pattern, suggestion, re.DOTALL)

    for sev, exc in zip(severities, excerpts):
        suggestions_list.append({
            "excerpt": exc.strip(),
            "severity": sev
        })

    highlighted_text, _ = highlight_severity(doc["text"], suggestions_list)

    return (
        f"{doc.get('title','Untitled')} - {doc.get('author','Unknown')}",
        category_str,
        suggestion,
        category_str,
        f"Document {current_index+1} of {len(docs_to_review)}",
        highlighted_text
    )

def save_current_document(label, suggestion, current_doc_idx, full_texts):
    """保存当前文档的评审结果"""
    if not full_texts or current_doc_idx >= len(full_texts):
        return "No document to save"

    doc_info = {
        "document_index": current_doc_idx,
        "label": label,
        "suggestion": suggestion,
        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    }

    global saved_data
    saved_data.append(doc_info)
    return f"Document {current_doc_idx + 1} saved!"

# 导出所有评审结果

def export_report(format_type):
    """导出所有保存的评审结果"""
    if not saved_data:
        raise gr.Error("No data to export. Please save some documents first.")

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"DMP_Report_{timestamp}"

    if format_type == "CSV":
        df = pd.DataFrame(saved_data)
        output = df.to_csv(index=False)
        return gr.File(value=output, visible=True, filename=f"{filename}.csv")
    elif format_type == "JSON":
        output = json.dumps(saved_data, indent=2)
        return gr.File(value=output, visible=True, filename=f"{filename}.json")
    elif format_type == "Excel":
        df = pd.DataFrame(saved_data)
        output = df.to_excel(f"{filename}.xlsx", index=False)
        return gr.File(value=f"{filename}.xlsx", visible=True, filename=f"{filename}.xlsx")
    else:
        raise gr.Error("Unsupported export format")

pdf_files = gr.File(file_types=[".pdf"], file_count="multiple")

# ===== Gradio界面搭建 =====

with gr.Blocks(title="DMP Document Review System", theme="soft") as demo:
    gr.Markdown("## DMP Document Classification and Recommendation System")
    gr.Markdown("Supports batch PDF upload, auto classification, manual review, and exporting results.")

    gr.Markdown("""
    <style>
    /* 让左右两栏对齐 */
    .gr-box .gr-column {
        display: flex;
        flex-direction: column;
        height: 100%;
    }

    /* 左侧高亮显示区域高度拉满 */
    .scrollable-html-box {
        height: 640px !important;
        overflow-y: auto;
        padding: 10px;
        border: 1px solid #ccc;
        font-size: 14px;
        line-height: 1.6;
    }

    /* 右侧建议输入框高度 */
    .scrollable-textbox textarea {
        height: 640px !important;
        overflow-y: auto;
        font-size: 14px;
        line-height: 1.6;
    }
    </style>
    """)

    current_doc_index = gr.State(0)
    current_page_num = gr.State(1)
    full_texts = gr.State([])
    titles_list = gr.State([])

    with gr.Row():
        pdf_input = gr.File(
            label="Upload PDF Files",
            file_types=[".pdf"],
            file_count="multiple",
            height="50px"
        )

    with gr.Row():
        with gr.Column(scale=2):
            highlighted_display = gr.HTML(
                label="Document Content with Highlighted Issues",
                value="<div>Waiting...</div>",
                elem_classes=["scrollable-html-box"]
            )
            with gr.Row():
                prev_page_btn = gr.Button("← Previous Page", variant="secondary", min_width=120)
                page_info = gr.Textbox(
                    value="Page 1/1",
                    interactive=False,
                    show_label=False,
                    lines=1,
                    elem_classes=["no-border"],
                    min_width=100,
                    scale=1
                )
                next_page_btn = gr.Button("Next Page →", variant="secondary", min_width=120)
            with gr.Row():
                next_doc_btn = gr.Button("Next Document →", variant="primary", min_width=360)

        with gr.Column(scale=1):
            process_btn = gr.Button("Auto Classify and Generate Suggestions", variant="primary")

            with gr.Group():
                title_author_output = gr.Textbox(label="Title & Author", lines=1, interactive=False)
                auto_category = gr.Textbox(label="Model Auto Classification", interactive=False)
                suggestion_box = gr.Textbox(
                    label="Suggestions (editable)",
                    lines=25,  # 增加高度
                    max_lines=35,  # 可选：允许更长的建议
                    interactive=True,
                    elem_classes=["scrollable-textbox"],
                    show_copy_button=True  # 可选：增加复制按钮
                )

                manual_label = gr.Dropdown(
                    choices=["Good", "Not Good", "Intermediate DMP Needed", "Good with Suggestions", "Other"],
                    label="Final Reviewed Label",
                    value="Good",
                    allow_custom_value=True
                )
            with gr.Row():
                save_btn = gr.Button("Save Current", variant="primary")
                export_btn = gr.Button("Download Report", variant="secondary")
                export_file = gr.File(label="Exported File", visible=False)

    def get_current_page_text(full_text, page_num):
        lines = full_text.split('\n')
        lines_per_page = 50
        max_page = max(1, (len(lines) + lines_per_page - 1) // lines_per_page)
        page_num = max(1, min(page_num, max_page))
        start = (page_num - 1) * lines_per_page
        end = min(start + lines_per_page, len(lines))
        return "\n".join(lines[start:end]), page_num, max_page

    def start_review(pdf_files):
        if not pdf_files:
            return [gr.update()]*6 + [[], [], 0, 1, "No files"]

        import fitz
        docs = []
        titles = []
        for pdf_file in pdf_files:
            doc = fitz.open(pdf_file.name)
            text = "\n".join(page.get_text() for page in doc)
            metadata = doc.metadata
            title = metadata.get("title", "").strip()
            author = metadata.get("author", "").strip()

            # 尝试从第一页提取标题
            if not title:
                first_page = doc[0].get_text("text")
                lines = [line.strip() for line in first_page.split('\n') if line.strip()]
                if lines:
                    title = lines[0]  # 第一行为标题
            if not author:
                if len(lines) > 1:
                    author = lines[1]  # 第二行为作者

            title_author = f"{title} - {author}" if author else title

            docs.append(text)
            titles.append(title)

        category_num = classify_text(docs[0])
        category_str = num_to_category.get(category_num, "Unknown")
        suggestion = generate_suggestion(docs[0], category_num)

        page_text, page_num, max_page = get_current_page_text(docs[0], 1)
        page_info_str = f"Document 1/{len(docs)} | Page {page_num}/{max_page}"

        return [
            title_author,
            category_str,
            suggestion,
            category_str,
            page_info_str,
            page_text,
            docs,
            titles,
            0,
            1
        ]

    def navigate_page(direction, current_doc_idx, current_page, full_texts):
        if not full_texts:
            return current_page, gr.update(), gr.update()
        page_text, new_page, max_page = get_current_page_text(full_texts[current_doc_idx], current_page + direction)
        page_info = f"Document {current_doc_idx+1}/{len(full_texts)} | Page {new_page}/{max_page}"
        return new_page, page_info, page_text

    def next_document(current_doc_idx, current_page, full_texts, titles):
        if not full_texts or current_doc_idx >= len(full_texts) - 1:
            return [gr.update()] * 9

        new_idx = current_doc_idx + 1
        category_num = classify_text(full_texts[new_idx])
        category_str = num_to_category.get(category_num, "Unknown")
        suggestion = generate_suggestion(full_texts[new_idx], category_num)

        page_text, page_num, max_page = get_current_page_text(full_texts[new_idx], 1)
        page_info = f"Document {new_idx+1}/{len(full_texts)} | Page {page_num}/{max_page}"

        return [
            f"{titles[new_idx]}",
            category_str,
            suggestion,
            category_str,
            page_info,
            page_text,
            new_idx,
            1
        ]

    def save_current_document(manual_label, suggestion, current_doc_idx, full_texts):
        if not full_texts:
            return gr.Info("Nothing to save")
        doc_info = {
            "index": current_doc_idx,
            "label": manual_label,
            "suggestion": suggestion,
            "text": full_texts[current_doc_idx]
        }
        saved_data.append(doc_info)
        return gr.Info("Saved successfully!")

    def export_report():
        import pandas as pd, tempfile
        from datetime import datetime
        if not saved_data:
            raise gr.Error("No data to export")

        df = pd.DataFrame(saved_data)
        now = datetime.now().strftime("%Y%m%d_%H%M%S")
        with tempfile.NamedTemporaryFile(delete=False, suffix=".xlsx") as tmp:
            df.to_excel(tmp.name, index=False)
            return tmp.name

    process_btn.click(
        fn=start_review,
        inputs=[pdf_input],
        outputs=[
            title_author_output, auto_category, suggestion_box, manual_label,
            page_info, highlighted_display,
            full_texts, titles_list, current_doc_index, current_page_num
        ]
    )

    prev_page_btn.click(
        fn=navigate_page,
        inputs=[gr.State(-1), current_doc_index, current_page_num, full_texts],
        outputs=[current_page_num, page_info, highlighted_display]
    )

    next_page_btn.click(
        fn=navigate_page,
        inputs=[gr.State(1), current_doc_index, current_page_num, full_texts],
        outputs=[current_page_num, page_info, highlighted_display]
    )

    next_doc_btn.click(
        fn=next_document,
        inputs=[current_doc_index, current_page_num, full_texts, titles_list],
        outputs=[title_author_output, auto_category, suggestion_box, manual_label, page_info, highlighted_display, current_doc_index, current_page_num]
    )

    save_btn.click(
        fn=save_current_document,
        inputs=[manual_label, suggestion_box, current_doc_index, full_texts],
        outputs=[]
    )

    # 替换原来 export_btn.click(...) 这一段
    export_btn.click(
        fn=lambda: export_report(export_format.value),
        inputs=[],
        outputs=[export_file]
    )

# ===== 启动 =====
demo.launch(debug=True)

pip install pandas openpyxl