import os

# Still init with the fake one
if 'OPENAI_API_KEY' not in os.environ:
    os.environ['OPENAI_API_KEY'] = 'none'

import pandas as pd
import requests
import streamlit as st
from IPython.core.display import HTML
from langchain.agents import AgentType, initialize_agent, load_tools
from langchain.chat_models import ChatOpenAI
from PIL import Image
from rmrkl import ChatZeroShotAgent, RetryAgentExecutor

from chemcrow.agents import ChemCrow, make_tools
from chemcrow.agents.prompts import (FORMAT_INSTRUCTIONS, QUESTION_PROMPT,
                                     SUFFIX)
from chemcrow.frontend.streamlit_callback_handler import \
    StreamlitCallbackHandlerChem
from langchain.callbacks import StreamlitCallbackHandler
from chemcrow.frontend.utils import cdk

from dotenv import load_dotenv

load_dotenv()
ss = st.session_state


icon = Image.open('assets/logo0.png')
st.set_page_config(
    page_title="ChemCrow",
    page_icon = icon
)

# Set width of sidebar
st.markdown(
    """
    <style>
    [data-testid="stSidebar"][aria-expanded="true"]{
        min-width: 450px;
        max-width: 450px;
    }
    """,
    unsafe_allow_html=True,
)

# Session state
# st.session_state['molecule'] = "CCO"

agent = ChemCrow(
    model='gpt-4',
    temp=0.1,
    api_key=ss.get('api_key')
).agent_executor

tools = agent.tools

tool_list = pd.Series(
    {f"✅ {t.name}":t.description for t in tools}
).reset_index()
tool_list.columns = ['Tool', 'Description']

# sidebar
with st.sidebar:
    chemcrow_logo = Image.open('assets/chemcrow-logo-bold-new.png')
    st.image(chemcrow_logo)

    # Input OpenAI api key
    st.markdown('Input your OpenAI API key.')
    st.text_input('OpenAI API key', type='password', key='api_key', label_visibility="collapsed")

    # Display available tools
    st.markdown(f"# Available tools: {len(tool_list)}")
    st.dataframe(
        tool_list,
        use_container_width=True,
        hide_index=True,
        height=300
    )


# Agent execution
if prompt := st.chat_input():
    st.chat_message("user").write(prompt)
    with st.chat_message("assistant"):
        st_callback = StreamlitCallbackHandlerChem(
            st.container(),
            max_thought_containers = 4,
            collapse_completed_thoughts = False,
            output_placeholder=st.session_state
        )
        try:
            response = agent.run(prompt, callbacks=[st_callback])
            st.write(response)
        except:
            st.write("Please input a valid OpenAI API key.")