import streamlit as st
from openai import OpenAI
import duckdb

EXAMPLE_PROMPTS = [
    "What are the unique values of Mang Type?",
    "What is the mean species richness of each GAP_Sts?",
    "STAT, LOC, FED, and DIST Mang types are all 'Public' land types. PVT and NGO are 'Private' land types. What is the average richness in public, private, and other land types?",
]

TITLE = "Protected Lands SQL Assistant"
# DESCRIPTION = """[DBRX Instruct](https://huggingface.co/databricks/dbrx-instruct) is a mixture-of-experts (MoE) large language model trained by the Mosaic Research team at Databricks. Users can interact with this model in the [DBRX Playground](https://huggingface.co/spaces/databricks/dbrx-instruct), subject to the terms and conditions below. 
# This demo is powered by [Databricks Foundation Model APIs](https://docs.databricks.com/en/machine-learning/foundation-models/index.html).
DESCRIPTION="""
This is a test
"""



duckdb.install_extension("spatial")
duckdb.load_extension("spatial")
duckdb.install_extension("httpfs")
duckdb.load_extension("httpfs")

duckdb.sql("create or replace view pad as select * from read_parquet('https://data.source.coop/cboettig/pad-us-3/pad-mobi.parquet')")

st.title(TITLE)

# Set OpenAI API key from Streamlit secrets
client = OpenAI(api_key=st.secrets["OPENAI_API_KEY"])

# Set a default model
if "openai_model" not in st.session_state:
    st.session_state["openai_model"] = "gpt-3.5-turbo"
# "gpt-4"

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []


setup = '''
You are a database administrator, and expert in SQL. You will be helping me write complex SQL queries. I will explain you my needs, you will generate SQL queries against my database.

My application does: Conservation prioritization of protected areas to help meet US 30x30 conservation goals. 

Please reply only with the SQL code that I will need to execute. Do not include an explanation of the code.
Please reply with raw SQL, do not use markdown formatting.

The database is a POSTGIS Postgres database, please take it into consideration when generating PLSQL/SQL. Please avoid ST_Within queries if possible, because they are so slow.

I will provide you with a description of the structure of my tables. You must remember them and use them for generating SQL queries. Once you read them all, just remember them for future and say nothing else.

Here are the tables :

Table "pad"
┌─────────────┬─────────────┬─────────┬─────────┬─────────┬─────────┐
│ column_name │ column_type │  null   │   key   │ default │  extra  │
│   varchar   │   varchar   │ varchar │ varchar │ varchar │ varchar │
├─────────────┼─────────────┼─────────┼─────────┼─────────┼─────────┤
│ FID         │ INTEGER     │ YES     │ NULL    │ NULL    │ NULL    │
│ time        │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
│ rsr         │ DOUBLE      │ YES     │ NULL    │ NULL    │ NULL    │
│ richness    │ DOUBLE      │ YES     │ NULL    │ NULL    │ NULL    │
│ bucket      │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
│ FeatClass   │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
│ Mang_Name   │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
│ Mang_Type   │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
│ Des_Tp      │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
│ Pub_Access  │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
│ GAP_Sts     │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
│ IUCN_Cat    │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
│ Unit_Nm     │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
│ area        │ DOUBLE      │ YES     │ NULL    │ NULL    │ NULL    │
│ geometry    │ BLOB        │ YES     │ NULL    │ NULL    │ NULL    │
├─────────────┴─────────────┴─────────┴─────────┴─────────┴─────────┤
└───────────────────────────────────────────────────────────────────┘
'''

# Display chat messages from history on app rerun
#for message in st.session_state.messages:
#    with st.chat_message(message["role"]):
#        st.markdown(message["content"])

# Accept user input
if prompt := st.chat_input("What is the total area in each GAP_Sts?"):
    # Add user message to chat history
    st.session_state.messages.append({"role": "system", "content": setup})

    st.session_state.messages.append({"role": "user", "content": prompt})
    # Display user message in chat message container
    # with st.chat_message("user"):
        #st.markdown(prompt)

  # Display assistant response in chat message container
    with st.chat_message("assistant"):
        stream = client.chat.completions.create(
            model=st.session_state["openai_model"],
            messages=[
                {"role": m["role"], "content": m["content"]}
                for m in st.session_state.messages
            ],
            stream=True,
        )
        response = st.write_stream(stream)
        st.divider()
        # st.write(response)
        df = duckdb.sql(response).df()
        st.table(df)
    st.session_state.messages.append({"role": "assistant", "content": response})



with st.sidebar:
    with st.container():
        st.title("Examples")
        for prompt in EXAMPLE_PROMPTS:
            st.markdown(prompt)


            # st.button(prompt, args=(prompt,)), on_click=handle_user_input)