print("start1")
import os
import sys
import subprocess
import gradio as gr
from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.schema import Document
print("start2")


# Check if already installed to avoid reinstalling
try:
    import llama_cpp
    print("llama_cpp already installed.")
except ImportError:
    print("Installing llama_cpp from wheel...")
    subprocess.check_call([
    sys.executable, "-m", "pip", "install",
    "llama-cpp-python", "--no-binary", ":all:", "--force-reinstall"
])


from llama_cpp import Llama
print("start3")
import warnings
warnings.filterwarnings("ignore")

print("Start")
import subprocess

subprocess.run([
    "huggingface-cli", "download",
    "TheBloke/Mistral-7B-Instruct-v0.1-GGUF",
    "mistral-7b-instruct-v0.1.Q2_K.gguf",
    "--local-dir", "./models",
    "--local-dir-use-symlinks", "False"
], check=True)

# ------------------------------
# Device and Embedding Setup (CPU optimized)
# ------------------------------
modelPath = "sentence-transformers/all-mpnet-base-v2"
model_kwargs = {"device": "cpu"}  # Force CPU usage
encode_kwargs = {"normalize_embedding": False}

embeddings = HuggingFaceEmbeddings(
    model_name=modelPath,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs
)

# ------------------------------
# Load Mistral GGUF via llama.cpp (CPU optimized)
# ------------------------------
llm_cpp = Llama(
    model_path="./models/mistral-7b-instruct-v0.1.Q2_K.gguf",
    n_ctx=2048,
    n_threads=4,  # Adjust based on your CPU cores
    n_gpu_layers=0,  # Force CPU-only
    temperature=0.7,
    top_p=0.9,
    repeat_penalty=1.1
)

# ------------------------------
# LangChain-compatible wrapper
# ------------------------------
def mistral_llm(prompt):
    output = llm_cpp(
        prompt,
        max_tokens=512,  # Reduced for CPU performance
        stop=["</s>", "[INST]", "[/INST]"]
    )
    return output["choices"][0]["text"].strip()

# ------------------------------
# Prompt Template (unchanged)
# ------------------------------
def get_qa_prompt():
    template = """<s>[INST] \
You are a helpful, knowledgeable AI assistant. Answer the user's question based on the provided context.

Guidelines:
- Respond in a natural, conversational tone
- Be detailed but concise
- Use paragraphs and bullet points when appropriate
- If you don't know, say so
- Maintain a friendly and professional demeanor

Conversation History:
{chat_history}

Relevant Context:
{context}

Current Question: {question} 

Provide a helpful response: [/INST]"""
    return PromptTemplate(
        template=template,
        input_variables=["context", "question", "chat_history"]
    )

# ------------------------------
# PDF and Chat Logic (optimized for CPU)
# ------------------------------
def pdf_text(pdf_docs):
    text = ""
    for doc in pdf_docs:
        reader = PdfReader(doc)
        for page in reader.pages:
            page_text = page.extract_text()
            if page_text:
                text += page_text + "\n"
    return text

def get_chunks(text):
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=800,  # Smaller chunks for CPU
        chunk_overlap=100,
        length_function=len
    )
    chunks = splitter.split_text(text)
    return [Document(page_content=chunk) for chunk in chunks]

def get_vectorstore(documents):
    db = FAISS.from_documents(documents, embedding=embeddings)
    db.save_local("faiss_index")

def format_chat_history(history):
    return "\n".join([f"User: {q}\nAssistant: {a}" for q, a in history[-2:]])  # Shorter history

def handle_pdf_upload(pdf_files):
    if not pdf_files:
        return "⚠️ Upload at least one PDF"
    try:
        text = pdf_text(pdf_files)
        if not text.strip():
            return "⚠️ Could not extract text"
        chunks = get_chunks(text)
        get_vectorstore(chunks)
        return f"✅ Processed {len(pdf_files)} PDF(s) with {len(chunks)} chunks"
    except Exception as e:
        return f"❌ Error: {str(e)}"

def user_query(msg, chat_history):
    if not os.path.exists("faiss_index"):
        chat_history.append((msg, "Please upload PDF documents first."))
        return "", chat_history

    try:
        db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
        retriever = db.as_retriever(search_kwargs={"k": 2})  # Fewer documents for CPU
        docs = retriever.get_relevant_documents(msg)
        context = "\n\n".join([d.page_content for d in docs][:2])  # Limit context

        prompt = get_qa_prompt()
        final_prompt = prompt.format(
            context=context[:1500],  # Further limit context size
            question=msg,
            chat_history=format_chat_history(chat_history)
        )

        response = mistral_llm(final_prompt)
        chat_history.append((msg, response))
        return "", chat_history
    except Exception as e:
        error_msg = f"Sorry, I encountered an error: {str(e)}"
        chat_history.append((msg, error_msg))
        return "", chat_history

# ------------------------------
# Gradio Interface (your exact requested format)
# ------------------------------
with gr.Blocks(theme=gr.themes.Soft(), title="PDF Chat Assistant") as demo:
    with gr.Row():
        gr.Markdown("""
        # 📚 PDF Chat Assistant
        ### Have natural conversations with your documents ((Note: This Space runs on CPU, so responses may take a few mins.))
        """)
    with gr.Row():
        with gr.Column(scale=1, min_width=300):
            gr.Markdown("### Document Upload")
            pdf_input = gr.File(
                file_types=[".pdf"],
                file_count="multiple",
                label="Upload PDFs",
                height=100
            )
            upload_btn = gr.Button("Process Documents", variant="primary")
            status_box = gr.Textbox(label="Status", interactive=False)
            gr.Markdown("""
            **Instructions:**
            1. Upload PDF documents
            2. Click Process Documents
            3. Start chatting in the right panel
            """)

        with gr.Column(scale=2):
            chatbot = gr.Chatbot(
                height=600,
                bubble_full_width=False,
                avatar_images=(
                    "user.png", 
                    "bot.png"
                )
            )
            
            with gr.Row():
                message = gr.Textbox(
                    placeholder="Type your question about the documents...",
                    show_label=False,
                    container=False,
                    scale=7,
                    autofocus=True
                )
                submit_btn = gr.Button("Send", variant="primary", scale=1)
            
            with gr.Row():
                clear_chat = gr.Button("🧹 Clear Conversation")
                examples = gr.Examples(
                    examples=[
                        "Summarize the key points from the documents",
                        "What are the main findings?",
                        "Explain this in simpler terms"
                    ],
                    inputs=message,
                    label="Example Questions"
                )

    upload_btn.click(handle_pdf_upload, inputs=pdf_input, outputs=status_box)
    submit_btn.click(user_query, inputs=[message, chatbot], outputs=[message, chatbot])
    message.submit(user_query, inputs=[message, chatbot], outputs=[message, chatbot])
    clear_chat.click(lambda: [], None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch()  # Disable sharing for local CPU use