import os
import gradio as gr
from sqlalchemy import text
from smolagents import tool, CodeAgent, HfApiModel
import spaces
import pandas as pd
from database import (
    engine,
    create_dynamic_table,
    clear_database,
    insert_rows_into_table,
    get_table_schema
)

def process_sql_file(file_path):
    """
    Process an SQL file and execute its contents.
    """
    try:
        # Read the SQL file
        with open(file_path, 'r') as file:
            sql_content = file.read()
            
        # Replace AUTO_INCREMENT with AUTOINCREMENT for SQLite compatibility
        sql_content = sql_content.replace('AUTO_INCREMENT', 'AUTOINCREMENT')
            
        # Split into individual statements
        statements = [stmt.strip() for stmt in sql_content.split(';') if stmt.strip()]
        
        # Clear existing database
        clear_database()
        
        # Execute each statement
        with engine.begin() as conn:
            for statement in statements:
                if statement.strip():
                    conn.execute(text(statement))
                    
        return True, "SQL file successfully executed! Click 'Continue' to proceed to query interface..."
        
    except Exception as e:
        return False, f"Error processing SQL file: {str(e)}"

# ... (other functions remain the same) ...

# Create the Gradio interface
with gr.Blocks() as demo:
    # Upload Interface Components
    with gr.Group() as upload_group:
        gr.Markdown("""
        # Data Query Interface

        Upload your data file to begin.
        
        ### Supported File Types:
        - SQL (.sql): SQL file containing CREATE TABLE and INSERT statements
        - CSV (.csv): CSV file with headers that will be automatically converted to a table
        
        ### CSV Requirements:
        - Must include headers
        - First column will be used as the primary key
        - Column types will be automatically detected
        
        ### SQL Requirements:
        - Must contain valid SQL statements
        - Statements must be separated by semicolons
        - Should include CREATE TABLE and data insertion statements
        """)
        
        file_input = gr.File(
            label="Upload Data File",
            file_types=[".csv", ".sql"],
            type="filepath"
        )
        upload_status = gr.Textbox(label="Status", interactive=False)
        continue_btn = gr.Button("Continue", visible=False)

    # Query Interface Components
    with gr.Group(visible=False) as query_group:
        gr.Markdown("## Data Query Interface")
        
        # Data Display Section
        gr.Markdown("### Current Data")
        data_table = gr.Dataframe(
            value=get_data_table(),
            label="Data Table",
            interactive=False
        )
        
        schema_display = gr.Markdown(value="Loading schema...")
        
        # Query Section
        with gr.Row():
            with gr.Column():
                user_input = gr.Textbox(
                    label="Ask a question about the data",
                    placeholder="Enter your question here..."
                )
                query_output = gr.Textbox(
                    label="Result",
                    interactive=False
                )
        
        with gr.Row():
            refresh_table_btn = gr.Button("Refresh Table")
            refresh_schema_btn = gr.Button("Refresh Schema")
            back_btn = gr.Button("Upload New File")

    def handle_upload(file):
        success, message = process_uploaded_file(file)
        if success:
            df = get_data_table()
            schema = get_table_schema()
            return {
                upload_status: message,
                continue_btn: gr.update(visible=True),
                data_table: df,
                schema_display: f"### Current Schema:\n```\n{schema}\n```" if schema else "No schema available"
            }
        return {
            upload_status: message,
            continue_btn: gr.update(visible=False)
        }

    def switch_to_query():
        df = get_data_table()
        schema = get_table_schema()
        return {
            upload_group: gr.update(visible=False),
            query_group: gr.update(visible=True),
            data_table: df,
            schema_display: f"### Current Schema:\n```\n{schema}\n```" if schema else "No schema available"
        }

    def switch_to_upload():
        return {
            upload_group: gr.update(visible=True),
            query_group: gr.update(visible=False),
            continue_btn: gr.update(visible=False),
            upload_status: gr.update(value="")
        }

    # Event handlers
    file_input.upload(
        fn=handle_upload,
        outputs=[
            upload_status,
            continue_btn,
            data_table,
            schema_display
        ]
    )

    continue_btn.click(
        fn=switch_to_query,
        outputs=[
            upload_group,
            query_group,
            data_table,
            schema_display
        ]
    )

    back_btn.click(
        fn=switch_to_upload,
        outputs=[
            upload_group,
            query_group,
            continue_btn,
            upload_status
        ]
    )

    user_input.change(
        fn=query_sql,
        inputs=[user_input],
        outputs=[query_output]
    )
    
    refresh_table_btn.click(
        fn=get_data_table,
        outputs=[data_table]
    )
    
    refresh_schema_btn.click(
        fn=update_schema,
        outputs=[schema_display]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)