AA

import json
import ast
import pandas as pd
import plotly.express as px
import streamlit as st
from st_aggrid import AgGrid, GridOptionsBuilder, GridUpdateMode, JsCode
from core.query_executor import query_executor
from queries.QueryBuilder import QueryBuilder
from queries.filter import CommonUI
from queries.final_last import USER_360_QUERIES

# ========================
# CONSTANTS & CONFIGURATION
# ========================
CSS_STYLES = """
<style>
    :root {
        --primary: #4361ee;
        --secondary: #3a0ca3;
        --accent: #4895ef;
        --success: #4cc9f0;
        --warning: #f72585;
        --light: #f8f9fa;
        --dark: #212529;
    }

    .main-header, .metric-card, .status-badge, .query-detail-card, 
    .recommendation-card, .grid-container {
        border-radius: 12px;
        box-shadow: 0 4px 12px rgba(0,0,0,0.08);
        transition: all 0.3s ease;
    }

    .main-header {
        background: linear-gradient(135deg, var(--primary) 0%, var(--secondary) 100%);
        padding: 1.5rem;
        margin-bottom: 1.5rem;
        text-align: center;
        color: white;
    }

    .metric-card {
        background: white;
        padding: 1.2rem;
        border-left: 4px solid var(--primary);
    }

    .metric-card:hover {
        transform: translateY(-3px);
        box-shadow: 0 6px 16px rgba(0,0,0,0.12);
    }

    .metric-value {
        font-size: 1.8rem;
        font-weight: 700;
        color: var(--dark);
    }

    .metric-label {
        color: #6c757d;
        font-size: 0.9rem;
    }

    .status-normal { background-color: #d8f3dc; color: #2d6a4f; }
    .status-warning { background-color: #fff3cd; color: #856404; }
    .status-critical { background-color: #f8d7da; color: #721c24; }
    .status-badge {
        padding: 0.3rem 0.8rem;
        border-radius: 20px;
        font-weight: 600;
        display: inline-block;
    }

    .query-detail-card {
        background: var(--light);
        padding: 1.5rem;
        border-left: 4px solid var(--accent);
    }

    .recommendation-card {
        background: #fff9db;
        padding: 1.2rem;
        border-left: 4px solid #ffd43b;
    }

    .grid-container {
        background: white;
        padding: 1rem;
        margin-top: 1.5rem;
    }

    .action-button {
        background: var(--accent) !important;
        color: white !important;
        border-radius: 8px !important;
        padding: 0.5rem 1rem !important;
        border: none !important;
        transition: all 0.3s !important;
        font-weight: 600 !important;
    }

    .action-button:hover {
        background: var(--secondary) !important;
        transform: scale(1.02);
    }

    .chart-container {
        background: white;
        border-radius: 12px;
        padding: 1.5rem;
        box-shadow: 0 4px 12px rgba(0,0,0,0.08);
    }

    .detail-header {
        display: flex;
        justify-content: space-between;
        align-items: center;
        margin-bottom: 1.5rem;
        padding-bottom: 0.5rem;
        border-bottom: 2px solid var(--light);
    }
</style>
"""

CLICKABLE_COLUMNS = [
    "SPILLED_QUERIES", "OVER_PROVISIONED_QUERIES", "PEAK_HOUR_LONG_RUNNING_QUERIES",
    "SELECT_STAR_QUERIES", "UNPARTITIONED_SCAN_QUERIES", "REPEATED_QUERIES",
    "COMPLEX_JOIN_QUERIES", "ZERO_RESULT_QUERIES", "HIGH_COMPILE_QUERIES",
    "UNTAGGED_QUERIES", "UNLIMITED_ORDER_BY_QUERIES", "LARGE_GROUP_BY_QUERIES",
    "SLOW_QUERIES", "EXPENSIVE_DISTINCT_QUERIES", "INEFFICIENT_LIKE_QUERIES",
    "NO_RESULTS_WITH_SCAN_QUERIES", "CARTESIAN_JOIN_QUERIES", "HIGH_COMPILE_RATIO_QUERIES"
]

