import os
import tempfile
import logging
from typing import List, Dict
from langchain.embeddings import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.chains import ConversationalRetrievalChain
from langchain.document_loaders import PyPDFLoader
from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
from langchain.schema import Document
from langchain.memory import ConversationBufferMemory

import chainlit as cl
from chainlit.types import AskFileResponse

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Set Chainlit config directory
os.environ["CHAINLIT_CONFIG_DIR"] = tempfile.gettempdir()

# App configuration
APP_CONFIG = {
    "model_name": "gpt-4o",
    "temperature": 0,
    "chunk_size": 2000,  # Increased for book-length content
    "chunk_overlap": 200,
    "max_file_size_mb": 100,  # Increased for larger books
    "max_output_tokens": 4096,
    "embedding_model": "text-embedding-3-large"
}

# Initialize text splitter with more sophisticated settings
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=APP_CONFIG["chunk_size"],
    chunk_overlap=APP_CONFIG["chunk_overlap"],
    length_function=len,
    separators=["\n\n", "\n", " ", ""]  # Prioritize splitting on paragraph breaks
)

# System template
system_template = """You are GPT-4o, an advanced AI model with multimodal capabilities and high intelligence. Use the following pieces of context to answer the user's question.
If you don't know the answer, just say that you don't know, don't try to make up an answer.
ALWAYS return a "SOURCES" part in your answer.
The "SOURCES" part should be a reference to the source of the document from which you got your answer.

If the user greets with phrases like Hi, hello, How are you, etc., reply accordingly as well.

Example of your response should be:

The answer is foo
SOURCES: xyz

Begin!
----------------
{context}"""

# Create prompt template
messages = [
    SystemMessagePromptTemplate.from_template(system_template),
    HumanMessagePromptTemplate.from_template("{question}")
]
prompt = ChatPromptTemplate.from_messages(messages)

def process_file(file: AskFileResponse) -> List[Dict[str, str]]:
    """Process uploaded PDF file and return list of text chunks with metadata."""
    with tempfile.NamedTemporaryFile(mode="wb", delete=False, suffix=".pdf") as temp_file:
        temp_file.write(file.content)
        temp_file_path = temp_file.name

    loader = PyPDFLoader(temp_file_path)
    pages = loader.load_and_split()
    
    chunks_with_metadata = []
    for i, page in enumerate(pages):
        chunks = text_splitter.split_text(page.page_content)
        for j, chunk in enumerate(chunks):
            chunks_with_metadata.append({
                "content": chunk,
                "metadata": {
                    "page": i + 1,
                    "chunk": j + 1,
                    "source": f"page_{i+1}_chunk_{j+1}"
                }
            })
    
    return chunks_with_metadata

@cl.on_chat_start
async def start():
    await cl.Message(f"Welcome! Chat session started using {APP_CONFIG['model_name']} for text generation and {APP_CONFIG['embedding_model']} for embeddings. This combination is optimized for processing and querying book-length content. You can now upload a PDF of a book to begin!").send()
    
    files = await cl.AskFileMessage(
        content="Please upload a PDF file of a book to begin!",
        accept=["application/pdf"],
        max_size_mb=APP_CONFIG["max_file_size_mb"],
        timeout=3600,
    ).send()

    if not files:
        await cl.Message("No file uploaded. Please try again.").send()
        return

    file = files[0]
    await cl.Message(f"Starting to process `{file.name}`. This may take a few minutes...").send()

    try:
        # Process the file
        await cl.Message("Reading and splitting the document...").send()
        chunks_with_metadata = process_file(file)
        
        await cl.Message("Creating embeddings and setting up the retrieval system...").send()
        embeddings = OpenAIEmbeddings(model=APP_CONFIG["embedding_model"])
        docsearch = await cl.make_async(Chroma.from_documents)(
            [Document(page_content=chunk["content"], metadata=chunk["metadata"]) for chunk in chunks_with_metadata],
            embeddings
        )

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

        await cl.Message("Finalizing the chat model...").send()
        chain = ConversationalRetrievalChain.from_llm(
            ChatOpenAI(
                model_name=APP_CONFIG["model_name"],
                temperature=APP_CONFIG["temperature"],
                max_tokens=APP_CONFIG["max_output_tokens"],
                streaming=True
            ),
            chain_type="stuff",
            retriever=docsearch.as_retriever(search_kwargs={"k": 3}),
            memory=memory,
            return_source_documents=True,
            combine_docs_chain_kwargs={"prompt": prompt}
        )

        cl.user_session.set("chain", chain)
        await cl.Message(content=f"Processing `{file.name}` complete. You can now ask questions about the book!").send()

    except Exception as e:
        error_msg = f"An error occurred while processing the file: {str(e)}"
        logger.error(error_msg)
        await cl.Message(content=error_msg).send()

@cl.on_message
async def main(message: cl.Message):
    """Handle user messages and generate responses."""
    try:
        chain = cl.user_session.get("chain")
        cb = cl.AsyncLangchainCallbackHandler()
        
        res = await chain.acall(message.content, callbacks=[cb])
        answer = res["answer"]
        source_documents = res["source_documents"]
        
        text_elements = []
        for source_doc in source_documents:
            source_name = f"Page {source_doc.metadata['page']}, Chunk {source_doc.metadata['chunk']}"
            text_elements.append(cl.Text(content=source_doc.page_content, name=source_name))
        
        source_names = [text_el.name for text_el in text_elements]
        if source_names:
            answer += f"\nSources: {', '.join(source_names)}"
        else:
            answer += "\nNo sources found"
        
        await cl.Message(content=answer, elements=text_elements).send()
    
    except Exception as e:
        error_msg = f"An error occurred: {str(e)}"
        logger.error(error_msg)
        await cl.Message(content=error_msg).send()

if __name__ == "__main__":
    logger.info(f"Starting Chainlit application with model: {APP_CONFIG['model_name']} and embedding model: {APP_CONFIG['embedding_model']}")
    cl.run()