import os
import fitz  # PyMuPDF
import streamlit as st
import google.generativeai as genai
from dotenv import load_dotenv
from google.api_core.exceptions import GoogleAPIError, InvalidArgument

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
from langchain_community.vectorstores import FAISS
from langchain.chains.question_answering import load_qa_chain
from langchain.prompts import PromptTemplate

# Load environment variables
load_dotenv()
api_key = os.getenv("GOOGLE_API_KEY")
genai.configure(api_key=api_key)


# ✅ Function to read all PDF files (Farsi + English)
def get_pdf_text(pdf_docs):
    text = ""
    for pdf in pdf_docs:
        with fitz.open(stream=pdf.read(), filetype="pdf") as doc:
            for page in doc:
                page_text = page.get_text("text")
                if page_text:
                    text += page_text + "\n"
    return text


# ✅ Function to split text into chunks
def get_text_chunks(text):
    splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=300)
    chunks = splitter.split_text(text)
    return chunks


# ✅ Function to get embeddings for each chunk and save to vector store
def get_vector_store(chunks):
    try:
        embeddings = GoogleGenerativeAIEmbeddings(model="models/text-embedding-004")
        vector_store = FAISS.from_texts(chunks, embedding=embeddings)
        vector_store.save_local("faiss_index")
    except Exception as e:
        raise RuntimeError(f"Error creating vector store: {e}")


# ✅ Function to get conversational chain
def get_conversational_chain():
    prompt_template = """
    You are a helpful assistant. 
    Answer the question as detailed as possible using the provided context. 
    If the answer is unclear, summarize the most relevant part of the context. 
    Do not make up answers. if the question is asked in Farsi, provide the answer in the same asking language.

    Context:
    {context}

    Question:
    {question}

    Answer:
    """
    try:
        model = ChatGoogleGenerativeAI(model="gemini-1.5-flash", client=genai, temperature=0.3)
        prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
        chain = load_qa_chain(llm=model, chain_type="stuff", prompt=prompt)
        return chain
    except Exception as e:
        raise RuntimeError(f"Error creating conversational chain: {e}")


# ✅ Function to clear chat history
def clear_chat_history():
    st.session_state.messages = [{"role": "assistant", "content": "در خدمتیم"}]


# ✅ Function to handle user input
def user_input(user_question):
    try:
        embeddings = GoogleGenerativeAIEmbeddings(model="models/text-embedding-004")
        new_db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
        docs = new_db.similarity_search(user_question, k=15)
        chain = get_conversational_chain()
        response = chain({"input_documents": docs, "question": user_question}, return_only_outputs=True)
        return response
    except Exception as e:
        raise RuntimeError(f"Error while answering: {e}")


# ✅ Main function to run the Streamlit app
def main():
    st.set_page_config(
        page_title="Chatbot",
        layout="wide",
        initial_sidebar_state="expanded"
    )

    if "uploaded" not in st.session_state:
        st.session_state.uploaded = False

    if not st.session_state.uploaded:
        # Upload Page
        st.title("Your personal assistant ...")
        pdf_docs = st.file_uploader("فایل پی دی اف مورد نظر را آپلود کنید ", accept_multiple_files=True)
        if st.button("تایید"):
            if pdf_docs:
                try:
                    st.info("در حال پردازش ...")
                    raw_text = get_pdf_text(pdf_docs)
                    text_chunks = get_text_chunks(raw_text)
                    get_vector_store(text_chunks)
                    st.session_state.uploaded = True
                    st.success("پردازش موفق شد ✅")
                except RuntimeError as e:
                    st.error(str(e))
            else:
                st.error("لطفا حداقل یک فایل را آپلود کنید")
    else:
        # Chat Page
        st.title("Assistant ready ...")
        st.write("میتونین سوالتونو بپرسین 👇")

        if st.button("بازگشت به صفحه آپلود"):
            st.session_state.uploaded = False
            clear_chat_history()
            st.rerun()

        st.button('حذف مکالمه', on_click=clear_chat_history)

        if "messages" not in st.session_state:
            st.session_state.messages = [{"role": "assistant", "content": "در خدمتیم"}]

        # Show history
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(
                    f"""
                    <div style="direction: rtl; text-align: right; font-size: 16px;">
                        {message["content"]}
                    </div>
                    """,
                    unsafe_allow_html=True
                )

        if prompt := st.chat_input():
            st.session_state.messages.append({"role": "user", "content": prompt})
            with st.chat_message("user"):
                st.markdown(
                    f"""
                    <div style="direction: rtl; text-align: right; font-size: 16px;">
                        {prompt}
                    </div>
                    """,
                    unsafe_allow_html=True
                )

            if st.session_state.messages[-1]["role"] != "assistant":
                try:
                    with st.chat_message("assistant"):
                        response = user_input(prompt)
                        if response:
                            full_response = response['output_text']
                            st.markdown(
                                f"""
                                <div style="direction: rtl; text-align: right; font-size: 16px;">
                                    {full_response}
                                </div>
                                """,
                                unsafe_allow_html=True
                            )
                            st.session_state.messages.append({"role": "assistant", "content": full_response})
                except RuntimeError as e:
                    st.error(str(e))


if __name__ == "__main__":
    main()