# ========================
# HELPER FUNCTIONS
# ========================
def init_session_state():
    """Initialize session state variables"""
    defaults = {
        "selected_user": None,
        "selected_metric": None,
        "show_table": True,
        "df": pd.DataFrame(),
        "aggrid_key": 0
    }
    for key, value in defaults.items():
        if key not in st.session_state:
            st.session_state[key] = value

def parse_uploaded_data(uploaded_file):
    """Parse uploaded JSON file into DataFrame"""
    if uploaded_file is None:
        return None

    try:
        data = json.load(uploaded_file)
        return pd.DataFrame(data) if isinstance(data, list) else pd.DataFrame([data])
    except Exception as e:
        st.error(f"Error loading data: {str(e)}")
        return None

def parse_json_input(json_str):
    """Parse JSON string input into DataFrame"""
    if not json_str.strip():
        return None

    try:
        data = json.loads(json_str)
        return pd.DataFrame(data) if isinstance(data, list) else pd.DataFrame([data])
    except Exception as e:
        st.error(f"Error parsing JSON: {str(e)}")
        return None

def format_metric_value(value, format_type):
    """Format metric values based on type"""
    if format_type == "currency":
        return f"${value:.2f}"
    elif format_type == "percentage":
        return f"{value:.2f}%"
    elif format_type == "time":
        return f"{value:.1f}ms"
    elif format_type == "size":
        return f"{value:.2f}GB"
    return str(value)

def safe_parse_samples(data):
    """Safely parse QUERY_SAMPLES data"""
    if isinstance(data, str):
        try:
            return ast.literal_eval(data)
        except:
            return {}
    return data

def safe_parse_recommendations(data):
    """Safely parse RECOMMENDATIONS data"""
    if isinstance(data, str):
        try:
            return ast.literal_eval(data)
        except:
            return []
    return data

def normalize_metric_name(name):
    """Normalize metric names for display"""
    return name.lower().replace("_queries", "").replace("_", " ").replace(" ", "_")

# ========================
# COMPONENT RENDERING
# ========================
def render_header():
    """Render dashboard header"""
    st.markdown(
        """
        <div class="main-header">
            <h1>❄️ Snowflake Query Analytics</h1>
            <p>Monitor and optimize query performance in real-time</p>
        </div>
        """,
        unsafe_allow_html=True
    )

def render_sidebar():
    """Render sidebar components"""
    with st.sidebar:
        st.header("📊 Data Input")
        uploaded_df = None

        # File uploader
        uploaded_file = st.file_uploader("Upload JSON Data", type=["json"])
        if uploaded_file:
            uploaded_df = parse_uploaded_data(uploaded_file)

        # JSON input
        st.subheader("Or paste JSON data:")
        json_input = st.text_area("JSON Data", height=150, placeholder='{"key": "value"}')
        if st.button("Load JSON Data", use_container_width=True):
            uploaded_df = parse_json_input(json_input)

        # Table visibility toggle
        st.markdown("---")
        st.subheader("Table Controls")
        if st.button("Toggle Table Visibility", use_container_width=True):
            st.session_state.show_table = not st.session_state.show_table

        return uploaded_df

