import streamlit as st
import pandas as pd
from typing import List, Dict, Any
import math


class DataPaginator:
    """数据分页器组件
    
    提供优雅的分页界面和导航功能
    """
    
    def __init__(self, key_prefix: str):
        """初始化分页器
        
        Args:
            key_prefix: streamlit组件的key前缀，确保唯一性
        """
        self.key_prefix = key_prefix
    
    def paginate_data(self, data: pd.Series, page_size: int = 20) -> Dict[str, Any]:
        """对数据进行分页处理
        
        Args:
            data: 要分页的数据Series
            page_size: 每页显示的数量
            
        Returns:
            包含分页信息的字典
        """
        if data.empty:
            return {
                'current_page_data': pd.Series(dtype='object'),
                'total_pages': 0,
                'current_page': 1,
                'total_items': 0,
                'start_idx': 0,
                'end_idx': 0
            }
        
        # 获取value_counts
        value_counts = data.value_counts()
        total_items = len(value_counts)
        total_pages = math.ceil(total_items / page_size)
        
        # 获取当前页码
        current_page = st.session_state.get(f"{self.key_prefix}_current_page", 1)
        current_page = max(1, min(current_page, total_pages))  # 确保页码在有效范围内
        
        # 计算索引范围
        start_idx = (current_page - 1) * page_size
        end_idx = min(start_idx + page_size, total_items)
        
        # 获取当前页数据
        current_page_data = value_counts.iloc[start_idx:end_idx]
        
        return {
            'current_page_data': current_page_data,
            'total_pages': total_pages,
            'current_page': current_page,
            'total_items': total_items,
            'start_idx': start_idx,
            'end_idx': end_idx,
            'original_data': data
        }
    
    def render_pagination_controls(self, pagination_info: Dict[str, Any]) -> Dict[str, Any]:
        """渲染分页控制器
        
        Args:
            pagination_info: 分页信息字典
            
        Returns:
            包含可能更新的分页设置的字典
        """
        if pagination_info['total_pages'] <= 1:
            return {'page_size_changed': False}
        
        current_page = pagination_info['current_page']
        total_pages = pagination_info['total_pages']
        
        # 分页控制布局：首页/上页 | 页码/分页数量 | 下页/末页
        col1, col2, col3 = st.columns([2, 3, 2])
        
        page_changed = False
        page_size_changed = False
        
        with col1:
            # 首页和上一页按钮
            nav_col1, nav_col2 = st.columns(2)
            with nav_col1:
                if st.button("⏮️", key=f"{self.key_prefix}_first_page", 
                           disabled=current_page == 1, help="首页"):
                    st.session_state[f"{self.key_prefix}_current_page"] = 1
                    page_changed = True
            
            with nav_col2:
                if st.button("◀️", key=f"{self.key_prefix}_prev_page", 
                           disabled=current_page == 1, help="上一页"):
                    st.session_state[f"{self.key_prefix}_current_page"] = current_page - 1
                    page_changed = True
        
        with col2:
            # 页码输入和分页数量选择
            page_col1, page_col2 = st.columns([3, 2])
            
            with page_col1:
                new_page = st.number_input(
                    f"页码 (共{total_pages}页)",
                    min_value=1,
                    max_value=total_pages,
                    value=current_page,
                    key=f"{self.key_prefix}_page_input"
                )
                if new_page != current_page:
                    st.session_state[f"{self.key_prefix}_current_page"] = new_page
                    page_changed = True
            
            with page_col2:
                # 获取当前分页大小
                current_page_size = st.session_state.get(f"{self.key_prefix}_page_size", 20)
                new_page_size = st.selectbox(
                    "每页",
                    options=[10, 20, 50, 100],
                    index=[10, 20, 50, 100].index(current_page_size) if current_page_size in [10, 20, 50, 100] else 1,
                    key=f"{self.key_prefix}_page_size_select"
                )
                
                if new_page_size != current_page_size:
                    st.session_state[f"{self.key_prefix}_page_size"] = new_page_size
                    # 重置到第一页，因为分页大小变了
                    st.session_state[f"{self.key_prefix}_current_page"] = 1
                    page_size_changed = True
        
        with col3:
            # 下一页和末页按钮
            nav_col3, nav_col4 = st.columns(2)
            with nav_col3:
                if st.button("▶️", key=f"{self.key_prefix}_next_page", 
                           disabled=current_page == total_pages, help="下一页"):
                    st.session_state[f"{self.key_prefix}_current_page"] = current_page + 1
                    page_changed = True
            
            with nav_col4:
                if st.button("⏭️", key=f"{self.key_prefix}_last_page", 
                           disabled=current_page == total_pages, help="末页"):
                    st.session_state[f"{self.key_prefix}_current_page"] = total_pages
                    page_changed = True
        
        # 如果页面或分页大小发生变化，需要重新运行
        if page_changed or page_size_changed:
            st.rerun()
        
        return {
            'page_size_changed': page_size_changed,
            'new_page_size': st.session_state.get(f"{self.key_prefix}_page_size", 20)
        }
    
    def render_page_info(self, pagination_info: Dict[str, Any]) -> None:
        """渲染页面信息
        
        Args:
            pagination_info: 分页信息字典
        """
        total_items = pagination_info['total_items']
        start_idx = pagination_info['start_idx'] + 1  # 从1开始计数
        end_idx = pagination_info['end_idx']
        
        st.info(f"📊 显示第 {start_idx}-{end_idx} 项，共 {total_items} 项")


