"""
可复用的图表构建组件
支持多种图表类型和数据分析模式
"""

import streamlit as st
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import duckdb
from typing import List, Dict, Any, Optional, Union
from enum import Enum

from super_datascope.components.sql_editor import render_sql_editor    

class ChartType(Enum):
    BAR = "bar"
    PIE = "pie"
    LINE = "line"
    SCATTER = "scatter"
    HISTOGRAM = "histogram"
    BOX = "box"
    AREA = "area"
    HEATMAP = "heatmap"

class AggregationFunction(Enum):
    COUNT = "count"
    SUM = "sum"
    MEAN = "mean"
    MEDIAN = "median"
    MIN = "min"
    MAX = "max"
    STD = "std"
    VAR = "var"

class ChartBuilder:
    """图表构建器 - 支持多种数据分析和可视化模式"""
    
    def __init__(self, data: pd.DataFrame, key_prefix: str = "chart_builder"):
        self.data = data
        self.key_prefix = key_prefix
        
    def _get_key(self, suffix: str) -> str:
        """生成唯一的组件key"""
        return f"{self.key_prefix}_{suffix}"
    
    def _get_column_types(self) -> Dict[str, List[str]]:
        """获取按类型分类的列名"""
        numeric_cols = self.data.select_dtypes(include=['number']).columns.tolist()
        categorical_cols = self.data.select_dtypes(include=['object', 'category']).columns.tolist()
        datetime_cols = self.data.select_dtypes(include=['datetime64']).columns.tolist()
        
        return {
            'numeric': numeric_cols,
            'categorical': categorical_cols,
            'datetime': datetime_cols,
            'all': self.data.columns.tolist()
        }
    
    def render_column_selector(self, 
                              label: str,
                              column_types: Union[str, List[str]] = 'all',
                              key_suffix: str = "",
                              allow_none: bool = False) -> Optional[str]:
        """渲染列选择器
        
        Args:
            label: 选择器标签
            column_types: 列类型过滤 ('all', 'numeric', 'categorical', 'datetime' 或列表)
            key_suffix: key后缀
            allow_none: 是否允许选择无
        """
        cols_by_type = self._get_column_types()
        
        if isinstance(column_types, str):
            available_cols = cols_by_type.get(column_types, [])
        else:
            available_cols = column_types
        
        options = ['无'] + available_cols if allow_none else available_cols
        
        if not available_cols:
            st.warning(f"没有可用的列用于{label}")
            return None
            
        selected = st.selectbox(
            label,
            options=options,
            key=self._get_key(f"col_select_{key_suffix}")
        )
        
        return None if selected == '无' else selected
    
    def render_chart_type_selector(self, 
                                  key_suffix: str = "",
                                  allowed_types: List[ChartType] = None) -> ChartType:
        """渲染图表类型选择器"""
        if allowed_types is None:
            allowed_types = list(ChartType)
        
        type_names = {
            ChartType.BAR: "柱状图",
            ChartType.PIE: "饼图", 
            ChartType.LINE: "折线图",
            ChartType.SCATTER: "散点图",
            ChartType.HISTOGRAM: "直方图",
            ChartType.BOX: "箱线图",
            ChartType.AREA: "面积图",
            ChartType.HEATMAP: "热力图"
        }
        
        selected_name = st.selectbox(
            "图表类型",
            options=[type_names[t] for t in allowed_types],
            key=self._get_key(f"chart_type_{key_suffix}")
        )
        
        # 反向查找枚举值
        for chart_type, name in type_names.items():
            if name == selected_name:
                return chart_type
        
        return allowed_types[0]
    
    def render_aggregation_selector(self, key_suffix: str = "") -> AggregationFunction:
        """渲染聚合函数选择器"""
        agg_names = {
            AggregationFunction.COUNT: "计数",
            AggregationFunction.SUM: "求和",
            AggregationFunction.MEAN: "平均值",
            AggregationFunction.MEDIAN: "中位数",
            AggregationFunction.MIN: "最小值",
            AggregationFunction.MAX: "最大值",
            AggregationFunction.STD: "标准差",
            AggregationFunction.VAR: "方差"
        }
        
        selected_name = st.selectbox(
            "聚合函数",
            options=list(agg_names.values()),
            key=self._get_key(f"agg_func_{key_suffix}")
        )
        
        # 反向查找枚举值
        for agg_func, name in agg_names.items():
            if name == selected_name:
                return agg_func
        
        return AggregationFunction.COUNT
    
    def create_chart(self, 
                    chart_type: ChartType,
                    x_col: Optional[str] = None,
                    y_col: Optional[str] = None,
                    color_col: Optional[str] = None,
                    group_col: Optional[str] = None,
                    agg_func: AggregationFunction = AggregationFunction.COUNT,
                    title: str = "",
                    **kwargs) -> go.Figure:
        """创建图表
        
        Args:
            chart_type: 图表类型
            x_col: X轴列
            y_col: Y轴列
            color_col: 颜色分组列
            group_col: 分组统计列
            agg_func: 聚合函数
            title: 图表标题
            **kwargs: 其他plotly参数
        """
        # 准备数据
        plot_data = self._prepare_data(x_col, y_col, color_col, group_col, agg_func)
        
        if plot_data.empty:
            st.warning("没有数据可用于绘制图表")
            return go.Figure()
        
        # 如果y_col为None但数据中有count列，使用count作为y_col
        if y_col is None and 'count' in plot_data.columns:
            y_col = 'count'
        
        # 根据图表类型创建图表
        if chart_type == ChartType.BAR:
            return self._create_bar_chart(plot_data, x_col, y_col, color_col, title, **kwargs)
        elif chart_type == ChartType.PIE:
            return self._create_pie_chart(plot_data, x_col, y_col, title, **kwargs)
        elif chart_type == ChartType.LINE:
            return self._create_line_chart(plot_data, x_col, y_col, color_col, title, **kwargs)
        elif chart_type == ChartType.SCATTER:
            return self._create_scatter_chart(plot_data, x_col, y_col, color_col, title, **kwargs)
        elif chart_type == ChartType.HISTOGRAM:
            return self._create_histogram_chart(plot_data, x_col, title, **kwargs)
        elif chart_type == ChartType.BOX:
            return self._create_box_chart(plot_data, x_col, y_col, color_col, title, **kwargs)
        elif chart_type == ChartType.AREA:
            return self._create_area_chart(plot_data, x_col, y_col, color_col, title, **kwargs)
        elif chart_type == ChartType.HEATMAP:
            return self._create_heatmap_chart(plot_data, x_col, y_col, title, **kwargs)
        else:
            st.error(f"不支持的图表类型: {chart_type}")
            return go.Figure()
    
    def _prepare_data(self, x_col, y_col, color_col, group_col, agg_func) -> pd.DataFrame:
        """准备绘图数据"""
        if group_col and x_col:
            # 分组统计模式
            if y_col and y_col in self.data.columns and agg_func != AggregationFunction.COUNT:
                # 有数值列的分组聚合
                if agg_func == AggregationFunction.SUM:
                    result = self.data.groupby([x_col, group_col])[y_col].sum().reset_index()
                elif agg_func == AggregationFunction.MEAN:
                    result = self.data.groupby([x_col, group_col])[y_col].mean().reset_index()
                elif agg_func == AggregationFunction.MEDIAN:
                    result = self.data.groupby([x_col, group_col])[y_col].median().reset_index()
                elif agg_func == AggregationFunction.MIN:
                    result = self.data.groupby([x_col, group_col])[y_col].min().reset_index()
                elif agg_func == AggregationFunction.MAX:
                    result = self.data.groupby([x_col, group_col])[y_col].max().reset_index()
                elif agg_func == AggregationFunction.STD:
                    result = self.data.groupby([x_col, group_col])[y_col].std().reset_index()
                elif agg_func == AggregationFunction.VAR:
                    result = self.data.groupby([x_col, group_col])[y_col].var().reset_index()
                else:
                    result = self.data.groupby([x_col, group_col]).size().reset_index(name='count')
            else:
                # 分组计数
                result = self.data.groupby([x_col, group_col]).size().reset_index(name='count')
            return result
        elif x_col and (y_col is None or agg_func == AggregationFunction.COUNT):
            # 单列统计
            result = self.data[x_col].value_counts().reset_index()
            result.columns = [x_col, 'count']
            return result
        elif x_col and y_col and y_col in self.data.columns:
            # 双列分析 - 检查y_col是否存在于原数据中
            if self.data[y_col].dtype in ['object', 'category']:
                # Y列是分类数据，进行计数
                result = self.data.groupby([x_col, y_col]).size().reset_index(name='count')
                return result
            else:
                # Y列是数值数据，可以直接使用或进行分组聚合
                if x_col == y_col:
                    return self.data[[x_col]].copy()
                else:
                    return self.data[[x_col, y_col]].copy()
        elif x_col and y_col == 'count':
            # 当y_col是'count'时，说明需要进行计数统计
            result = self.data[x_col].value_counts().reset_index()
            result.columns = [x_col, 'count']
            return result
        elif x_col:
            # 仅有x_col，进行计数统计
            result = self.data[x_col].value_counts().reset_index()
            result.columns = [x_col, 'count']
            return result
        
        return self.data.copy()
    
    def _create_bar_chart(self, data, x_col, y_col, color_col, title, **kwargs):
        return px.bar(data, x=x_col, y=y_col, color=color_col, title=title, **kwargs)
    
    def _create_pie_chart(self, data, x_col, y_col, title, **kwargs):
        return px.pie(data, names=x_col, values=y_col, title=title, **kwargs)
    
    def _create_line_chart(self, data, x_col, y_col, color_col, title, **kwargs):
        return px.line(data, x=x_col, y=y_col, color=color_col, title=title, **kwargs)
    
    def _create_scatter_chart(self, data, x_col, y_col, color_col, title, **kwargs):
        return px.scatter(data, x=x_col, y=y_col, color=color_col, title=title, **kwargs)
    
    def _create_histogram_chart(self, data, x_col, title, **kwargs):
        return px.histogram(data, x=x_col, title=title, **kwargs)
    
    def _create_box_chart(self, data, x_col, y_col, color_col, title, **kwargs):
        return px.box(data, x=x_col, y=y_col, color=color_col, title=title, **kwargs)
    
    def _create_area_chart(self, data, x_col, y_col, color_col, title, **kwargs):
        return px.area(data, x=x_col, y=y_col, color=color_col, title=title, **kwargs)
    
    def _create_heatmap_chart(self, data, x_col, y_col, title, **kwargs):
        if x_col and y_col:
            pivot_data = data.pivot_table(values='count' if 'count' in data.columns else data.columns[-1], 
                                        index=y_col, columns=x_col, fill_value=0)
            return px.imshow(pivot_data, title=title, **kwargs)
        return go.Figure()