def render_overall_metrics(df):
    """Render top-level metrics cards"""
    st.header("📈 Performance Overview")

    # Define metrics grid
    metrics = [
        ("Total Users", len(df), None, "count"),
        ("Total Queries", df["TOTAL_QUERIES"].sum(), None, "number"),
        ("Total Credits", df["TOTAL_CREDITS"].sum(), None, "currency"),
        ("Avg Weighted Score", df["WEIGHTED_SCORE"].mean(), None, "score"),
        ("Total Spilled Queries", df["SPILLED_QUERIES"].sum(), None, "number"),
        ("Total Slow Queries", df["SLOW_QUERIES"].sum(), None, "number"),
        ("SELECT * Queries", df["SELECT_STAR_QUERIES"].sum(), None, "number"),
        ("Avg Failure Rate", df["FAILURE_CANCELLATION_RATE_PCT"].mean(), None, "percentage")
    ]

    # Render in responsive grid
    cols = st.columns(4)
    for i, (label, value, delta, fmt) in enumerate(metrics):
        with cols[i % 4]:
            formatted_value = format_metric_value(value, fmt)
            st.markdown(
                f"""
                <div class="metric-card">
                    <div class="metric-value">{formatted_value}</div>
                    <div class="metric-label">{label}</div>
                </div>
                """,
                unsafe_allow_html=True
            )

def render_query_details(df):
    """Render detailed query view when a cell is selected"""
    if not st.session_state.selected_user or not st.session_state.selected_metric:
        return

    user_data = df[df["USER_NAME"] == st.session_state.selected_user].iloc[0]
    metric_value = user_data[st.session_state.selected_metric]

    # Header section
    with st.container():
        st.markdown(
            f"""
            <div class="detail-header">
                <div>
                    <h2>🔍 Detailed Analysis: {st.session_state.selected_user}</h2>
                    <h4>Metric: {st.session_state.selected_metric.replace('_', ' ').title()}</h4>
                </div>
                <div>
                    <button class="action-button" onclick="window.scrollTo(0,0)">Back to Top</button>
                </div>
            </div>
            """,
            unsafe_allow_html=True
        )

    # Main content columns
    col1, col2 = st.columns([3, 1])

    with col1:
        # User metrics card
        st.markdown(
            f"""
            <div class="query-detail-card">
                <div class="metric-value">{metric_value}</div>
                <div class="metric-label">Current Metric Value</div>
                <div style="margin-top: 1.5rem">
                    <div><strong>User:</strong> {st.session_state.selected_user}</div>
                    <div><strong>Total Queries:</strong> {user_data['TOTAL_QUERIES']}</div>
                    <div><strong>Weighted Score:</strong> {user_data['WEIGHTED_SCORE']:.1f}</div>
                    <div><strong>Status:</strong> 
                        <span class="status-badge status-{user_data['COST_STATUS'].lower()}">
                            {user_data['COST_STATUS']}
                        </span>
                    </div>
                </div>
            </div>
            """,
            unsafe_allow_html=True
        )

        # Sample queries section
        if "QUERY_SAMPLES" in user_data:
            query_samples = safe_parse_samples(user_data["QUERY_SAMPLES"])
            sample_key = normalize_metric_name(st.session_state.selected_metric)

            if sample_key in query_samples and query_samples[sample_key]:
                st.subheader("📋 Query Examples")
                for i, sample in enumerate(query_samples[sample_key][:5]):
                    with st.expander(f"Query {i+1}: {sample.get('query_id', 'ID')}", expanded=False):
                        st.code(sample.get("query_text", "No query available"), language="sql")

                        # Metrics grid
                        cols = st.columns(3)
                        metrics = [
                            ("Execution Time", f"{sample.get('execution_time_ms', 0):,}ms", None),
                            ("Bytes Scanned", f"{sample.get('bytes_scanned', 0):,}", None),
                            ("Warehouse Size", sample.get('warehouse_size', 'N/A'), None)
                        ]

                        for col, (label, value, delta) in zip(cols, metrics):
                            with col:
                                st.metric(label, value, delta)

                        st.caption(f"Started at: {sample.get('start_time', 'N/A')}")

        # Recommendations section
        if "RECOMMENDATIONS" in user_data:
            recommendations = safe_parse_recommendations(user_data["RECOMMENDATIONS"])
            if recommendations:
                st.subheader("💡 Optimization Suggestions")
                for rec in recommendations[:3]:
                    st.markdown(
                        f'<div class="recommendation-card">🔍 {rec}</div>',
                        unsafe_allow_html=True
                    )

    with col2:
        # Action buttons
        if st.button("Clear Selection", use_container_width=True, type="primary"):
            st.session_state.selected_user = None
            st.session_state.selected_metric = None
            st.rerun()

        if st.button("Show Table", use_container_width=True):
            st.session_state.show_table = True
            st.rerun()

