import os
import getpass
import streamlit as st
from typing import List
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import HumanMessage, AIMessage
from langchain_google_genai import ChatGoogleGenerativeAI

class Chat:
    def __init__(self):
        self.chat = ChatGoogleGenerativeAI(model="gemini-pro", convert_system_message_to_human=True)
        self.history = []

        if "GOOGLE_API_KEY" not in os.environ:
            os.environ["GOOGLE_API_KEY"] = st.text_input("Provide your API key: ", type="password")

        self.prompt = ChatPromptTemplate.from_messages([
            ("system", "You are a senior software engineer. Answer all questions to the best of your ability."),
            MessagesPlaceholder(variable_name="messages"),
        ])

    def add_message(self, message: str):
        self.history.append(HumanMessage(content=message))

    def add_ai_message(self, message: str):
        self.history.append(AIMessage(content=message))

    def chat_with_user(self):
        st.title("Chat with AI")
        
        user_input = st.text_input("You: ", "")
        if st.button("Send"):
            if user_input:
                self.add_message(user_input)
                response = self.chat.invoke(self.prompt.format_prompt(messages=self.history).to_messages())
                self.add_ai_message(response.content)

                st.write(f"AI: {response.content}")

                # Display chat history
                for msg in self.history:
                    if isinstance(msg, HumanMessage):
                        st.write(f"User: {msg.content}")
                    else:
                        st.write(f"AI: {msg.content}")

if __name__ == "__main__":
    chat_instance = Chat()
    chat_instance.chat_with_user()