"""
SQL 查询视图 - SQL 编辑、执行、结果分析与导出
"""

import io
import hashlib
import streamlit as st
import duckdb

from super_datascope.components.json_table import JsonDataTable
from super_datascope.components.sql_editor import render_sql_editor
from super_datascope.components.data_insight import render_reusable_analysis_tabs


@st.cache_data
def execute_cached_sql_query(query, tables_hash):
    """缓存的SQL查询执行函数"""
    conn = duckdb.connect()

    # 根据当前状态注册表
    if 'tables' in st.session_state and st.session_state.tables:
        for i, (_, table_info) in enumerate(st.session_state.tables.items(), 1):
            table_data = table_info['data'].copy()
            if 'index' not in table_data.columns:
                table_data = table_data.reset_index(drop=False)
                if 'index' not in table_data.columns:
                    table_data['index'] = range(len(table_data))
            conn.register(f'data{i}', table_data)
    elif 'data' in st.session_state and st.session_state.data is not None:
        table_data = st.session_state.data.copy()
        if 'index' not in table_data.columns:
            table_data = table_data.reset_index(drop=False)
            if 'index' not in table_data.columns:
                table_data['index'] = range(len(table_data))
        conn.register('data', table_data)

    result = conn.execute(query).df()
    return result


def sql_query_view():
    """SQL查询视图页面"""

    has_single_data = 'data' in st.session_state and st.session_state.data is not None
    has_multiple_tables = 'tables' in st.session_state and st.session_state.tables

    if not has_single_data and not has_multiple_tables:
        st.warning("请先选择并加载数据文件")
        return

    st.subheader("🔍 SQL 查询")

    # 显示可用表信息
    if has_multiple_tables:
        tables = st.session_state.tables
        default_query = _render_multi_table_info(tables)
    else:
        default_query = _render_single_table_info(st.session_state.data)

    # 示例查询
    sample_queries = _build_sample_queries(default_query, has_multiple_tables,
                                           tables if has_multiple_tables else None)

    if 'current_query' not in st.session_state:
        st.session_state.current_query = default_query

    # SQL 编辑器
    editor_data = st.session_state.data if has_single_data else list(tables.values())[0]['data']
    result = render_sql_editor(
        data=editor_data,
        value=st.session_state.current_query,
        height=200,
        key="sql_query_editor",
        sample_queries=sample_queries,
        placeholder="输入SQL查询语句..."
    )

    query = result["query"]
    should_execute = result["should_execute"]

    if query != st.session_state.current_query:
        st.session_state.current_query = query

    # 执行查询
    if should_execute and query:
        _execute_and_display(query, has_multiple_tables)
        st.session_state.query_executed = True
        st.session_state.executed_query = query
    elif st.session_state.get('query_executed') and st.session_state.get('executed_query'):
        _execute_and_display(st.session_state.executed_query, has_multiple_tables)


# ---------------------------------------------------------------------------
# 内部函数
# ---------------------------------------------------------------------------

def _render_multi_table_info(tables) -> str:
    """渲染多表信息，返回默认查询"""
    st.write("**可用的表：**")

    for i, (original_table_name, table_info) in enumerate(tables.items(), 1):
        table_name = f"data{i}"
        table_data = table_info['data']
        display_cols = list(table_data.columns)
        if 'index' not in display_cols:
            display_cols = ['index'] + display_cols

        with st.expander(
            f"📋 **{table_name}** (原名: {original_table_name}) "
            f"- {len(table_data)} 行, {len(display_cols)} 列",
            expanded=False
        ):
            st.write(f"**数据源：** {table_info.get('source', '')}")
            st.write("**列名：**")
            for chunk_start in range(0, len(display_cols), 3):
                chunk = display_cols[chunk_start:chunk_start + 3]
                st.text(", ".join(f"`{c}`" for c in chunk))

    st.write("您可以使用 data1, data2, ... 等表名进行 JOIN 查询。")
    return "SELECT * FROM data1 LIMIT 10"


def _render_single_table_info(df) -> str:
    """渲染单表信息，返回默认查询"""
    st.write("您可以使用SQL语句查询数据（表名为 'data'）")

    display_cols = list(df.columns)
    if 'index' not in display_cols:
        display_cols = ['index'] + display_cols

    with st.expander(f"📋 **data** ({len(df)} 行, {len(display_cols)} 列)", expanded=False):
        st.write("**列名：**")
        for chunk_start in range(0, len(display_cols), 3):
            chunk = display_cols[chunk_start:chunk_start + 3]
            st.text(", ".join(f"`{c}`" for c in chunk))

    return "SELECT * FROM data LIMIT 10"


def _build_sample_queries(default_query, is_multi, tables=None):
    """构建示例查询列表"""
    samples = [default_query]
    if is_multi and tables:
        samples.append("SELECT COUNT(*) as total_records FROM data1")
        if len(tables) >= 2:
            samples.append(
                "-- 示例JOIN查询\n"
                "-- SELECT a.*, b.* FROM data1 a JOIN data2 b "
                "ON a.common_field = b.common_field LIMIT 10"
            )
    else:
        samples.append("SELECT COUNT(*) as total_records FROM data")
    return samples


def _execute_and_display(query, is_multi):
    """执行查询并显示结果"""
    try:
        if is_multi:
            tables = st.session_state.tables
            tables_hash = hash(tuple(
                (f"data{i}", tuple(info['data'].columns), info['data'].shape[0])
                for i, (_, info) in enumerate(tables.items(), 1)
            ))
        else:
            df = st.session_state.data
            tables_hash = hash(tuple(df.columns)) + hash(df.shape[0])

        result = execute_cached_sql_query(query, tables_hash)

        st.subheader("查询结果")
        query_hash = hashlib.md5(query.encode()).hexdigest()[:8]

        json_table = JsonDataTable(result, name=f"sql_query_{query_hash}")
        json_table.display(page_size=20, height=500)

        if not result.empty:
            st.divider()
            st.subheader("📊 查询结果分析")
            render_reusable_analysis_tabs(result, tab_prefix=f"query_{query_hash}")

        # 导出
        _render_export_section(result, query_hash)

    except Exception as e:
        st.error(f"SQL查询执行错误: {str(e)}")


def _render_export_section(result_df, query_hash):
    """渲染本地导出功能"""
    if len(result_df) == 0:
        return

    st.divider()
    st.subheader("📤 导出查询结果")

    col1, col2 = st.columns(2)

    with col1:
        csv_data = result_df.to_csv(index=False)
        st.download_button(
            "📥 下载 CSV",
            csv_data,
            file_name=f"query_result_{query_hash}.csv",
            mime="text/csv",
            key=f"download_csv_{query_hash}"
        )

    with col2:
        buffer = io.BytesIO()
        result_df.to_parquet(buffer, index=False)
        st.download_button(
            "📥 下载 Parquet",
            buffer.getvalue(),
            file_name=f"query_result_{query_hash}.parquet",
            mime="application/octet-stream",
            key=f"download_parquet_{query_hash}"
        )