def configure_aggrid(df):
    """Configure AgGrid options for user table"""
    gb = GridOptionsBuilder.from_dataframe(df)

    # Define column types and formatting
    gb.configure_column("USER_NAME", headerName="User", pinned="left", width=200)

    for col in CLICKABLE_COLUMNS:
        gb.configure_column(
            col, 
            headerName=col.replace("_", " ").title(),
            type=["numericColumn"],
            width=150
        )

    gb.configure_column(
        "WEIGHTED_SCORE", 
        headerName="Score", 
        type=["numericColumn"], 
        valueFormatter="value.toFixed(1)",
        width=120
    )

    gb.configure_column(
        "COST_STATUS", 
        headerName="Status", 
        cellStyle=JsCode(
            """
            function(params) {
                const status = params.value.toLowerCase();
                const colors = {
                    'normal': { bg: '#d8f3dc', text: '#2d6a4f' },
                    'warning': { bg: '#fff3cd', text: '#856404' },
                    'critical': { bg: '#f8d7da', text: '#721c24' }
                };
                const color = colors[status] || { bg: '#e9ecef', text: '#495057' };
                return {
                    'backgroundColor': color.bg,
                    'color': color.text,
                    'borderRadius': '20px',
                    'fontWeight': '600',
                    'textAlign': 'center'
                };
            }
            """
        ),
        width=120
    )

    # Add selection functionality
    gb.configure_selection(
        selection_mode="single",
        use_checkbox=False,
        suppressRowDeselection=True,
        suppressRowClickSelection=False
    )

    # Add pagination
    gb.configure_pagination(
        paginationAutoPageSize=False,
        paginationPageSize=10
    )

    # Add grid options
    grid_options = gb.build()
    grid_options["rowHeight"] = 50
    grid_options["domLayout"] = "autoHeight"

    return grid_options

def render_user_table(df):
    """Render interactive user analytics table using AgGrid"""
    st.header("👥 User Performance Analysis")

    # Create display dataframe
    display_df = df[["USER_NAME"] + CLICKABLE_COLUMNS + ["WEIGHTED_SCORE", "COST_STATUS"]].copy()

    # Configure AgGrid
    grid_options = configure_aggrid(display_df)

    # Render AgGrid component
    with st.container():
        st.markdown('<div class="grid-container">', unsafe_allow_html=True)
        grid_response = AgGrid(
            display_df,
            gridOptions=grid_options,
            update_mode=GridUpdateMode.SELECTION_CHANGED,
            allow_unsafe_jscode=True,
            theme="streamlit",
            height=500,
            key=f"aggrid-{st.session_state.aggrid_key}"
        )
        st.markdown("</div>", unsafe_allow_html=True)

    # Handle row selection
    if grid_response["selected_rows"]:
        selected_row = grid_response["selected_rows"][0]
        st.session_state.selected_user = selected_row["USER_NAME"]
        st.session_state.selected_metric = next(
            (col for col in CLICKABLE_COLUMNS if selected_row[col] > 0),
            CLICKABLE_COLUMNS[0]
        )
        st.session_state.show_table = False
        st.rerun()

