import streamlit as st
from pathlib import Path
from langchain.llms.openai import OpenAI
from langchain.agents import create_sql_agent
from langchain.sql_database import SQLDatabase
from langchain.agents.agent_types import AgentType
from langchain.callbacks import StreamlitCallbackHandler
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
from sqlalchemy import create_engine
import sqlite3
import os
from langchain_openai import ChatOpenAI
os.environ["OPENAI_API_KEY"] = st.secrets["OPENAI_API_KEY"]
st.set_page_config(page_title="Chat with Protected Areas Database", page_icon="🦜")
st.title("🦜 LangChain: Chat with Protected Areas Database")
db_uri = "duckdb:///pad.duckdb"

# User inputs
radio_opt = ["US Protected Areas v3"]
selected_opt = st.sidebar.radio(label="Choose suitable option", options=radio_opt)

# Setup agent
@st.cache_resource(ttl="2h")
def configure_db(db_uri):        
    return SQLDatabase.from_uri(database_uri=db_uri, view_support=True)
db = configure_db(db_uri)

llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
agent = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=True)

def handle_user_input(user_query):
    with history:
        st.session_state.messages.append({"role": "user", "content": user_query})
        st.chat_message("user").write(user_query)

        with st.chat_message("assistant"):
            st_cb = StreamlitCallbackHandler(st.container())
            response = agent.run(user_query, callbacks=[st_cb])
            st.session_state.messages.append({"role": "assistant", "content": response})
            st.write(response)


if "messages" not in st.session_state:
    st.session_state["messages"] = []

main = st.container()
with main:
    history = st.container(height=800)
    #with history:
    #    for msg in st.session_state.messages:
    #        st.chat_message(msg["role"]).write(msg["content"])  
    if user_query := st.chat_input(placeholder="Ask me about US Protected areas!"):
        handle_user_input(user_query)
    
    st.markdown("\n") #add some space for iphone users


#if "messages" not in st.session_state or st.sidebar.button("Clear message history"):
#    st.session_state["messages"] = [{"role": "assistant", "content": "How can I help you?"}]



EXAMPLE_PROMPTS = ["What is the total area in each GAP_Sts category in the fee table?",
                   "List the name of each table in the database",
                   "How much BLM land (BLM is a Mang_Name in the fee table) is in each GAP_Sts category?",
                   "Federal agencies are identified as 'FED' in the Mang_Type column in the 'combined' data table. The Mang_Name column indicates the different agencies. The full name of each agency is given in the agency_name table. Which federal agencies, by full name, manage the greatest area of GAP_Sts 1 or 2 land?"]

with st.sidebar:
    with st.container():
        st.title("Examples")
        for prompt in EXAMPLE_PROMPTS:
            st.button(prompt, args=(prompt,), on_click=handle_user_input)