import os
import streamlit as st
from dotenv import load_dotenv
from langchain_community.document_loaders.pdf import PyPDFLoader
from langchain.indexes.vectorstore import VectorstoreIndexCreator
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import TextLoader
from langchain.chains.llm import LLMChain
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.messages import AIMessage, HumanMessage
from langchain_community.llms import HuggingFaceEndpoint
from langchain_core.output_parsers import StrOutputParser
from langchain.embeddings import HuggingFaceEmbeddings
from langchain_core.prompts import ChatPromptTemplate
#from langchain.embeddings import OpenAIEmbeddings
#from langchain.vectorstores import Chroma
from langchain.chains import create_retrieval_chain
load_dotenv()

api_token = os.getenv("HUGGINGFACE_API_KEY")

# Define the repository ID and task
repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
task = "text-generation"

# App config
st.set_page_config(page_title="INTERVIEW BOT",page_icon= "🌍")
st.title("INTERVIEW BOT")

@st.cache_resource
def load_text_and_create_index():
    # Load the text file
    loader = PyPDFLoader("res.pdf")
    documents = loader.load()
    
    # Split the text into chunks
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=10)
    texts = text_splitter.split_documents(documents)
    
    # Create embeddings
    embeddings = HuggingFaceEmbeddings(model_name='all-MiniLM-L12-v2')
    
    # Create the FAISS index
    db = FAISS.from_documents(texts, embeddings)
    
    return db

db = load_text_and_create_index()

template = """
You are INTERVIEWER CHATBOT. PEOPLE WILL ASK YOU FOR SUGGESTION AND YOU WILL HELP THEM IN EVERY WAY TO LAND THEM A PERFECT JOB.
Chat history:
{chat_history}

User question:
{user_question}
"""
prompt = ChatPromptTemplate.from_template(template)

# Function to get a response from the model
def get_response(user_query, chat_history,db):
    # Initialize the Hugging Face Endpoint
    llm = HuggingFaceEndpoint(
        huggingfacehub_api_token=api_token,
        repo_id=repo_id,
        task=task
    )
    chain = prompt | llm | StrOutputParser()
    retriever = db.as_retriever()
    chain_ret = create_retrieval_chain(retriever,chain)
    response = chain.invoke({
        "chat_history": chat_history,
        "user_question": user_query,
    })
    return response

# Initialize session state.
if "chat_history" not in st.session_state:
    st.session_state.chat_history = [
        AIMessage(content="Hello, I am Interview Bot How can I help you?"),
    ]

# Display chat history.
for message in st.session_state.chat_history:
    if isinstance(message, AIMessage):
        with st.chat_message("AI"):
            st.write(message.content)
    elif isinstance(message, HumanMessage):
        with st.chat_message("Human"):
            st.write(message.content)

# User input
db = load_text_and_create_index()
user_query = st.chat_input("Type your message here...")
if user_query is not None and user_query != "":
    st.session_state.chat_history.append(HumanMessage(content=user_query))

    with st.chat_message("Human"):
        st.markdown(user_query)

    response = get_response(user_query, st.session_state.chat_history,db)

    # Remove any unwanted prefixes from the response u should use these function but 
#before using it I requestto[replace("bot response:", "").strip()] combine 1&2 to run without error.

    #1.response = response.replace("AI response:", "").replace("chat response:", "").
    #2.replace("bot response:", "").strip()

    with st.chat_message("AI"):
        st.write(response)

    st.session_state.chat_history.append(AIMessage(content=response))