import os
import getpass
from langchain_community.document_loaders import ConfluenceLoader
from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores.faiss import FAISS
import google.generativeai as genai
from langchain.prompts import PromptTemplate
from langchain.chains.question_answering import load_qa_chain
import streamlit as st

confluence_api_key = os.environ["CONFLUENCE_API_KEY"]

if "GOOGLE_API_KEY" not in os.environ:
    os.environ["GOOGLE_API_KEY"] = getpass.getpass("Please provide Google API Key")

google_api_key = os.environ['GOOGLE_API_KEY']
genai.configure(api_key=google_api_key)


loader = ConfluenceLoader(
    url=os.environ["CONFLUENCE_URL"], space_key=os.environ['SPACE_KEY'], username=os.environ['USERNAME'], api_key=confluence_api_key
)

conf_docs = loader.load(page_id=os.environ["PAGE_ID"])


text_splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=1000)
chunks = text_splitter.split_text(conf_docs[-1].page_content)

embeddings = GoogleGenerativeAIEmbeddings(model='models/embedding-001')
llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash-latest")

vector_store = FAISS.from_texts(chunks, embedding=embeddings)
vector_store.save_local("faiss_index")

def get_response(query):
    prompt_template = """
        Answer the question as detailed as possible from the provided context, make sure to provide all the details, if the answer is not in
        provided context just say, "answer is not available in the context", don't provide the wrong answer\n\n
        Context:\n {context}?\n
        Question: \n{question}\n

        Answer:
    """
    prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
    chain = load_qa_chain(llm, chain_type="stuff", prompt=prompt)
    db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
    docs = db.similarity_search(query)
    response =  chain({"input_documents" : docs, "question": query}, return_only_outputs = True)
    return response["output_text"]


if __name__ == '__main__':
    st.set_page_config("Chat with Confluence Page")
    st.header("Chat with Confluence Page using AI")

    question = st.text_input("Ask questions related to login and registration")
    answer = get_response(question)
    st.write("Reply: ", answer)