class InteractiveDataSelector:
    """交互式数据选择器
    
    使用st.data_editor提供优雅的数据选择和统计显示功能
    """
    
    def __init__(self, key_prefix: str):
        """初始化选择器
        
        Args:
            key_prefix: streamlit组件的key前缀，确保唯一性
        """
        self.key_prefix = key_prefix
        self.paginator = DataPaginator(key_prefix)
    
    def render_data_selector(self, 
                            data_series: pd.Series, 
                            title: str = "数据选择与统计",
                            page_size: int = 20) -> List[str]:
        """渲染数据选择器界面
        
        Args:
            data_series: 要展示的数据Series
            title: 标题
            page_size: 默认每页显示的数量（可被用户选择覆盖）
            
        Returns:
            用户选中的值列表
        """
        if data_series.empty:
            st.warning("没有数据可供选择")
            return []
        
        st.write(f"**{title}**")
        
        # 获取用户选择的分页大小，如果没有则使用默认值
        actual_page_size = st.session_state.get(f"{self.key_prefix}_page_size", page_size)
        
        # 获取分页信息
        pagination_info = self.paginator.paginate_data(data_series, actual_page_size)
        
        if pagination_info['total_items'] == 0:
            st.warning("没有数据可供分析")
            return []
        
        # 渲染页面信息
        self.paginator.render_page_info(pagination_info)
        
        # 渲染分页控制器并获取更新信息
        pagination_result = self.paginator.render_pagination_controls(pagination_info)
        
        # 如果分页大小发生变化，重新获取分页信息
        if pagination_result['page_size_changed']:
            actual_page_size = pagination_result['new_page_size']
            pagination_info = self.paginator.paginate_data(data_series, actual_page_size)
        
        # 准备当前页的数据用于data_editor
        current_page_data = pagination_info['current_page_data']
        total_count = len(data_series)
        
        # 构建DataFrame用于data_editor
        editor_data = []
        for value, count in current_page_data.items():
            percentage = (count / total_count * 100)
            editor_data.append({
                '选择': False,  # 默认不选中
                '值': str(value)[:80] + ('...' if len(str(value)) > 80 else ''),  # 截断长值
                '数量': count,
                '占比(%)': round(percentage, 2),
                '_original_value': str(value)  # 保存原始值用于返回
            })
        
        if not editor_data:
            st.warning("当前页没有数据")
            return []
        
        editor_df = pd.DataFrame(editor_data)
        
        # 使用data_editor显示和编辑数据
        st.write(f"**第 {pagination_info['current_page']} 页数据（可选择多行）:**")
        
        edited_df = st.data_editor(
            editor_df[['选择', '值', '数量', '占比(%)']],  # 不显示原始值列
            key=f"_data_editor_{self.key_prefix}__{pagination_info['current_page']}_{actual_page_size}",
            hide_index=True,
            width='stretch',
            column_config={
                "选择": st.column_config.CheckboxColumn(
                    "选择",
                    help="选择此行数据用于筛选",
                    default=False
                ),
                "值": st.column_config.TextColumn(
                    "值",
                    help="数据值",
                    disabled=True
                ),
                "数量": st.column_config.NumberColumn(
                    "数量",
                    help="出现次数",
                    disabled=True
                ),
                "占比(%)": st.column_config.NumberColumn(
                    "占比(%)",
                    help="在总数据中的占比",
                    disabled=True,
                    format="%.2f%%"
                )
            },
            height=min(400, (len(editor_df) + 1) * 35 + 50)  # 动态高度
        )
        
        # 获取选中的值
        selected_indices = edited_df[edited_df['选择'] == True].index.tolist()
        selected_values = []
        
        for idx in selected_indices:
            if idx < len(editor_data):  # 确保索引有效
                original_value = editor_data[idx]['_original_value']
                selected_values.append(original_value)
        
        # 显示选择统计
        if selected_values:
            selected_count = sum(
                current_page_data[v] for v in selected_values 
                if v in current_page_data.index
            )
            selected_percentage = (selected_count / total_count * 100)
            
            st.success(f"✅ 当前页已选择 {len(selected_values)} 个值，"
                      f"覆盖 {selected_count} 个数据点 ({selected_percentage:.1f}%)")
        
        return selected_values
    
    def render_selection_summary(self, selected_values: List[str], data_series: pd.Series) -> None:
        """渲染选择总结
        
        Args:
            selected_values: 选中的值列表
            data_series: 原始数据Series
        """
        if not selected_values:
            return
        
        total_count = len(data_series)
        value_counts = data_series.value_counts()
        
        # 计算总的选择统计
        total_selected_count = sum(
            value_counts.get(v, 0) for v in selected_values
        )
        total_selected_percentage = (total_selected_count / total_count * 100)
        
        # 显示汇总信息
        col1, col2, col3 = st.columns(3)
        
        with col1:
            st.metric("选中值数量", len(selected_values))
        with col2:
            st.metric("覆盖数据点", total_selected_count)
        with col3:
            st.metric("覆盖比例", f"{total_selected_percentage:.1f}%")
        
        # 显示选中值的详细列表
        with st.expander("查看选中的值", expanded=False):
            selected_details = []
            for value in selected_values:
                count = value_counts.get(value, 0)
                percentage = (count / total_count * 100)
                selected_details.append({
                    '值': str(value)[:100] + ('...' if len(str(value)) > 100 else ''),
                    '数量': count,
                    '占比(%)': round(percentage, 2)
                })
            
            if selected_details:
                details_df = pd.DataFrame(selected_details)
                st.dataframe(details_df, width='stretch', hide_index=True)