import gradio as gr
import os
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings.fastembed import FastEmbedEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.llms import LlamaCpp
from langchain.prompts import ChatPromptTemplate
from huggingface_hub import hf_hub_download

# --- 1. Model Setup (The "Brain") ---
# We download a "Quantized" (compressed) version of Llama 3 to run on CPU
MODEL_REPO = "QuantFactory/Meta-Llama-3-8B-Instruct-GGUF"
MODEL_FILE = "Meta-Llama-3-8B-Instruct.Q4_K_M.gguf"

print("Downloading model... this may take a minute on first run.")
model_path = hf_hub_download(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE,
    repo_type="model"
)

# Initialize the Model
llm = LlamaCpp(
    model_path=model_path,
    n_ctx=4096,       # Context window size
    temperature=0.7,  # Creativity
    max_tokens=2000,  # Max length of output
    n_batch=512,
    verbose=True
)

# --- 2. The Core Logic ---
def generate_question_paper(pdf_file, difficulty, num_questions):
    if not pdf_file:
        return "Please upload a PDF file first."
    
    try:
        # A. Load PDF
        loader = PyPDFLoader(pdf_file.name)
        pages = loader.load()
        
        # B. Split Text
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=100
        )
        chunks = text_splitter.split_documents(pages)
        
        # C. Vector Store (FAISS)
        # We use FAISS (Ram-based) instead of Chroma for better Cloud compatibility
        embeddings = FastEmbedEmbeddings()
        vector_store = FAISS.from_documents(chunks, embeddings)
        
        # D. Retrieve Context
        # Get the top 7 most relevant chunks
        retriever = vector_store.as_retriever(search_kwargs={"k": 7})
        context_docs = retriever.invoke("Key concepts and definitions")
        context_text = "\n\n".join([doc.page_content for doc in context_docs])
        
        # E. Prompt
        template = """
        You are an expert academic examiner. Create a formal Question Paper based ONLY on the context provided below.
        
        CONTEXT:
        {context}
        
        INSTRUCTIONS:
        - Difficulty: {difficulty}
        - Total Questions: {num_questions}
        - Format:
          Section A: Multiple Choice Questions (MCQs)
          Section B: Short Answer Questions
          Section C: Long Answer/Essay Questions
        - Provide the Answer Key for MCQs at the very end.
        
        Do not output conversational text. Output ONLY the exam paper.
        """
        
        prompt = ChatPromptTemplate.from_template(template)
        
        # F. Generate
        chain = prompt | llm
        response = chain.invoke({
            "context": context_text, 
            "difficulty": difficulty,
            "num_questions": num_questions
        })
        
        return response

    except Exception as e:
        return f"Error processing PDF: {str(e)}"

# --- 3. The UI ---
theme = gr.themes.Soft(primary_hue="blue")

with gr.Blocks(theme=theme, title="AI Question Paper Generator") as demo:
    gr.Markdown("# 📄 AI Question Paper Generator")
    gr.Markdown("Hosted on Hugging Face • Powered by Llama 3 (GGUF)")
    
    with gr.Row():
        with gr.Column(scale=1):
            pdf_input = gr.File(label="Upload Study Material (PDF)")
            
            with gr.Group():
                difficulty = gr.Radio(
                    ["Easy", "Medium", "Hard"], 
                    label="Difficulty", 
                    value="Medium"
                )
                num_questions = gr.Slider(
                    5, 20, value=10, step=1, 
                    label="Total Questions"
                )
            
            btn = gr.Button("Generate Question Paper", variant="primary")
        
        with gr.Column(scale=2):
            output = gr.Markdown(label="Generated Paper")

    btn.click(
        fn=generate_question_paper,
        inputs=[pdf_input, difficulty, num_questions],
        outputs=output
    )

if __name__ == "__main__":
    demo.launch()