# app.py
import streamlit as st
import pandas as pd
from typing import List, Dict
import plotly.express as px
import re


class LiteratureSearchApp:
    """文献搜索Web应用"""

    def __init__(self, corpus):
        self.corpus = corpus

    def run(self):
        st.set_page_config(
            page_title="文献语料库检索系统",
            page_icon="📚",
            layout="wide"
        )

        st.title("📚 个人文献语料库检索系统")
        st.markdown("---")

        # 侧边栏
        with st.sidebar:
            st.header("搜索设置")

            # 搜索查询
            query = st.text_input(
                "搜索关键词",
                placeholder="例如: lutein, antioxidant, clinical trial"
            )

            # 章节筛选
            section_options = {
                "所有章节": None,
                "摘要": "abstract",
                "引言": "introduction",
                "方法": "methods",
                "结果": "results",
                "讨论": "discussion",
                "结论": "conclusion"
            }

            selected_section = st.selectbox(
                "限定章节",
                options=list(section_options.keys())
            )

            # 结果数量
            n_results = st.slider("结果数量", 5, 50, 10)

            # 相关性阈值
            min_relevance = st.slider("最小相关性", 0.0, 1.0, 0.3)

            # 执行搜索按钮
            search_button = st.button("搜索", type="primary")

        # 主界面
        col1, col2 = st.columns([2, 1])

        with col1:
            if search_button and query:
                st.subheader(f"搜索结果: '{query}'")

                # 执行搜索
                section_type = section_options[selected_section]
                results = self.corpus.search_by_section(
                    query,
                    section_type=section_type,
                    n_results=n_results
                )

                # 过滤相关性
                filtered_results = [
                    r for r in results
                    if r['relevance_score'] >= min_relevance
                ]

                if not filtered_results:
                    st.warning("未找到相关结果")
                else:
                    # 显示结果
                    for i, result in enumerate(filtered_results):
                        with st.expander(
                                f"结果 {i + 1}: {result['metadata']['article_title']} "
                                f"(相关性: {result['relevance_score']:.2f})"
                        ):
                            self._display_result(result)

            elif not search_button:
                st.info("在侧边栏输入搜索词开始检索")

        with col2:
            st.subheader("语料库统计")
            stats = self.corpus.get_statistics()

            st.metric("总段落数", stats['total_segments'])
            st.metric("文献数量", stats['unique_articles'])

            # 章节分布图
            if stats['section_distribution']:
                df_sections = pd.DataFrame(
                    list(stats['section_distribution'].items()),
                    columns=['章节', '段落数']
                )
                fig = px.pie(
                    df_sections,
                    values='段落数',
                    names='章节',
                    title="章节分布"
                )
                st.plotly_chart(fig, use_container_width=True)

    def _display_result(self, result: Dict):
        """显示单个结果"""
        metadata = result['metadata']
        text = result['document']

        # 基本信息
        col1, col2 = st.columns(2)
        with col1:
            st.markdown(f"**文献标题:** {metadata['article_title']}")
            st.markdown(f"**作者:** {metadata.get('authors', 'N/A')}")
            st.markdown(f"**期刊:** {metadata.get('journal', 'N/A')}")

        with col2:
            st.markdown(f"**章节:** {metadata['section_title']}")
            st.markdown(f"**页码:** {metadata['page_number']}")
            st.markdown(f"**相关性分数:** {result['relevance_score']:.3f}")

        st.markdown("---")

        # 高亮显示搜索词
        query_words = st.session_state.get('query', '').lower().split()
        highlighted_text = text

        for word in query_words:
            if len(word) > 2:  # 只高亮长度>2的词
                pattern = re.compile(re.escape(word), re.IGNORECASE)
                highlighted_text = pattern.sub(
                    f"**{word.upper()}**",
                    highlighted_text
                )

        # 显示文本（限制长度）
        st.markdown("### 相关文本")
        text_preview = highlighted_text[:2000] + "..." if len(highlighted_text) > 2000 else highlighted_text
        st.markdown(text_preview)

        # 显示完整文本的选项
        if len(highlighted_text) > 2000:
            with st.expander("显示完整文本"):
                st.markdown(highlighted_text)

        st.markdown("---")


class EnhancedLiteratureSearchApp(LiteratureSearchApp):
    """增强版搜索应用"""

    def __init__(self, corpus):
        super().__init__(corpus)
        self.enhanced_mode = hasattr(corpus, 'use_llm') and corpus.use_llm

    def run(self):
        # 原有代码...

        # 在侧边栏添加高级选项
        with st.sidebar:
            st.markdown("---")
            st.subheader("高级选项")

            if self.enhanced_mode:
                generate_summaries = st.checkbox("生成AI摘要", value=True)
                extract_key_points = st.checkbox("提取关键点", value=False)
                show_ai_insights = st.checkbox("显示AI分析", value=True)

                # 问答功能
                st.markdown("---")
                st.subheader("文献问答")
                article_id_for_qa = st.text_input("文献ID", placeholder="输入文献ID")
                question = st.text_input("问题", placeholder="输入关于文献的问题")
                if st.button("提问") and article_id_for_qa and question:
                    answer = self.corpus.ask_question_about_article(
                        article_id_for_qa,
                        question
                    )
                    st.success(f"回答: {answer}")

        # 在主界面中显示AI生成的内容
        if search_button and query:
            # 使用增强搜索
            if self.enhanced_mode:
                results = self.corpus.search_with_summary(
                    query,
                    n_results=n_results,
                    generate_summary=generate_summaries
                )
            else:
                results = self.corpus.search(query, n_results=n_results)

            # 显示结果时添加AI内容
            for i, result in enumerate(filtered_results):
                with st.expander(...):
                    self._display_enhanced_result(result, extract_key_points, show_ai_insights)

    def _display_enhanced_result(self, result: Dict, show_key_points=False, show_ai_insights=True):
        """显示增强结果"""
        # 原有显示代码...

        # 添加AI生成的内容
        if show_ai_insights and 'summary' in result:
            st.markdown("### 🤖 AI摘要")
            st.info(result['summary'])

            if show_key_points and 'key_points' in result:
                st.markdown("#### 🔑 关键点")
                for j, point in enumerate(result['key_points'], 1):
                    st.markdown(f"{j}. {point}")
