import streamlit as st
from openai import OpenAI
import sqlite3
import pandas as pd  # Add to top of file


# App UI
st.set_page_config(page_title="Zero SQL", layout="wide")
st.title("Zero SQL - Natural Language to SQL Query")


# Database initialization with caching
@st.cache_resource
def initialize_database():
    conn = sqlite3.connect('database.db')
    cursor = conn.cursor()

    # Create tables
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS Produkte (
            ProduktID INTEGER PRIMARY KEY AUTOINCREMENT,
            Produktname TEXT NOT NULL,
            Preis REAL NOT NULL
        )
    ''')
    
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS Bestellungen (
            BestellungID INTEGER PRIMARY KEY AUTOINCREMENT,
            ProduktID INTEGER NOT NULL,
            Menge INTEGER NOT NULL,
            Bestelldatum TEXT NOT NULL,
            Person TEXT NOT NULL,
            FOREIGN KEY (ProduktID) REFERENCES Produkte(ProduktID)
        )
    ''')

    # Insert sample products
    cursor.execute("SELECT COUNT(*) FROM Produkte")
    if cursor.fetchone()[0] == 0:
        products = [
            ('Laptop', 999.99),
            ('Smartphone', 699.99),
            ('Tablet', 399.99)
        ]
        cursor.executemany("INSERT INTO Produkte (Produktname, Preis) VALUES (?, ?)", products)

    # Insert sample orders
    cursor.execute("SELECT COUNT(*) FROM Bestellungen")
    if cursor.fetchone()[0] == 0:
        orders = [
            (1, 2, '2024-10-20', 'Max Mustermann'),
            (2, 1, '2024-10-21', 'Erika Musterfrau'),
            (3, 3, '2024-10-22', 'Hans Meier')
        ]
        cursor.executemany("INSERT INTO Bestellungen (ProduktID, Menge, Bestelldatum, Person) VALUES (?, ?, ?, ?)", orders)

    conn.commit()
    conn.close()

# Initialize the database
initialize_database()


# Sidebar for API key
with st.sidebar:
    st.header("Configuration")
    api_key = st.text_input("OpenAI API Key", type="password")
    st.markdown("---")
    st.markdown("**Sample Questions:**")
    st.markdown("- Show total sales per product")
    st.markdown("- List orders from Max Mustermann")
    st.markdown("- Find most popular product by quantity")

# Main form
with st.form("query_form"):
    user_input = st.text_area(
        "Enter your data request in natural language:",
        placeholder="e.g. Show all orders over €500",
        height=100
    )
    submitted = st.form_submit_button("🚀 Generate Query")

if submitted:
    if not api_key:
        st.error("🔑 API key is required!")
    elif not user_input:
        st.error("📝 Please enter your data request!")
    else:
        try:
            client = OpenAI(api_key=api_key)
            
            # System prompt with schema
            system_context = """You are a SQL expert. Given these SQL tables:
            CREATE TABLE Produkte (
                ProduktID INTEGER PRIMARY KEY,
                Produktname TEXT NOT NULL,
                Preis REAL NOT NULL
            );
            CREATE TABLE Bestellungen (
                BestellungID INTEGER PRIMARY KEY,
                ProduktID INTEGER,
                Menge INTEGER,
                Bestelldatum TEXT,
                Person TEXT,
                FOREIGN KEY (ProduktID) REFERENCES Produkte(ProduktID)
            );
            Generate ONLY the raw SQL query for the user's request. 
            NEVER use markdown code blocks or any formatting. 
            ONLY output the pure SQL statement."""
        
           # Generate SQL query
            response = client.chat.completions.create(
                model="gpt-4o",
                messages=[
                    {"role": "system", "content": system_context},
                    {"role": "user", "content": user_input}
                ],
                temperature=0.3
            )
            
            # Clean up the response
            sql_query = response.choices[0].message.content.strip()
            sql_query = sql_query.replace("```sql", "").replace("```", "").strip()  # Add this line
            
            # Execute and display results
            with sqlite3.connect('database.db') as conn:
                cursor = conn.cursor()
                cursor.execute(sql_query)
                results = cursor.fetchall()
                column_names = [desc[0] for desc in cursor.description]

            st.subheader("Generated SQL")
            st.code(sql_query, language="sql")

            st.subheader("Results")
            if results:
                df = pd.DataFrame(results, columns=column_names)
                st.dataframe(
                    data=df,
                    use_container_width=True,
                    hide_index=True
                )
            else:
                st.info("No results found", icon="ℹ️")

        except sqlite3.Error as e:
            st.error(f"🚨 SQL Error: {str(e)}")
        except Exception as e:
            st.error(f"💥 Unexpected Error: {str(e)}")