from dotenv import load_dotenv
load_dotenv()

import streamlit as st
import requests
import json

st.set_page_config(page_title="LangGraph Agent UI", layout="centered")
st.title("AI Chatbot Agents")
st.write("Create and Interact with the AI Agents!")

system_prompt = st.text_area(
    "Define your AI Agent: ",
    height=70,
    placeholder="Type your system prompt here...",
    value="You are a helpful AI assistant."
)

MODEL_NAMES_GROQ = ["llama-3.3-70b-versatile", "mixtral-8x7b-32768"]
selected_model = st.selectbox("Select Groq Model:", MODEL_NAMES_GROQ)

allow_web_search = st.checkbox("Allow Web Search")

user_query = st.text_area(
    "Enter your query: ",
    height=150,
    placeholder="Ask Anything!"
)

API_URL = "http://127.0.0.1:9999/chat"

if st.button("Ask Agent!"):
    if not user_query.strip():
        st.error("Please enter a query!")
    else:
        try:
            with st.spinner("Getting response..."):
                payload = {
                    "model_name": selected_model,
                    "model_provider": "Groq",
                    "system_prompt": system_prompt,
                    "messages": [user_query],
                    "allow_search": allow_web_search
                }

                # Add debug output
                st.write("Sending request with payload:", json.dumps(payload, indent=2))
                
                response = requests.post(API_URL, json=payload, timeout=120)
                
                if response.status_code == 200:
                    response_data = response.json()
                    st.subheader("Agent Response")
                    st.markdown(f"{response_data['response']}")
                else:
                    st.error(f"Server error: {response.text}")
                
        except requests.exceptions.ConnectionError:
            st.error("Could not connect to the backend server. Please make sure it's running on port 9999.")
        except Exception as e:
            st.error(f"An error occurred: {str(e)}")