src/streamlit_app.py

import streamlit as st
import requests
import uuid
import time
import os

# --- CONFIGURATION ---
# Your FastAPI server URL on Hugging Face
API_URL = "https://lightrt-text2sql-backend.hf.space" 

st.set_page_config(page_title="Text2SQL Agent", page_icon="🤖", layout="centered")

# --- HUGGING FACE FIREWALL BYPASS (CHROME SPOOFING) ---
HF_TOKEN = os.getenv("HF_TOKEN")

HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
    "Accept": "application/json",
    "Content-Type": "application/json"
}

if HF_TOKEN:
    HEADERS["Authorization"] = f"Bearer {HF_TOKEN}"

# (Optional) A quick UI check so you know the token loaded in the cloud
if not HF_TOKEN:
    st.sidebar.error("⚠️ HF_TOKEN missing from Space Secrets!")
else:
    st.sidebar.success("🔑 Token Loaded! Ready for cloud connection.")


# --- SESSION STATE INITIALIZATION ---
if "thread_id" not in st.session_state:
    st.session_state.thread_id = str(uuid.uuid4())
if "user_id" not in st.session_state:
    st.session_state.user_id = "tenant_" + str(uuid.uuid4())[:8]
if "is_db_connected" not in st.session_state:
    st.session_state.is_db_connected = False
if "connection_url" not in st.session_state:
    st.session_state.connection_url = ""
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# --- SIDEBAR: DATABASE CONNECTION ---
with st.sidebar:
    st.header("⚙️ Database Setup")
    
    db_input = st.text_input(
        "Enter Database URL:",
        disabled=st.session_state.is_db_connected
    )
    
    if not st.session_state.is_db_connected:
        if st.button("Connect & Initialize", type="primary", use_container_width=True):
            if not db_input:
                st.error("Please enter a valid URL.")
            else:
                with st.spinner("Building embeddings and initializing agent..."):
                    try:
                        payload = {"connection_url": db_input, "user_id": st.session_state.user_id}
                        
                        # Added headers=HEADERS here!
                        response = requests.post(f"{API_URL}/upload_url", json=payload, headers=HEADERS)
                        
                        if response.status_code == 200:
                            st.session_state.is_db_connected = True
                            st.session_state.connection_url = db_input
                            
                            time.sleep(15) 
                            
                            st.success("Database connected securely!")
                            st.rerun()
                        elif response.status_code == 429:
                            st.error("🚨 Rate Limited! Waiting for HF firewall cooldown.")
                        else:
                            st.error(f"Failed to connect: {response.text}")
                    except requests.exceptions.ConnectionError:
                        st.error("🚨 Cannot connect to backend. Is FastAPI running?")
    else:
        st.success("✅ Connected to Database")
        st.caption(f"URL: {st.session_state.connection_url}")
        
        if st.button("Disconnect & Reset", use_container_width=True):
            st.session_state.clear()
            st.rerun()

# --- MAIN CHAT INTERFACE ---
st.title("🗣️ Text2SQL Agent")

if not st.session_state.is_db_connected:
    st.info("👈 Please connect your database in the sidebar to begin analyzing data.")
else:
    for msg in st.session_state.chat_history:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])

    if user_query := st.chat_input("Ask a question about your data..."):
        
        st.session_state.chat_history.append({"role": "user", "content": user_query})
        with st.chat_message("user"):
            st.markdown(user_query)

        with st.chat_message("assistant"):
            with st.spinner("Analyzing schema and generating SQL..."):
                try:
                    payload = {
                        "message": user_query,
                        "thread_id": st.session_state.thread_id,
                        "user_id": st.session_state.user_id,
                        "connection_url": st.session_state.connection_url
                    }
                    
                    # Added headers=HEADERS here!
                    response = requests.post(f"{API_URL}/chat", json=payload, headers=HEADERS)
                    
                    if response.status_code == 200:
                        answer = response.json().get("response", "No response found.")
                        st.markdown(answer)
                        st.session_state.chat_history.append({"role": "assistant", "content": answer})
                    elif response.status_code == 429:
                        st.error("🚨 Rate Limited by Hugging Face firewall.")
                    else:
                        st.error(f"Agent Error: {response.text}")
                        
                except requests.exceptions.ConnectionError:
                    st.error("🚨 Connection dropped. Ensure FastAPI is running.")