import os
import streamlit as st
from langchain.memory.chat_message_histories import StreamlitChatMessageHistory
import time
from main import get_prompt

# Sidebar for selecting the model and entering the API key
st.sidebar.title("API Configuration")
api = st.sidebar.text_input("Enter your Anthropic API key", type="password")
st.sidebar.markdown("Get the API key from here: [https://console.anthropic.com/settings/keys](https://console.anthropic.com/settings/keys)")
icons = ""
# Chat interface
st.title("Prompt Writer")

# Store LLM-generated responses
if "messages" not in st.session_state.keys():
    st.session_state.messages = [{"role": "assistant", "content": "Hi. I'm an LLM powered prompting assistant helping you with writing efficient prompts to LLM for specific task. Give me the task as input."}]

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])

msgs = StreamlitChatMessageHistory(key="special_app_key")

for msg in msgs.messages:
    st.chat_message(msg.type).write(msg.content)

if prompt := st.chat_input():
    st.chat_message("human").write(prompt)
    msgs.add_user_message(prompt)

    with st.spinner("Waiting for response..."):
        res = get_prompt(prompt, api)
    if res:
        st.chat_message("ai").write(res)
        msgs.add_ai_message(res)
    else:
        st.error("No valid response received from the AI.")