import streamlit as st
from pipeline import text_to_sql

st.title("SQLCoder Text-to-SQL App")
st.write("Powered by defog/sqlcoder-7b-2 🚀")

# Sample queries for user guidance
st.sidebar.header("Sample Queries")
sample_queries = [
    "List 11 names of ships type schooner",
    "Show me the 5 oldest ships",
    "What are the different types of vessels?",
    "Count the number of ships by type",
    "Show ships built after 1900"
]

selected_sample = st.sidebar.selectbox("Choose a sample query:", [""] + sample_queries)

# Main input
nl_query = st.text_input(
    "Enter your natural language query:", 
    value=selected_sample if selected_sample else "List 11 names of ships type schooner",
    help="Ask questions about your database in plain English"
)

if st.button("🔄 Generate & Execute SQL"):
    if nl_query.strip():
        with st.spinner("Generating SQL and executing query..."):
            try:
                sql, results = text_to_sql(nl_query)
                
                # Display results
                st.success("Query executed successfully!")
                
                # Show generated SQL
                st.subheader("Generated SQL:")
                st.code(sql, language="sql")
                
                # Show results
                st.subheader("Results:")
                if results:
                    # Convert results to a more readable format
                    if isinstance(results[0], tuple):
                        # If results are tuples, display as table
                        st.write(f"Found {len(results)} rows:")
                        for i, row in enumerate(results[:50]):  # Show first 50 rows
                            st.write(f"Row {i+1}: {row}")
                        if len(results) > 50:
                            st.info(f"Showing first 50 rows out of {len(results)} total results.")
                    else:
                        st.write(results)
                else:
                    st.info("Query executed successfully but returned no results.")
                    
            except Exception as e:
                st.error(f"Error: {str(e)}")
                st.write("Please try rephrasing your query or check if the requested data exists in the database.")
    else:
        st.warning("Please enter a query to proceed.")