# main.py
import os
import json
import shutil
from pathlib import Path
import argparse
import sys
from llm_enhancer import LLMEnhancer
from corpus_builder import LiteratureCorpus


class EnhancedLiteratureCorpus(LiteratureCorpus):
    """增强版语料库，集成LLM功能"""

    def __init__(self, persist_directory="./corpus_db", use_llm=False):
        super().__init__(persist_directory)

        self.use_llm = use_llm
        if use_llm:
            self.llm = LLMEnhancer()

    def search_with_summary(self, query: str, n_results: int = 10, generate_summary=True):
        """搜索并生成摘要"""
        results = self.search(query, n_results=n_results)

        if self.use_llm and generate_summary:
            for result in results:
                # 为每个结果生成摘要
                text = result['document']
                summary = self.llm.generate_summary(text)
                result['summary'] = summary

                # 提取关键点
                key_points = self.llm.extract_key_points(text)
                result['key_points'] = key_points

        return results

    def ask_question_about_article(self, article_id: str, question: str):
        """对特定文献提问"""
        # 获取文献的所有段落
        all_docs = self.collection.get(where={"article_id": article_id})

        if not all_docs['documents']:
            return "未找到该文献"

        # 合并相关文本
        context = "\n".join(all_docs['documents'][:5])  # 只取前5段

        if self.use_llm and self.llm.qa_model:
            answer = self.llm.answer_question(context, question)
            return answer
        else:
            return "LLM问答功能未启用"

def main():
    parser = argparse.ArgumentParser(description="构建文献语料库")
    parser.add_argument('--export', action='store_true', help='从Zotero导出数据')
    parser.add_argument('--process', action='store_true', help='处理PDF文件')
    parser.add_argument('--build', action='store_true', help='构建向量数据库')
    parser.add_argument('--run', action='store_true', help='启动Web界面')
    parser.add_argument('--pdf-dir', type=str, default='./pdfs', help='PDF目录')

    args = parser.parse_args()

    if args.export:
        print("从Zotero导出数据...")
        from export_zotero_data import export_zotero_highlights
        articles = export_zotero_highlights()

        # 复制PDF文件到本地目录
        if articles:
            os.makedirs(args.pdf_dir, exist_ok=True)
            for article in articles:
                pdf_path = article['pdfPath']
                if os.path.exists(pdf_path):
                    # 简化路径处理
                    try:
                        dest_path = os.path.join(
                            args.pdf_dir,
                            f"{article['itemID']}.pdf"
                        )
                        shutil.copy2(pdf_path, dest_path)
                        article['local_pdf_path'] = dest_path
                    except:
                        print(f"复制失败: {pdf_path}")

    if args.process:
        print("处理PDF文件...")
        from pdf_processor import PDFSectionExtractor

        extractor = PDFSectionExtractor()
        processed_articles = []

        # 加载导出的数据
        if os.path.exists('zotero_starred_articles.json'):
            with open('zotero_starred_articles.json', 'r', encoding='utf-8') as f:
                articles = json.load(f)

            for article in articles:
                pdf_path = article.get('local_pdf_path')
                if pdf_path and os.path.exists(pdf_path):
                    result = extractor.process_pdf(pdf_path)
                    if result:
                        result['metadata'] = article['metadata']
                        result['zotero_id'] = article['itemID']

                        # 保存处理结果
                        output_file = f"processed_{article['itemID']}.json"
                        with open(output_file, 'w', encoding='utf-8') as f:
                            json.dump(result, f, ensure_ascii=False, indent=2)

                        processed_articles.append(result)
                        print(f"已处理: {pdf_path}")

        # 保存所有处理结果索引
        with open('processed_articles_index.json', 'w', encoding='utf-8') as f:
            json.dump(
                [{'zotero_id': a['zotero_id'], 'pdf_path': a['pdf_path']}
                 for a in processed_articles],
                f, indent=2
            )

    if args.build:
        print("构建向量数据库...")
        from corpus_builder import LiteratureCorpus

        corpus = LiteratureCorpus()

        # 加载处理过的文章
        if os.path.exists('processed_articles_index.json'):
            with open('processed_articles_index.json', 'r') as f:
                articles_index = json.load(f)

            for article_info in articles_index:
                zotero_id = article_info['zotero_id']
                processed_file = f"processed_{zotero_id}.json"

                if os.path.exists(processed_file):
                    with open(processed_file, 'r', encoding='utf-8') as f:
                        article_data = json.load(f)

                    # 添加到语料库
                    corpus.add_document(
                        article_id=zotero_id,
                        metadata=article_data['metadata'],
                        segments=article_data['segments']
                    )

                    print(f"已添加到语料库: {article_data['metadata'].get('title', zotero_id)}")

        # 保存数据库
        corpus.client.persist()
        print("语料库构建完成！")

    if args.run:
        print("启动Web界面...")
        from corpus_builder import LiteratureCorpus
        from app import LiteratureSearchApp

        # 加载现有语料库
        corpus = LiteratureCorpus()

        # 启动应用
        app = LiteratureSearchApp(corpus)
        app.run()


if __name__ == "__main__":
    """
# 1. 从Zotero导出数据
python main.py --export

# 2. 处理PDF文件
python main.py --process

# 3. 构建向量数据库
python main.py --build

# 4. 启动Web界面
streamlit run app.py

# 使用LLM增强功能
python main_enhanced.py --llm

# 批量处理
python main_enhanced.py --batch

# 聚类分析并可视化
python main_enhanced.py --cluster 5 --visualize

# 启动增强版Web应用
streamlit run app.py -- --llm
"""
    main()