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

try:
    from streamlit_ace import st_ace
    ACE_AVAILABLE = True
except ImportError:
    ACE_AVAILABLE = False
    st.warning("streamlit-ace package not installed. Using fallback text area.")

class SqlEditor:
    """
    SQL编辑器组件，支持语法高亮和自动补全
    
    注意：当前版本的streamlit-ace不支持自定义completers参数，
    因此暂时只能提供语法高亮功能。未来版本可能会支持列名自动补全。
    """
    
    def __init__(self, 
                 data: Optional[pd.DataFrame] = None,
                 height: int = 200,
                 theme: str = "monokai",
                 font_size: int = 14,
                 sample_queries: List[str] = None,
                 key: str = "sql_editor"):
        """
        初始化SQL编辑器
        
        Args:
            data: 数据框，用于提取列名进行自动补全
            height: 编辑器高度
            theme: 编辑器主题
            font_size: 字体大小
            key: Streamlit组件key
            selected_sample: 选中的示例查询
        """
        self.data = data
        self.height = height
        self.theme = theme
        self.font_size = font_size
        self._key = key
        self.sample_queries = sample_queries if sample_queries else []
        self.selected_sample = None
        self._selected_sample_changed = False
    
    @property
    def key(self) -> str:
        _key_inc_id = st.session_state.get(('inc_id', self._key), 0)
        if _key_inc_id:
            return self._key+'.'+str(_key_inc_id)
        else:
            return self._key
    
    def _get_column_completers(self) -> List[Dict[str, Any]]:
        """获取列名自动补全项"""
        completers = []
        
        if self.data is not None:
            # 添加列名补全
            for col in self.data.columns:
                completers.append({
                    "name": col,
                    "value": col,
                    "caption": f"Column: {col}",
                    "meta": "column",
                    "score": 100
                })
        
        # 添加SQL关键字补全
        sql_keywords = [
            "SELECT", "FROM", "WHERE", "GROUP BY", "ORDER BY", "HAVING",
            "JOIN", "INNER JOIN", "LEFT JOIN", "RIGHT JOIN", "FULL JOIN",
            "COUNT", "SUM", "AVG", "MIN", "MAX", "DISTINCT",
            "AND", "OR", "NOT", "IN", "LIKE", "BETWEEN", "IS NULL", "IS NOT NULL",
            "LIMIT", "OFFSET", "AS", "CASE", "WHEN", "THEN", "ELSE", "END"
        ]
        
        for keyword in sql_keywords:
            completers.append({
                "name": keyword,
                "value": keyword,
                "caption": f"Keyword: {keyword}",
                "meta": "keyword",
                "score": 50
            })
        
        # 添加DuckDB特定函数
        duckdb_functions = [
            "STRFTIME", "EXTRACT", "DATE_PART", "REGEXP_MATCHES", "ARRAY_AGG",
            "UNNEST", "JSON_EXTRACT", "JSON_EXTRACT_PATH", "STRUCT_PACK",
            "LIST_CONTAINS", "LIST_TRANSFORM", "LIST_FILTER"
        ]
        
        for func in duckdb_functions:
            completers.append({
                "name": func,
                "value": func + "()",
                "caption": f"Function: {func}",
                "meta": "function",
                "score": 80
            })
        
        return completers

    def _on_selected_sample_changed(self):
        # 执行顺序：先触发on_change事件，然后再执行st.selectbox的返回值；
        #   所以本函数内只能获取到旧值，不能获取到新值；
        st.session_state[('selected_sample_changed', self._key)] = True

    def render(self, 
               value: str = "",
               placeholder: str = "输入SQL查询语句...",
               wrap: bool = True,
               auto_update: bool = False) -> Dict[str, Any]:
        """
        渲染SQL编辑器
        
        Args:
            value: 初始值
            placeholder: 占位符文本
            wrap: 是否自动换行
            auto_update: 是否自动更新
            sample_queries: 示例查询列表
            
        Returns:
            包含query和should_execute的字典
        """
        if self.sample_queries:  
            self.selected_sample = st.selectbox(
                "选择示例查询", 
                ["自定义"] + self.sample_queries,
                on_change=self._on_selected_sample_changed
            )

        # 处理选中的示例查询
        has_selected_sample = bool(self.selected_sample and self.selected_sample != "自定义")
        if has_selected_sample:
            value = self.selected_sample

        if st.session_state.get(('selected_sample_changed', self._key)):
            # 代表selectbox变动了
            del st.session_state[('selected_sample_changed', self._key)]
            if has_selected_sample:
                # 让key变动，以使得st_ace组件更新输入数据
                st.session_state[('inc_id', self._key)] = st.session_state.get(('inc_id', self._key), 0) + 1
        
        if ACE_AVAILABLE:
            # 使用ACE编辑器
            try:
                sql_content = st_ace(
                    value=value,
                    language="sql",
                    theme=self.theme,
                    key=self.key,
                    height=self.height,
                    font_size=self.font_size,
                    wrap=wrap,
                    auto_update=auto_update
                )
                should_execute = bool(sql_content) 
                return {
                    "query": sql_content,
                    "should_execute": should_execute
                }
            except Exception as e:
                st.warning(f"SQL编辑器加载失败，使用基础文本框: {str(e)}")
        else:
            # 回退到普通文本框
            query = st.text_area(
                "SQL查询语句",
                value=value,
                height=self.height,
                placeholder=placeholder,
                key=f"{self.key}_fallback"
            )
            
            should_execute = st.button("🚀 执行查询", key=f"{self.key}_fallback_execute") and bool(query)
            
            return {
                "query": query,
                "should_execute": should_execute
            }
    
    @staticmethod
    def create_with_data(data: pd.DataFrame, **kwargs) -> 'SqlEditor':
        """
        使用数据框创建SQL编辑器的便捷方法
        
        Args:
            data: 数据框
            **kwargs: 其他参数
            
        Returns:
            配置好的SqlEditor实例
        """
        return SqlEditor(data=data, **kwargs)


def render_sql_editor(data: Optional[pd.DataFrame] = None,
                     value: str = "",
                     height: int = 200,
                     theme: str = "monokai",
                     key: str = "sql_editor",
                     sample_queries: List[str] = None,
                     **kwargs) -> Dict[str, Any]:
    """
    渲染SQL编辑器的便捷函数
    
    Args:
        data: 数据框，用于列名自动补全
        value: 初始SQL内容
        height: 编辑器高度
        theme: 编辑器主题
        key: Streamlit组件key
        sample_queries: 示例查询列表
        **kwargs: 其他参数
        
    Returns:
        包含query和should_execute的字典
    """
    editor = SqlEditor(
        data=data,
        height=height,
        theme=theme,
        key=key,
        sample_queries=sample_queries
    )
    
    return editor.render(value=value, **kwargs)