class InteractiveChartBuilder:
    """交互式图表构建器 - 提供完整的UI界面"""
    
    def __init__(self, data: pd.DataFrame, key_prefix: str = "interactive_chart"):
        self.data = data
        self.chart_builder = ChartBuilder(data, key_prefix)
        self.key_prefix = key_prefix
    
    def render(self) -> Optional[go.Figure]:
        """渲染完整的交互式图表构建界面"""
        st.subheader("📊 自定义数据可视化")
        
        # 分析模式选择
        analysis_mode = st.radio(
            "分析模式",
            options=["分类统计", "数值分析", "自定义SQL"],
            key=f"{self.key_prefix}_analysis_mode"
        )
        
        if analysis_mode == "分类统计":
            return self._render_categorical_analysis()
        elif analysis_mode == "数值分析":
            return self._render_numerical_analysis()
        elif analysis_mode == "自定义SQL":
            return self._render_sql_analysis()
        
        return None
    
    def _render_categorical_analysis(self) -> Optional[go.Figure]:
        """渲染分类数据分析界面"""
        col1, col2 = st.columns(2)
        
        with col1:
            x_col = self.chart_builder.render_column_selector(
                "选择主要分析列", 'categorical', "cat_x"
            )
            group_col = self.chart_builder.render_column_selector(
                "选择分组列（可选）", 'categorical', "cat_group", allow_none=True
            )
        
        with col2:
            chart_type = self.chart_builder.render_chart_type_selector(
                "cat", [ChartType.BAR, ChartType.PIE, ChartType.HEATMAP]
            )
        
        if x_col:
            title = f"{x_col} 分布统计"
            if group_col:
                title += f" (按 {group_col} 分组)"
            
            fig = self.chart_builder.create_chart(
                chart_type=chart_type,
                x_col=x_col,
                y_col=None,  # Let the chart builder determine y_col
                color_col=group_col,
                group_col=group_col,
                agg_func=AggregationFunction.COUNT,
                title=title
            )
            
            st.plotly_chart(fig, use_container_width=True)
            return fig
        
        return None
    
    def _render_numerical_analysis(self) -> Optional[go.Figure]:
        """渲染数值数据分析界面"""
        col1, col2, col3 = st.columns(3)
        
        with col1:
            y_col = self.chart_builder.render_column_selector(
                "选择数值列", 'numeric', "num_y"
            )
            x_col = self.chart_builder.render_column_selector(
                "选择X轴列（可选）", 'all', "num_x", allow_none=True
            )
        
        with col2:
            group_col = self.chart_builder.render_column_selector(
                "选择分组列（可选）", ['categorical', 'datetime'], "num_group", allow_none=True
            )
            agg_func = self.chart_builder.render_aggregation_selector("num")
        
        with col3:
            chart_type = self.chart_builder.render_chart_type_selector(
                "num", [ChartType.BAR, ChartType.LINE, ChartType.SCATTER, 
                       ChartType.HISTOGRAM, ChartType.BOX, ChartType.AREA]
            )
        
        if y_col:
            title = f"{y_col} "
            if agg_func != AggregationFunction.COUNT:
                agg_names = {
                    AggregationFunction.SUM: "求和",
                    AggregationFunction.MEAN: "平均值",
                    AggregationFunction.MEDIAN: "中位数",
                    AggregationFunction.MIN: "最小值",
                    AggregationFunction.MAX: "最大值",
                    AggregationFunction.STD: "标准差",
                    AggregationFunction.VAR: "方差"
                }
                title += agg_names.get(agg_func, "分析")
            else:
                title += "分析"
                
            if x_col:
                title += f" (按 {x_col})"
            if group_col:
                title += f" (分组: {group_col})"
            
            fig = self.chart_builder.create_chart(
                chart_type=chart_type,
                x_col=x_col or group_col,
                y_col=y_col,
                color_col=group_col if x_col else None,
                group_col=group_col,
                agg_func=agg_func,
                title=title
            )
            
            st.plotly_chart(fig, use_container_width=True)
            return fig
        
        return None
    
    def _render_sql_analysis(self) -> Optional[go.Figure]:
        """渲染SQL自定义分析界面"""     
        st.info("💡 使用SQL查询数据，然后将结果可视化。表名为 'data'。")
        
        # SQL编辑器
        result = render_sql_editor(
            data=self.data,
            value="SELECT * FROM data LIMIT 100",
            height=150,
            key=f"{self.key_prefix}_sql_editor",
            sample_queries=[
                "SELECT column_name, COUNT(*) as count FROM data GROUP BY column_name",
                "SELECT * FROM data WHERE numeric_column > 100",
                "SELECT category, AVG(value) as avg_value FROM data GROUP BY category"
            ]
        )
        
        query = result["query"]
        should_execute = result["should_execute"]
        
        if should_execute and query:
            try:
                # 执行SQL查询
                conn = duckdb.connect()
                conn.register('data', self.data)
                sql_result = conn.execute(query).df()
                
                if not sql_result.empty:
                    st.subheader("查询结果预览")
                    st.dataframe(sql_result.head(10), width='stretch')
                    
                    # 图表配置
                    col1, col2, col3 = st.columns(3)
                    
                    with col1:
                        chart_builder = ChartBuilder(sql_result, f"{self.key_prefix}_sql")
                        x_col = chart_builder.render_column_selector("X轴列", 'all', "sql_x")
                    
                    with col2:
                        y_col = chart_builder.render_column_selector("Y轴列（可选）", 'all', "sql_y", allow_none=True)
                    
                    with col3:
                        chart_type = chart_builder.render_chart_type_selector("sql")
                    
                    if x_col:
                        fig = chart_builder.create_chart(
                            chart_type=chart_type,
                            x_col=x_col,
                            y_col=y_col,
                            title=f"SQL查询结果可视化 - {x_col}"
                        )
                        
                        st.plotly_chart(fig, use_container_width=True)
                        return fig
                else:
                    st.warning("查询结果为空")
                    
            except Exception as e:
                st.error(f"SQL查询执行错误: {str(e)}")
        
        return None