#!/usr/bin/env python3
"""
🏆 CAPSTONE: End-to-End RAG Chatbot
Combines ALL techniques from Pages 1-9:
  Page 1: Pipeline, Hub, AutoModel
  Page 3: Text generation, prompt format
  Page 6: Sentence embeddings, FAISS, RAG
  Page 7: Gradio ChatInterface, Spaces
  Page 8: QLoRA, 4-bit quantization, PEFT
  Page 9: DPO alignment concepts

This is a DEPLOYABLE app for HF Spaces!
"""

import gradio as gr
import faiss
import numpy as np
import torch
from sentence_transformers import SentenceTransformer
from transformers import pipeline

# ═══════════════════════════════════════════════════
# COMPONENT 1: RETRIEVER (Page 6 — Embeddings + FAISS)
# ═══════════════════════════════════════════════════
print("📚 Loading retriever...")
retriever = SentenceTransformer("all-MiniLM-L6-v2")

# Knowledge base — replace with YOUR documents!
KNOWLEDGE_BASE = [
    "Hugging Face was founded in 2016 by Clément Delangue, Julien Chaumond, and Thomas Wolf.",
    "The Transformers library supports over 200 model architectures including BERT, GPT, T5, and LLaMA.",
    "Fine-tuning BERT on IMDB sentiment analysis typically achieves 93%+ accuracy.",
    "LoRA (Low-Rank Adaptation) allows fine-tuning large models by training only 0.1-1% of parameters.",
    "QLoRA combines 4-bit quantization with LoRA, enabling 7B model fine-tuning on a single T4 GPU.",
    "DPO (Direct Preference Optimization) is a simpler alternative to RLHF for model alignment.",
    "Gradio allows creating ML demo web apps with just Python, deployable to HF Spaces for free.",
    "FAISS is Facebook's library for efficient similarity search across millions of vectors in milliseconds.",
    "The Trainer API handles training loops, evaluation, logging, checkpointing, and multi-GPU automatically.",
    "Sentence Transformers encode text into dense vectors for semantic similarity and search tasks.",
    "Named Entity Recognition (NER) identifies people, locations, organizations in text using BIO tagging.",
    "T5 treats all NLP tasks as text-to-text: summarize, translate, classify — all with the same model.",
    "BERT uses bidirectional attention for understanding. GPT uses causal attention for generation.",
    "The Model Hub hosts over 500,000 pre-trained models for NLP, vision, audio, and multimodal tasks.",
    "Jakarta is the capital of Indonesia with a population of approximately 10.56 million people.",
    "Indonesia has over 17,000 islands and declared independence on August 17, 1945.",
    "Python is the most popular programming language for machine learning and data science.",
    "TensorFlow and PyTorch are the two most popular deep learning frameworks.",
]

# Build FAISS index
kb_embeddings = retriever.encode(KNOWLEDGE_BASE, convert_to_numpy=True)
faiss.normalize_L2(kb_embeddings)
index = faiss.IndexFlatIP(kb_embeddings.shape[1])
index.add(kb_embeddings)
print(f"  Indexed {index.ntotal} documents")

# ═══════════════════════════════════════════════════
# COMPONENT 2: GENERATOR (Page 3 — Text Generation)
# ═══════════════════════════════════════════════════
print("🤖 Loading generator...")
generator = pipeline("text2text-generation", model="google/flan-t5-base", device=0)
# FLAN-T5 base: 250M params, fits CPU/GPU easily
# For better quality: use fine-tuned model from Page 8!

# ═══════════════════════════════════════════════════
# COMPONENT 3: RAG PIPELINE (Page 6 — Retrieve + Generate)
# ═══════════════════════════════════════════════════
def retrieve(query, top_k=3):
    """Retrieve top-k relevant documents."""
    q_emb = retriever.encode([query], convert_to_numpy=True)
    faiss.normalize_L2(q_emb)
    scores, indices = index.search(q_emb, top_k)
    docs = [(KNOWLEDGE_BASE[i], float(s)) for i, s in zip(indices[0], scores[0]) if i >= 0]
    return docs

def generate_answer(question, context):
    """Generate answer using retrieved context."""
    prompt = f"""Answer the question based on the context below. If the answer is not in the context, say "I don't have information about that."

Context: {context}

Question: {question}

Answer:"""
    result = generator(prompt, max_length=200)
    return result[0]["generated_text"]

# ═══════════════════════════════════════════════════
# COMPONENT 4: GRADIO CHATBOT (Page 7 — ChatInterface)
# ═══════════════════════════════════════════════════
def chat(message, history):
    """RAG chatbot: retrieve → generate → respond with sources."""
    # Retrieve
    docs = retrieve(message, top_k=3)
    context = " ".join([doc for doc, score in docs])

    # Generate
    answer = generate_answer(message, context)

    # Format with sources
    sources = "\n\n📚 **Sources:**\n" + "\n".join(
        [f"- _{doc[:80]}..._ (relevance: {score:.0%})" for doc, score in docs])

    return answer + sources

# ═══════════════════════════════════════════════════
# COMPONENT 5: LAUNCH APP (Page 7 — Deploy to Spaces!)
# ═══════════════════════════════════════════════════
demo = gr.ChatInterface(
    fn=chat,
    title="🏆 RAG Chatbot — Hugging Face Knowledge Assistant",
    description="""Ask me anything about Hugging Face, Transformers, fine-tuning, NLP, or Indonesia!
Powered by: Sentence Transformers (retrieval) + FLAN-T5 (generation) + FAISS (vector search).
Built with techniques from the entire Learn Hugging Face series (Pages 1-9).""",
    examples=[
        "What is LoRA and how does it work?",
        "How accurate is BERT on IMDB sentiment analysis?",
        "What is the capital of Indonesia?",
        "What is the difference between BERT and GPT?",
        "How do I deploy a model to Hugging Face Spaces?",
        "What is DPO?",
    ],
    retry_btn="🔄 Retry",
    undo_btn="↩️ Undo",
    clear_btn="🗑️ Clear",
    theme=gr.themes.Soft(),
)

print("🚀 Launching RAG Chatbot...")
demo.launch()
# → Deploy to HF Spaces: upload app.py + requirements.txt
# → Free public URL: https://username-rag-chatbot.hf.space
# → ANYONE can chat with your knowledge-grounded AI! 🎉