import streamlit as st
import requests
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Get Lambda URL from environment variable
LAMBDA_URL = os.getenv("LAMBDA_URL") or "http://localhost:8000/chats/messages"
if not LAMBDA_URL:
    st.error("LAMBDA_URL environment variable is not set. Please check your .env file.")
    st.stop()

st.set_page_config(page_title="Ask Redshift", page_icon="🤖")

st.title("🧠 Ask your Redshift Database Anything...")

# Initialize conversation_id in session state if not exists
if "conversation_id" not in st.session_state:
    st.session_state.conversation_id = None

# Add system prompt input
system_prompt = st.text_area(
    "System Prompt (Optional)",
    help="Customize the AI's behavior and instructions. Leave empty to use default behavior.",
    height=100
)

question = st.text_input(
    "Enter your question about your data (e.g., 'What is the average cost per lead for March?')"
)

if st.button("Run Query"):
    if question:
        with st.spinner("Querying..."):
            # Prepare payload
            payload = {"question": question}
            if st.session_state.conversation_id:
                payload["conversation_id"] = st.session_state.conversation_id
            if system_prompt:
                payload["system_prompt"] = system_prompt

            # Call the API
            response = requests.post(
                LAMBDA_URL,
                json=payload,
            )
            response.raise_for_status()
            response = response.json()

            # Store conversation_id for future requests
            if "conversation_id" in response:
                st.session_state.conversation_id = response["conversation_id"]

        st.success(response["answer"])
    else:
        st.error("Please enter a question.")