import gradio as gr
from langchain.chains import create_retrieval_chain
from langchain.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationalRetrievalChain
from langchain.memory.chat_message_histories import ChatMessageHistory
from langchain_openai import ChatOpenAI
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.embeddings import HuggingFaceEmbeddings

# Embedding model
embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

# Vector store setup
persist_directory = 'vec_db'
vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding_model)
vectordb_retriever = vectordb.as_retriever(search_kwargs={'k': 5})

# LLM
llm = ChatOpenAI(model="gpt-4.1-nano", temperature=0.7)

# Load instructions
with open("instructions.txt", 'r') as file:
    instructions = file.read()

# Custom prompt
custom_prompt = ChatPromptTemplate.from_messages([
    ("system", instructions),
    MessagesPlaceholder(variable_name="chat_history"),
    ("user", "Question: {input}\nContext: {context}")
])

# Memory
memory = ConversationBufferMemory(
    memory_key="chat_history",
    return_messages=True
)

# Chains
question_answer_chain = create_stuff_documents_chain(llm, custom_prompt)
chain = create_retrieval_chain(vectordb_retriever, question_answer_chain)

# Main interaction function
def conversate_assistant(query, history):
    greetings = {"hey", "hi", "hello"}
    normalized_query = query.strip().lower()

    # Load the last 6 messages from memory
    chat_history = memory.load_memory_variables({})["chat_history"]
    chat_history = chat_history[-6:] if len(chat_history) >= 6 else chat_history

    # If greeting, skip context retrieval
    if normalized_query in greetings:
        response = question_answer_chain.invoke({
            "input": query,
            "context": [],  # Empty context for greetings
            "chat_history": chat_history
        })
        answer = response
    else:
        response = chain.invoke({
            "input": query,
            "chat_history": chat_history
        })
        answer = response['answer']

    # Save the interaction in memory
    memory.save_context({"input": query}, {"output": answer})

    return answer

# Gradio interface
demo = gr.ChatInterface(
    conversate_assistant,
    type="messages",
    title="Mini-Sid Ai Assistant"
)

# Launch the app
demo.launch()