def render_analytics_charts(df):
    """Render data visualization charts"""
    st.header("📊 Data Visualization")

    col1, col2 = st.columns(2)

    with col1:
        with st.container():
            st.markdown('<div class="chart-container">', unsafe_allow_html=True)
            fig = px.scatter(
                df,
                x="TOTAL_QUERIES",
                y="TOTAL_CREDITS",
                size="TOTAL_DATA_SCANNED_GB",
                color="WEIGHTED_SCORE",
                hover_name="USER_NAME",
                title="Cost vs Query Volume",
                labels={
                    "TOTAL_QUERIES": "Total Queries",
                    "TOTAL_CREDITS": "Total Cost ($)",
                    "WEIGHTED_SCORE": "Performance Score"
                }
            )
            st.plotly_chart(fig, use_container_width=True)
            st.markdown("</div>", unsafe_allow_html=True)

    with col2:
        with st.container():
            st.markdown('<div class="chart-container">', unsafe_allow_html=True)
            status_counts = df["COST_STATUS"].value_counts()
            fig = px.pie(
                names=status_counts.index,
                values=status_counts.values,
                title="Performance Distribution",
                color=status_counts.index,
                color_discrete_map={
                    "Normal": "#d8f3dc",
                    "Warning": "#fff3cd",
                    "Critical": "#f8d7da"
                }
            )
            st.plotly_chart(fig, use_container_width=True)
            st.markdown("</div>", unsafe_allow_html=True)

def render_no_data():
    """Render empty state"""
    st.info("📭 No data available. Upload JSON data to begin analysis")
    st.markdown("""
    ### Expected Data Format:
    ```

json
    [
        {
            "USER_NAME": "user1",
            "TOTAL_QUERIES": 150,
            "TOTAL_CREDITS": 25.50,
            "WEIGHTED_SCORE": 82.5,
            "COST_STATUS": "Normal",
            "SPILLED_QUERIES": 5,
            // ...other metrics,
            "QUERY_SAMPLES": {
                "spilled_queries": [
                    {
                        "query_id": "01a2b3c4",
                        "query_text": "SELECT ...",
                        "execution_time_ms": 4500
                    }
                ]
            },
            "RECOMMENDATIONS": [
                "Optimize join operations",
                "Increase warehouse size"
            ]
        }
    ]


    ```
    """)

# ========================
# MAIN DASHBOARD FUNCTION
# ========================
def create_snowflake_dashboard(df):
    """Main dashboard rendering function"""
    # Apply global styles
    st.markdown(CSS_STYLES, unsafe_allow_html=True)

    # Initialize session state
    init_session_state()

    # Render header
    render_header()

    # Sidebar and data loading
    uploaded_df = render_sidebar()
    current_df = uploaded_df if uploaded_df is not None else df

    # Main content
    if not current_df.empty:
        render_overall_metrics(current_df)
        st.markdown("---")

        # Detailed view if selected
        if st.session_state.selected_user:
            render_query_details(current_df)
            st.markdown("---")

        # Table view toggle
        if st.session_state.show_table:
            render_user_table(current_df)
            st.markdown("---")

        # Charts section
        render_analytics_charts(current_df)
    else:
        render_no_data()

# ========================
# EXECUTION FLOW
# ========================
if __name__ == "__main__":
    # Filters section
    st.subheader("🔧 Filters")
    col1, col2 = st.columns(2)

    with col1:
        date_filter, custom_start, custom_end = CommonUI.render_date_filter()

    with col2:
        selected_user = CommonUI.render_object_filter("user", query_executor)

    # Prepare query parameters
    start_date, end_date = query_executor.get_date_range(date_filter, custom_start, custom_end)
    object_filter = QueryBuilder.build_object_filter("user", selected_user)

    # Execute query
    df = pd.DataFrame()
    query_key = "user_all_queries"

    try:
        query = QueryBuilder.prepare_query(
            USER_360_QUERIES, 
            query_key, 
            start_date, 
            end_date, 
            object_filter
        )

        if query:
            data = query_executor._execute_single_query(query, {}, query_key)
            df = data if data is not None else pd.DataFrame()
    except Exception as e:
        st.error(f"Query execution failed: {str(e)}")

    # Render dashboard
    create_snowflake_dashboard(df)