import gradio as gr
from datasets import load_dataset
import pandas as pd

DATASETS = {
    "CS1": "withmartian/cs1_dataset",
    "CS2": "withmartian/cs2_dataset",
    "CS3": "withmartian/cs3_dataset",
    "CS2 Synonyms": "withmartian/cs2_dataset_synonyms",
    "CS3 Synonyms": "withmartian/cs3_dataset_synonyms",
    "CS4 Synonyms": "withmartian/cs4_dataset_synonyms",
}

COLUMNS = ["create_statement", "english_prompt", "sql_statement"]

def load_preview(dataset_name):
    try:
        ds = load_dataset(DATASETS[dataset_name], split="train")
        df = pd.DataFrame(ds).head(500)
        if all(col in df.columns for col in COLUMNS):
            df = df[COLUMNS]
        # Add index column
        df.insert(0, 'index', range(len(df)))
        return df
    except Exception as e:
        return pd.DataFrame({"Error": [str(e)]})

def filter_dataframe(df, search_query):
    if not search_query or df.empty or "Error" in df.columns:
        return df
    
    mask = df.astype(str).apply(
        lambda row: row.str.contains(search_query, case=False, na=False).any(), 
        axis=1
    )
    return df[mask]

def dataset_viewer(shared_instruction, shared_schema):
    gr.HTML("""
        <div style="text-align: center; padding: 2rem 1.5rem; background: linear-gradient(135deg, #2A2A2A 0%, #3A3A3A 100%); border-radius: 16px; margin-bottom: 1.5rem; box-shadow: 0 4px 12px rgba(0,0,0,0.3);">
            <h2 style="font-size: 2rem; font-weight: 700; margin-bottom: 0.5rem; color: #FF6B4A;">Dataset Explorer</h2>
            <p style="font-size: 1rem; opacity: 0.9; line-height: 1.6; color: #D0D0D0;">
                Browse, search, and explore TinySQL datasets
            </p>
        </div>
    """)
    
    gr.HTML("""
        <div style="background: linear-gradient(135deg, #2A2A2A 0%, #3A3A3A 100%); border-radius: 12px; padding: 1.5rem; margin: 1rem 0; border-left: 4px solid #FF6B4A;">
            <p style="color: #D0D0D0; margin: 0; line-height: 1.6;">
                <strong style="color: #FF6B4A;">Quick Start:</strong> Select a dataset, click Load Dataset, then use search to filter. Pick any row and send it to the Model Demo tab.
            </p>
        </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Controls")
            
            dataset_dropdown = gr.Dropdown(
                choices=list(DATASETS.keys()),
                value="CS1",
                label="Choose Dataset",
                info="Select complexity level"
            )
            
            # Simpler dataset guide - no colors, no beginner/intermediate
            gr.HTML("""
                <div style="background: #2A2A2A; border-radius: 12px; padding: 1.5rem; margin: 1.5rem 0; border: 1px solid #3A3A3A;">
                    <h4 style="color: #FF6B4A; font-size: 1rem; margin: 0 0 1.25rem 0; font-weight: 700; border-bottom: 2px solid #3A3A3A; padding-bottom: 0.75rem;">Dataset Complexity Levels</h4>
                    
                    <div style="color: #D0D0D0; font-size: 0.9rem; line-height: 2;">
                        <div><strong>CS1:</strong> Basic SELECT-FROM queries</div>
                        <div><strong>CS2:</strong> Adds ORDER BY clauses</div>
                        <div><strong>CS3:</strong> Aggregations (COUNT, SUM, AVG)</div>
                        <div><strong>CS4:</strong> Adds WHERE filters</div>
                        <div><strong>CS5:</strong> Multi-table JOINs</div>
                    </div>
                    
                    <div style="margin-top: 1.5rem; padding-top: 1.25rem; border-top: 1px solid #3A3A3A;">
                        <div style="color: #FF6B4A; font-weight: 600; font-size: 0.9rem; margin-bottom: 0.5rem;">Synonym Variants</div>
                        <div style="color: #999; font-size: 0.85rem; line-height: 1.6;">Natural language variations with semantic mappings</div>
                    </div>
                </div>
            """)
            
            load_btn = gr.Button("Load Dataset", variant="primary", size="lg")
            
            gr.Markdown("### Test Example")
            row_selector = gr.Number(
                label="Row Number",
                value=0,
                minimum=0,
                precision=0,
                info="Pick a row to test"
            )
            
            send_to_model_btn = gr.Button("Run in Model Demo", variant="primary")
        
        with gr.Column(scale=3):
            gr.Markdown("### Dataset Preview")
            
            search_box = gr.Textbox(
                label="Search",
                placeholder="Search across all columns...",
                lines=1
            )
            
            # HuggingFace-style table with row index on hover
            gr.HTML("""
                <style>
                /* True HuggingFace-style table - NO "Results" label */
                .dataframe-container label {
                    display: none !important;
                }
                
                .dataframe-container {
                    border-radius: 8px !important;
                    overflow: hidden !important;
                    border: 1px solid #374151 !important;
                }

                .dataframe table {
                    border-collapse: collapse !important;
                    width: 100% !important;
                    font-size: 0.875rem !important;
                }

                .dataframe thead {
                    background: #1f2937 !important;
                }

                .dataframe thead th {
                    color: #9ca3af !important;
                    font-weight: 600 !important;
                    text-align: left !important;
                    padding: 0.75rem 1rem !important;
                    border-bottom: 1px solid #374151 !important;
                    font-size: 0.75rem !important;
                    text-transform: uppercase !important;
                    letter-spacing: 0.05em !important;
                }

                .dataframe tbody tr {
                    background: #111827 !important;
                    border-bottom: 1px solid #1f2937 !important;
                    transition: all 0.15s ease !important;
                    position: relative !important;
                }

                .dataframe tbody tr:hover {
                    background: #1f2937 !important;
                    box-shadow: 0 2px 8px rgba(255, 107, 74, 0.1) !important;
                }

                .dataframe tbody tr:hover::before {
                    content: "Row " attr(data-row-index);
                    position: absolute;
                    left: -60px;
                    top: 50%;
                    transform: translateY(-50%);
                    background: #FF6B4A;
                    color: white;
                    padding: 0.25rem 0.5rem;
                    border-radius: 4px;
                    font-size: 0.75rem;
                    font-weight: 600;
                    white-space: nowrap;
                    opacity: 0.9;
                }

                .dataframe tbody td {
                    padding: 0.75rem 1rem !important;
                    color: #d1d5db !important;
                    font-size: 0.875rem !important;
                    line-height: 1.5 !important;
                    max-width: 400px !important;
                    overflow: hidden !important;
                    text-overflow: ellipsis !important;
                }

                .dataframe tbody tr:last-child {
                    border-bottom: none !important;
                }
                
                /* Hide index column but keep it for reference */
                .dataframe tbody td:first-child,
                .dataframe thead th:first-child {
                    width: 0 !important;
                    padding: 0 !important;
                    opacity: 0 !important;
                    position: absolute !important;
                }
                </style>
            """)
            
            df_display = gr.Dataframe(
                headers=["index"] + COLUMNS,
                datatype=["number", "str", "str", "str"],
                interactive=False,
                wrap=True,
                elem_classes="dataframe-container"
            )
            
            stats_display = gr.Markdown("Click **Load Dataset** to begin exploring")
    
    df_state = gr.State(value=pd.DataFrame())
    
    def load_and_display(dataset_name):
        df = load_preview(dataset_name)
        if "Error" in df.columns:
            return df, df, "Error loading dataset"
        stats = f"**Loaded {len(df)} rows** • Columns: {', '.join(COLUMNS)}"
        return df, df, stats
    
    load_btn.click(
        fn=load_and_display,
        inputs=dataset_dropdown,
        outputs=[df_state, df_display, stats_display]
    )
    
    def search_and_display(df, query):
        if df.empty:
            return df, "Load a dataset first"
        
        filtered_df = filter_dataframe(df, query)
        stats = f"**Showing {len(filtered_df)} of {len(df)} rows**"
        if query:
            stats += f" • Search: '{query}'"
        return filtered_df, stats
    
    search_box.change(
        fn=search_and_display,
        inputs=[df_state, search_box],
        outputs=[df_display, stats_display]
    )
    
    def send_to_model(df, row_num):
        if df.empty or row_num >= len(df):
            return "", "", "Invalid row or no data loaded"
        
        row = df.iloc[int(row_num)]
        instruction = row['english_prompt'] if 'english_prompt' in row else ""
        schema = row['create_statement'] if 'create_statement' in row else ""
        
        return instruction, schema, f"**Row {row_num} loaded!** Switch to Model Demo tab"
    
    send_to_model_btn.click(
        fn=send_to_model,
        inputs=[df_state, row_selector],
        outputs=[shared_instruction, shared_schema, stats_display]
    )
    
    return {'df_state': df_state, 'df_display': df_display}