import os
import streamlit as st  # type: ignore

from dotenv import load_dotenv
from PyPDF2 import PdfReader

from langchain_classic.text_splitter import RecursiveCharacterTextSplitter
from langchain_cohere import CohereEmbeddings, ChatCohere
from langchain_community.vectorstores import FAISS
from langchain_classic.chains import ConversationalRetrievalChain
from langchain_classic.memory import ConversationBufferMemory

load_dotenv()

st.set_page_config(
    page_title="Chat with Documents",
    page_icon="📚"
)


def get_pdf_text(pdf_docs):
    text = ""

    for pdf in pdf_docs:
        pdf_reader = PdfReader(pdf)

        for page in pdf_reader.pages:
            text += page.extract_text() or ""

    return text


def get_text_chunks(text):
    text_splitter = RecursiveCharacterTextSplitter(
        separators=["\n\n", "\n", " ", ""],
        chunk_size=1000,
        chunk_overlap=200,
        length_function=len
    )

    return text_splitter.split_text(text)


def get_embeddings():
    return CohereEmbeddings(
        model='embed-english-v3.0',
        user_agent="langchain"
    )


def get_conversational_chain(vectorstore):
    llm = ChatCohere(
        model="command-r-08-2024",
        temperature=0
    )

    memory = ConversationBufferMemory(
        memory_key="chat_history",
        return_messages=True,
        output_key="answer"
    )

    conversational_chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=vectorstore.as_retriever(),
        memory=memory,
        return_source_documents=True
    )

    return conversational_chain


def handle_userinput(user_question):
    response = st.session_state.conversational_chain.invoke(
        {"question": user_question}
    )

    answer = response["answer"]

    st.session_state.chat_history = (
        st.session_state.conversational_chain.memory
        .load_memory_variables({})["chat_history"]
    )

    for i, message in enumerate(st.session_state.chat_history):
        if message.type == "human":
            with st.chat_message("user"):
                st.write(message.content)
        else:
            with st.chat_message("assistant"):
                st.write(message.content)


def main():
    st.markdown(
        """
        <style>
        .main {
            background: linear-gradient(135deg, #f8fbff 0%, #eef5ff 100%);
        }
        .block-container {
            padding-top: 2rem;
            padding-bottom: 2rem;
        }
        .hero {
            background: linear-gradient(135deg, #0f172a 0%, #1d4ed8 100%);
            padding: 2rem 2rem;
            border-radius: 22px;
            box-shadow: 0 12px 30px rgba(29, 78, 216, 0.18);
            margin-bottom: 1.5rem;
            margin-top: 1.5rem;
        }
        .hero h1 {
            color: white !important;
            font-size: 2.5rem !important;
            font-weight: 800 !important;
            margin: 0 !important;
        }
        .hero p {
            color: #dbeafe !important;
            font-size: 1.05rem !important;
            margin-top: 0.5rem !important;
            margin-bottom: 0 !important;
        }
        </style>
        """,
        unsafe_allow_html=True,
    )

    st.markdown(
        """
        <div class="hero">
            <h1> AI Chatbot for Your Documents</h1>
            <p>Upload a PDF and ask questions in natural language.</p>
        </div>
        """,
        unsafe_allow_html=True,
    )

    if "conversational_chain" not in st.session_state:
        st.session_state.conversational_chain = None

    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []

    with st.sidebar:
        st.header("Upload Documents")

        pdf_docs = st.file_uploader(
            "Upload your PDF files",
            type=["pdf"],
            accept_multiple_files=True
        )

        if st.button("Process Documents"):
            if not pdf_docs:
                st.warning("Please upload at least one PDF.")
            elif not os.getenv("COHERE_API_KEY"):
                st.error("Cohere API key is missing.")
            else:
                try:
                    with st.spinner("Processing documents..."):
                        raw_text = get_pdf_text(pdf_docs)

                        if not raw_text.strip():
                            st.error("No extractable text found in the PDFs.")
                            return

                        chunks = get_text_chunks(raw_text)

                        embeddings = get_embeddings()

                        vectorstore = FAISS.from_texts(
                            texts=chunks,
                            embedding=embeddings
                        )

                        st.session_state.conversational_chain = (
                            get_conversational_chain(vectorstore)
                        )

                        st.session_state.chat_history = []

                    st.success("Documents processed successfully!")

                except Exception as e:
                    st.error(f"Error processing documents: {e}")

    user_question = st.chat_input("Ask a question about your PDFs")

    if user_question:
        if st.session_state.conversational_chain is None:
            st.warning("Please upload and process your documents first.")
        else:
            with st.chat_message("user"):
                st.write(user_question)

            with st.chat_message("assistant"):
                with st.spinner("Thinking..."):
                    try:
                        response = (
                            st.session_state.conversational_chain.invoke(
                                {"question": user_question}
                            )
                        )
                        st.write(response["answer"])

                        st.session_state.chat_history = (
                            st.session_state.conversational_chain
                            .memory.load_memory_variables(
                                {}
                            )["chat_history"]
                        )

                    except Exception as e:
                        st.error(f"Error generating answer: {e}")


if __name__ == "__main__":
    main()