import os
import gradio as gr
import logging
from typing import List, Dict, Any

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 从环境变量中获取 API Key
def get_api_key() -> str:
    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        logger.error("环境变量中未找到 OPENAI_API_KEY")
        raise ValueError("环境变量中未找到 OPENAI_API_KEY，请先设置。")
    return api_key

# 初始化 LLM
def initialize_llm(api_key: str):
    try:
        from langchain_openai.chat_models import ChatOpenAI
        logger.info("初始化 LLM...")
        return ChatOpenAI(
            temperature=0,
            model="gpt-4o-mini",
            api_key=api_key
        )
    except Exception as e:
        logger.error(f"初始化 LLM 失败: {str(e)}")
        raise

# 定义 Prompt 模板
def initialize_prompt_template():
    from langchain_core.prompts import PromptTemplate
    template_text = """请根据以下 context 回答问题，答案请使用 Markdown 格式输出。
如果 context 存在latex表达式，请正确书写。
如果 context 中包含图表链接，请在回答中原封不动地加入图表，并在每个包含"/images"的链接前添加前缀 "https://huggingface.co/spaces/zliang/palynogeology/resolve/main"。
Question: {question}
Context: {context}"""
    return PromptTemplate(
        input_variables=["question", "context"],
        template=template_text
    )


# 将检索到的文档内容格式化为字符串
def format_docs(docs) -> str:
    return "\n\n".join(doc.page_content for doc in docs)

# 加载本地 FAISS 向量库
def initialize_vectorstore(api_key: str):
    try:
        from langchain_openai import OpenAIEmbeddings
        from langchain_community.vectorstores import FAISS

        logger.info("初始化向量存储...")
        embed = OpenAIEmbeddings(
            model="text-embedding-3-small",
            api_key=api_key
        )
        
        if not os.path.exists("faiss_db"):
            logger.error("未找到 'faiss_db' 目录")
            raise FileNotFoundError("未找到 'faiss_db' 目录，请先构建向量库。")
            
        return FAISS.load_local("faiss_db", embed, allow_dangerous_deserialization=True)
    except Exception as e:
        logger.error(f"初始化向量存储失败: {str(e)}")
        raise

# 构建检索器
def initialize_retriever(db):
    return db.as_retriever(
        search_type="mmr", 
        search_kwargs={"score_threshold": 0.5, "k": 3}
    )

# 构造问答链
def initialize_qa_chain(llm, prompt, retriever):
    from langchain_core.runnables import RunnablePassthrough
    from langchain_core.output_parsers import StrOutputParser

    qa_chain = (
        {
            "context": retriever | format_docs,
            "question": RunnablePassthrough(),
        }
        | prompt
        | llm
        | StrOutputParser()
    )
    return qa_chain

# 根据用户输入的问题调用问答链
def answer_question(qa_chain, question: str) -> str:
    try:
        logger.info(f"处理问题: {question}")
        answer = qa_chain.invoke(question)
        return answer
    except Exception as e:
        logger.error(f"调用问答链时出错: {str(e)}", exc_info=True)
        return "⚠️ 出错了，请稍后重试。如果问题持续存在，请联系管理员。"

# 返回《孢粉地质学》书籍简介
def get_book_introduction() -> str:
    introduction = """
# 《孢粉地质学》

![](https://huggingface.co/spaces/zliang/palynogeology/resolve/main/images/cover.jpg)

## 简介
《孢粉地质学》是一本系统介绍孢粉与孢子在地质学中应用的重要专著。本书由国内顶尖专家编撰，融合了最新的研究成果和实践经验。

## 内容亮点
- **孢粉与孢子的形成与保存**  
  阐述孢粉在沉积环境中的生成过程及保存机制。
- **鉴定与分类方法**  
  详细介绍如何通过形态学、化学特征对孢粉进行鉴定与分类。
- **地层对比与古环境重建**  
  探讨孢粉在地层划分、沉积环境重建、古气候研究等方面的应用。
- **案例分析与实践应用**  
  配合丰富的实例解析，为读者提供理论与实践相结合的指导。

## 适读人群
本书适合地质学、古生物学及相关领域的研究人员和学生参考，不仅为学术研究提供了坚实的理论基础，同时也为野外勘查和实际应用提供了实用工具。

👉 欢迎在"书籍问答"标签页中提问，获取更多关于书中内容的详细解读！
"""
    return introduction

