import numpy as np
import faiss, torch, gradio as gr
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForCausalLM

DOCS = [
    "The Generative AI Summer Bootcamp at Najran University runs for three weeks.",
    "Week 1 covers the foundations of generative AI and prompt engineering.",
    "Week 2 focuses on large language models, embeddings, and retrieval-augmented generation.",
    "Week 3 covers multimodal AI: vision, audio, and image generation.",
    "The bootcamp prepares students for the NVIDIA NCA-GENL and NCA-GENM associate exams.",
    "The NCA-GENL exam has 50 to 60 questions and a time limit of one hour.",
    "Students need a free Google Colab account and a free Hugging Face account.",
    "The refund policy: tuition is fully refundable up to seven days before the start date.",
    "All lab notebooks run on a free Colab T4 GPU with about 15 GB of memory.",
    "Certificates are issued to students who complete every hands-on lab.",
]

embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
emb = np.asarray(embedder.encode(DOCS, normalize_embeddings=True), dtype="float32")
index = faiss.IndexFlatIP(emb.shape[1]); index.add(emb)

GEN_ID = "Qwen/Qwen2.5-1.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(GEN_ID)
model = AutoModelForCausalLM.from_pretrained(GEN_ID, torch_dtype=torch.float32)

def retrieve(query, k=3):
    q = np.asarray(embedder.encode([query], normalize_embeddings=True), dtype="float32")
    scores, idxs = index.search(q, k)
    return [(DOCS[i], float(s)) for i, s in zip(idxs[0], scores[0])]

def generate(prompt, max_new_tokens=256):
    msgs = [{"role": "user", "content": prompt}]
    inp = tokenizer.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt")
    out = model.generate(inp, max_new_tokens=max_new_tokens, do_sample=False,
                         pad_token_id=tokenizer.eos_token_id)
    return tokenizer.decode(out[0][inp.shape[1]:], skip_special_tokens=True).strip()

def rag_answer(question, k=3, min_score=0.15):
    chunks = retrieve(question, k=k)
    if not chunks or chunks[0][1] < min_score:
        return "I dont know. (No relevant context was found.)", []
    ctx = "\n".join(f"[{i+1}] {t}" for i, (t, _) in enumerate(chunks))
    prompt = ("Answer using ONLY the context below. If the answer is not in the context, "
              "say: I dont know. Cite sources like [1], [2].\n\n"
              f"Context:\n{ctx}\n\nQuestion: {question}\n\nAnswer:")
    return generate(prompt), [t for t, _ in chunks]

def chat_fn(question):
    answer, sources = rag_answer(question)
    srcs = "\n".join(f"{i}. {s}" for i, s in enumerate(sources, 1)) or "(none)"
    return answer, srcs

demo = gr.Interface(
    fn=chat_fn,
    inputs=gr.Textbox(label="Ask about the bootcamp"),
    outputs=[gr.Textbox(label="Answer"), gr.Textbox(label="Sources")],
    title="Najran Bootcamp RAG Chatbot",
)

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