import os
import requests
import gradio as gr
from llama_cpp import Llama
from PyPDF2 import PdfReader
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity

# -------------------------------
# Model Setup
# -------------------------------
MODEL_URL = "https://huggingface.co/melsonop/mistral-chat-pdf/resolve/main/mistral-7b-instruct-v0.1.Q4_K_M.gguf"
MODEL_PATH = "mistral.gguf"

if not os.path.exists(MODEL_PATH):
    print("🔄 Downloading Mistral model...")
    with requests.get(MODEL_URL, stream=True) as r:
        r.raise_for_status()
        with open(MODEL_PATH, 'wb') as f:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)
    print("✅ Model downloaded successfully!")

llm = Llama(
    model_path=MODEL_PATH,
    n_ctx=2048,
    n_threads=4,
    n_batch=512,
    verbose=False
)

embed_model = SentenceTransformer("all-MiniLM-L6-v2")

# -------------------------------
# PDF + QA
# -------------------------------
docs = []
doc_embeddings = []

def load_pdf(file):
    global docs, doc_embeddings
    if file is None:
        return "❌ Please upload a valid PDF file."
    
    reader = PdfReader(file)
    text = ""
    for page in reader.pages:
        text += page.extract_text() + "\n"

    docs = text.split(". ")
    doc_embeddings = embed_model.encode(docs)
    return "✅ PDF loaded! You can now ask questions."

def chat_with_pdf(question):
    if not docs:
        return "❌ No PDF loaded yet. Please upload and load a PDF first."

    question_embedding = embed_model.encode([question])
    similarities = cosine_similarity(question_embedding, doc_embeddings)[0]
    top_indices = similarities.argsort()[-3:][::-1]
    context = "\n".join([docs[i] for i in top_indices])

    prompt = f"""You are a helpful assistant. Use the context below to answer the question.\n\nContext:\n{context}\n\nQuestion: {question}\nAnswer:"""

    output = llm(prompt=prompt, max_tokens=512, temperature=0.7)
    return output['choices'][0]['text'].strip()

# -------------------------------
# Gradio UI
# -------------------------------

with gr.Blocks() as demo:
    gr.Markdown("# 🤖 Chat with Your PDF using Mistral 7B (GGUF)")
    
    pdf = gr.File(label="Upload PDF", file_types=[".pdf"])
    load_output = gr.Textbox(label="Status", interactive=False)
    load_button = gr.Button("🔄 Load PDF")

    question = gr.Textbox(label="Ask a Question")
    answer = gr.Textbox(label="Answer")

    load_button.click(load_pdf, inputs=[pdf], outputs=[load_output])
    question.submit(chat_with_pdf, inputs=[question], outputs=[answer])

demo.launch()