# 定义常见问题列表
def get_faq_list() -> List[Dict[str, str]]:
    return [
        {"question": "孢粉是什么？", "category": "基础概念"},
        {"question": "如何采集孢粉样本？", "category": "实验方法"},
        {"question": "孢粉分析的主要步骤有哪些？", "category": "实验方法"},
        {"question": "如何鉴定孢粉的年代？", "category": "年代学"},
        {"question": "孢粉数据如何应用于地层对比？", "category": "地层学"},
        {"question": "常见的孢粉类型有哪些？", "category": "分类学"},
        {"question": "孢粉如何指示古气候变化？", "category": "古环境"},
        {"question": "孢粉研究在石油勘探中的应用", "category": "应用领域"},
        {"question": "孢粉与孢子有什么区别？", "category": "基础概念"},
        {"question": "电子显微镜在孢粉研究中的应用", "category": "技术方法"}
    ]

# 构建自定义 Gradio 界面
def create_custom_ui(qa_chain):
    # 主题配置
    theme = gr.Theme(
        primary_hue="green",
        secondary_hue="emerald",
        neutral_hue="gray",
        font=[gr.themes.GoogleFont("Source Sans Pro"), "system-ui", "sans-serif"],
    )
    
    # 自定义CSS
    custom_css = """
    .container {max-width: 1000px; margin: auto;}
    .header-logo {text-align: center; padding: 20px 0; margin-bottom: 20px;}
    .header-title {font-size: 2.5rem; font-weight: 700; margin: 0.5rem 0;}
    .header-subtitle {font-size: 1.25rem; color: #555; margin-bottom: 1rem;}
    .footer {text-align: center; margin-top: 40px; padding: 20px 0; font-size: 0.9rem; color: #666;}
    .card {border-radius: 10px; box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1); padding: 20px; margin-bottom: 20px; background-color: white;}
    .faq-item {cursor: pointer; padding: 10px; border-radius: 5px; margin: 5px 0; transition: all 0.2s;}
    .faq-item:hover {background-color: #f0f9f0;}
    .category-tag {display: inline-block; font-size: 0.8rem; padding: 2px 8px; border-radius: 12px; background-color: #e0f2e0; color: #2e7d32; margin-left: 10px;}
    .loading-indicator {display: flex; justify-content: center; align-items: center; height: 100px;}
    .sample-questions-header {font-weight: 600; margin: 15px 0 10px 0;}
    .search-box {margin-bottom: 15px;}
    
    /* 自定义滚动条 */
    ::-webkit-scrollbar {width: 8px; height: 8px;}
    ::-webkit-scrollbar-track {background: #f1f1f1; border-radius: 10px;}
    ::-webkit-scrollbar-thumb {background: #c1e0c1; border-radius: 10px;}
    ::-webkit-scrollbar-thumb:hover {background: #8bc48b;}
    
    /* 响应式调整 */
    @media (max-width: 768px) {
        .header-title {font-size: 2rem;}
        .header-subtitle {font-size: 1rem;}
    }
    """
    
    with gr.Blocks(theme=theme, css=custom_css) as demo:
        # 注入 KaTeX 样式表以支持LaTeX公式渲染
        gr.HTML(
            '<link rel="stylesheet" '
            'href="https://cdn.jsdelivr.net/npm/katex@0.16.8/dist/katex.min.css" '
            'integrity="sha384-GvrOXuhMATgEsSwCs4smul74iXGOixntILdUW9XmUC6+HX0sLNAK3q71HotJqlAn" '
            'crossorigin="anonymous">'
        )
        
        # 应用标题和介绍
        with gr.Row(elem_classes="header-logo"):
            gr.HTML("""
            <div style="text-align: center;">
                <h1 class="header-title">《孢粉地质学》数字化知识库</h1>
                <p class="header-subtitle">探索孢粉地质学的奥秘，获取专业知识解答</p>
            </div>
            """)
        
        # 主内容区域
        with gr.Tabs() as tabs:
            # 书籍简介标签页
            with gr.TabItem("📚 书籍简介", id="intro"):
                with gr.Box(elem_classes="card"):
                    gr.Markdown(get_book_introduction())
            
            # 书籍问答标签页
            with gr.TabItem("❓ 知识问答", id="qa"):
                with gr.Row():
                    # 左侧：问答区域
                    with gr.Column(scale=7):
                        with gr.Box(elem_classes="card"):
                            gr.Markdown("### 📝 提问区")
                            with gr.Row():
                                question_input = gr.Textbox(
                                    lines=3, 
                                    placeholder="请输入您关于孢粉地质学的问题...", 
                                    label="问题",
                                    elem_classes="search-box"
                                )
                            
                            with gr.Row():
                                submit_button = gr.Button("提交问题", variant="primary")
                                clear_button = gr.Button("清空", variant="secondary")
                            
                            with gr.Box():
                                with gr.Row():
                                    status_indicator = gr.Markdown("准备就绪，等待提问...")
                        
                        with gr.Box(elem_classes="card", visible=False) as answer_card:
                            gr.Markdown("### 🔍 回答结果")
                            answer_output = gr.Markdown(
                                label="回答", 
                                elem_id="answer-output",
                                latex_delimiters=[
                                    {"left": "$", "right": "$", "display": False},
                                    {"left": "$$", "right": "$$", "display": True}
                                ]
                            )
                    
                    # 右侧：常见问题
                    with gr.Column(scale=3):
                        with gr.Box(elem_classes="card"):
                            gr.Markdown("### 📋 常见问题")
                            
                            # 分类筛选下拉框
                            category_filter = gr.Dropdown(
                                ["全部", "基础概念", "实验方法", "分类学", "年代学", "地层学", "古环境", "应用领域", "技术方法"],
                                value="全部",
                                label="按类别筛选"
                            )
                            
                            faq_container = gr.HTML()  # 用于显示FAQ的容器
                            
                            # 更新FAQ显示的函数
                            def update_faq_display(category):
                                faq_list = get_faq_list()
                                html = "<div class='faq-list'>"
                                
                                for item in faq_list:
                                    if category == "全部" or item["category"] == category:
                                        html += f"""
                                        <div class='faq-item' onclick='document.querySelector("[data-testid=textbox]").value = this.getAttribute("data-question"); document.querySelector("[data-testid=button]").click()' data-question='{item["question"]}'>
                                            {item["question"]}
                                            <span class='category-tag'>{item["category"]}</span>
                                        </div>
                                        """
                                
                                html += "</div>"
                                return html
                            
                            # 绑定更新事件
                            category_filter.change(update_faq_display, inputs=[category_filter], outputs=[faq_container])
            
            # 使用指南标签页
            with gr.TabItem("📖 使用指南", id="guide"):
                with gr.Box(elem_classes="card"):
                    gr.Markdown("""
                    # 使用指南
                    
                    ## 🔍 如何有效提问
                    
                    为了获得最准确的回答，建议您：
                    
                    1. **使用专业术语** - 尽量使用孢粉学和地质学的专业术语
                    2. **具体明确** - 问题越具体，回答越精准
                    3. **一次一问** - 每次提交一个问题，而不是多个问题组合
                    4. **参考示例** - 可以参考右侧的常见问题示例
                    
                    ## 📊 功能介绍
                    
                    本平台提供以下功能：
                    
                    - **书籍内容检索** - 快速获取《孢粉地质学》中的知识点
                    - **专业问题解答** - 解答孢粉学相关的各类专业问题
                    - **图例与公式** - 支持显示专业图例和数学公式
                    - **常见问题库** - 提供常见问题的快速访问
                    
                    ## ⚠️ 注意事项
                    
                    - 本平台不替代专业人士的建议
                    - 复杂图表可能需要等待较长时间加载
                    - 若遇到技术问题，请刷新页面或稍后再试
                    """)
        
        # 页脚
        with gr.Row(elem_classes="footer"):
            gr.HTML("""
            <div>
                <p>© 2025 孢粉地质学数字平台 | 由 GPT-4o 提供支持</p>
                <p>如有问题或建议，请联系我们</p>
            </div>
            """)
        
        # 函数：处理问题提交
        def process_question(question):
            if not question or question.strip() == "":
                return ("请输入有效的问题！", gr.update(visible=False))
            
            status = "🔍 正在检索相关知识..."
            yield (status, gr.update(visible=False))
            
            try:
                # 调用问答链获取答案
                answer = answer_question(qa_chain, question)
                status = "✅ 回答已生成"
                
                # 显示答案卡片
                return (status, gr.update(visible=True, value=answer))
            except Exception as e:
                logger.error(f"处理问题时出错: {str(e)}")
                status = "❌ 出错了：" + str(e)
                return (status, gr.update(visible=False))
        
        # 函数：清空输入和结果
        def clear_input():
            return "", "准备就绪，等待提问...", gr.update(visible=False)
        
        # 事件绑定
        submit_button.click(
            process_question, 
            inputs=[question_input], 
            outputs=[status_indicator, answer_card]
        )
        
        clear_button.click(
            clear_input, 
            inputs=[], 
            outputs=[question_input, status_indicator, answer_card]
        )
        
        # 初始化FAQ显示
        demo.load(
            update_faq_display, 
            inputs=[category_filter], 
            outputs=[faq_container]
        )
    
    return demo

def main():
    try:
        logger.info("启动应用...")
        api_key = get_api_key()
        llm = initialize_llm(api_key)
        prompt = initialize_prompt_template()
        db = initialize_vectorstore(api_key)
        retriever = initialize_retriever(db)
        qa_chain = initialize_qa_chain(llm, prompt, retriever)
        app = create_custom_ui(qa_chain)
        app.launch(share=False)
    except Exception as e:
        logger.critical(f"应用启动失败: {str(e)}", exc_info=True)
        print(f"错误: {str(e)}")
        print("请检查日志以获取详细信息。")

if __name__ == "__main__":
    main()