streamlit_app.py

"""
Streamlit UI for AI-Assisted Database Design and CSV Import
Connects to FastAPI backend
"""
import streamlit as st
import requests
import pandas as pd
import io
import os

# Configure page
st.set_page_config(
    page_title="AI CSV Import",
    page_icon="📊",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Backend API configuration
API_URL = os.getenv("API_URL", "http://localhost:8000")
HEALTH_CHECK_URL = f"{API_URL}/health"

# Custom CSS for better UI
st.markdown("""
<style>
    .main-header {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        padding: 2rem;
        border-radius: 10px;
        color: white;
        margin-bottom: 2rem;
    }
    .success-box {
        background-color: #d4edda;
        border: 1px solid #c3e6cb;
        color: #155724;
        padding: 1rem;
        border-radius: 5px;
        margin: 1rem 0;
    }
    .error-box {
        background-color: #f8d7da;
        border: 1px solid #f5c6cb;
        color: #721c24;
        padding: 1rem;
        border-radius: 5px;
        margin: 1rem 0;
    }
    .info-box {
        background-color: #d1ecf1;
        border: 1px solid #bee5eb;
        color: #0c5460;
        padding: 1rem;
        border-radius: 5px;
        margin: 1rem 0;
    }
</style>
""", unsafe_allow_html=True)

# Session state
if "backend_connected" not in st.session_state:
    st.session_state.backend_connected = False
if "api_error" not in st.session_state:
    st.session_state.api_error = None


def check_backend():
    """Check if backend API is available"""
    try:
        response = requests.get(HEALTH_CHECK_URL, timeout=5)
        if response.status_code in [200, 503]:
            return True, response.json().get("message", "Unknown")
    except Exception as e:
        return False, str(e)
    return False, "Backend unavailable"


def get_collections():
    """Get list of tables from backend"""
    try:
        response = requests.get(f"{API_URL}/api/collections", timeout=10)
        if response.status_code == 200:
            data = response.json()
            return data.get("tables", [])
        elif response.status_code == 503:
            return None  # Database still initializing
        else:
            st.error(f"Error: {response.json().get('detail', 'Unknown error')}")
            return []
    except requests.exceptions.ConnectionError:
        return None
    except Exception as e:
        st.error(f"Error fetching collections: {str(e)}")
        return []


def get_table_data(table_name: str, limit: int = 100):
    """Get data from a specific table"""
    try:
        response = requests.get(
            f"{API_URL}/api/data/{table_name}",
            params={"limit": limit, "skip": 0},
            timeout=30
        )
        if response.status_code == 200:
            return response.json()
        else:
            try:
                detail = response.json().get("detail", "Unknown error")
                st.error(f"Error fetching table data: {detail}")
            except Exception:
                st.error(f"Error fetching table data: HTTP {response.status_code}")
            return None
    except Exception as e:
        st.error(f"Error fetching table data: {str(e)}")
        return None


# Header
st.markdown("""
<div class="main-header">
    <h1>📊 AI CSV Import Database</h1>
    <p>Upload CSV files with intelligent column mapping powered by Gemini AI</p>
</div>
""", unsafe_allow_html=True)

# Sidebar - Backend Status
with st.sidebar:
    st.header("🔧 Status")
    
    connected, message = check_backend()
    if connected:
        st.success(f"✅ Backend Connected\n{message}")
        st.session_state.backend_connected = True
    else:
        st.warning(f"⚠️ Backend: {message}\n\nPlease wait, the database may be initializing...")
        st.session_state.backend_connected = False
    
    st.divider()
    st.markdown("### Navigation")
    page = st.radio("Select Page", ["📤 Upload CSV", "📋 View Data", "🗂️ Manage Schemas"])

# Main content
if not st.session_state.backend_connected:
    st.info("""
    ⏳ **Initializing...**
    
    The backend is still connecting to the database. This typically takes 1-2 minutes on first startup.
    
    Please refresh this page in a moment. Your data will be persisted in Supabase PostgreSQL.
    """)
    st.stop()

# Page routing
if page == "📤 Upload CSV":
    st.header("Upload CSV File")
    
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.markdown("""
        Upload your CSV file and let AI automatically:
        - Map column names to database-friendly formats
        - Infer correct data types
        - Create PostgreSQL table
        """)
    
    uploaded_file = st.file_uploader("Choose CSV file", type="csv")
    
    if uploaded_file is not None:
        # Show preview
        try:
            preview_df = pd.read_csv(io.BytesIO(uploaded_file.getvalue()), nrows=5)
            st.subheader("📋 Preview (first 5 rows)")
            st.dataframe(preview_df, use_container_width=True)
        except Exception as e:
            st.error(f"Error reading CSV: {str(e)}")
            st.stop()
        
        # Configuration
        st.subheader("⚙️ Configuration")
        col1, col2, col3 = st.columns(3)
        
        with col1:
            table_name = st.text_input("Table Name", value=uploaded_file.name.replace(".csv", "").lower())
        
        with col2:
            schema_name = st.text_input("Schema Name (optional)", value="", help="Save this mapping for reuse")
        
        with col3:
            missing_strategy = st.selectbox(
                "Missing Values Strategy",
                ["drop", "fill_mean", "fill_empty"],
                help="How to handle empty values"
            )
        
        # Upload button
        if st.button("🚀 Upload & Process", use_container_width=True, type="primary"):
            with st.spinner("Processing CSV with AI mapping..."):
                try:
                    # Prepare files and data
                    files = {"file": (uploaded_file.name, uploaded_file.getvalue(), "text/csv")}
                    data = {
                        "table_name": table_name,
                        "schema_name": schema_name if schema_name else None,
                        "missing_value_strategy": missing_strategy
                    }
                    
                    response = requests.post(
                        f"{API_URL}/api/upload-csv",
                        files=files,
                        data=data,
                        timeout=30
                    )
                    
                    if response.status_code == 200:
                        payload = response.json()
                        result = payload.get("data", {})
                        st.markdown("""
                        <div class="success-box">
                        ✅ CSV uploaded successfully!
                        </div>
                        """, unsafe_allow_html=True)
                        
                        col1, col2, col3 = st.columns(3)
                        with col1:
                            st.metric("Rows Inserted", result.get("records_inserted", 0))
                        with col2:
                            st.metric("Columns Mapped", len(result.get("mapping", {})))
                        with col3:
                            st.metric("Table", result.get("table_name", "N/A"))
                        
                        # Show column mapping
                        st.subheader("📊 Column Mapping")
                        mapping = result.get("mapping", {})
                        if mapping:
                            mapping_df = pd.DataFrame([
                                {"Original": k, "Mapped": v.get("name"), "Type": v.get("type")}
                                for k, v in mapping.items()
                            ])
                            st.dataframe(mapping_df, use_container_width=True)
                    
                    elif response.status_code == 503:
                        st.markdown("""
                        <div class="error-box">
                        ⏳ Database is initializing. Please try again in a moment.
                        </div>
                        """, unsafe_allow_html=True)
                    else:
                        error_msg = response.json().get("detail", "Unknown error")
                        st.markdown(f"""
                        <div class="error-box">
                        ❌ Error: {error_msg}
                        </div>
                        """, unsafe_allow_html=True)
                
                except requests.exceptions.Timeout:
                    st.error("⏱️ Request timeout. Server might be busy.")
                except Exception as e:
                    st.error(f"❌ Error: {str(e)}")

elif page == "📋 View Data":
    st.header("View Table Data")
    
    collections = get_collections()
    
    if collections is None:
        st.info("⏳ Database is still initializing. Please try again in a moment.")
    elif not collections:
        st.info("📭 No tables found. Upload a CSV file first!")
    else:
        selected_table = st.selectbox("Select Table", collections)
        
        col1, col2 = st.columns(2)
        with col1:
            limit = st.slider("Rows to Display", 10, 1000, 100)
        
        if st.button("📥 Load Data", use_container_width=True, type="primary"):
            with st.spinner("Loading data..."):
                data = get_table_data(selected_table, limit)
                
                if data:
                    st.subheader(f"Table: {selected_table}")
                    col1, col2, col3 = st.columns(3)
                    with col1:
                        st.metric("Total Rows", data.get("total_count", 0))
                    with col2:
                        st.metric("Displayed", data.get("returned_count", 0))
                    with col3:
                        st.metric("Limit", limit)
                    
                    # Display data
                    if data.get("data"):
                        df = pd.DataFrame(data["data"])
                        st.dataframe(df, use_container_width=True)
                        
                        # Download option
                        csv = df.to_csv(index=False)
                        st.download_button(
                            label="📥 Download as CSV",
                            data=csv,
                            file_name=f"{selected_table}.csv",
                            mime="text/csv"
                        )
                    else:
                        st.info("📭 No data in this table")

elif page == "🗂️ Manage Schemas":
    st.header("Manage Saved Schemas")
    
    if st.button("🔄 Load Schemas", use_container_width=True, type="primary"):
        with st.spinner("Loading schemas..."):
            try:
                response = requests.get(f"{API_URL}/api/schemas", timeout=10)
                
                if response.status_code == 200:
                    data = response.json()
                    schemas = data.get("schemas", [])
                    
                    if not schemas:
                        st.info("📭 No saved schemas yet. Upload a CSV to save a schema!")
                    else:
                        st.success(f"Found {len(schemas)} saved schemas")
                        
                        for schema in schemas:
                            with st.expander(f"📋 {schema.get('name', 'Unnamed Schema')}"):
                                st.json(schema)
                
                elif response.status_code == 503:
                    st.info("⏳ Database is initializing. Please try again soon.")
                else:
                    st.error(f"Error: {response.json().get('detail', 'Unknown error')}")
            
            except Exception as e:
                st.error(f"Error: {str(e)}")

# Footer
st.divider()
st.markdown("""
<div style="text-align: center; color: gray; font-size: 0.8rem;">
    <p>AI-Assisted Database Design and CSV Import | Powered by Streamlit & FastAPI</p>
    <p>API: <code>%s</code></p>
</div>
""" % API_URL, unsafe_allow_html=True)