# ==========================================
# India AI Impact Summit 2026 Chatbot
# Deploy on Hugging Face Spaces (Gradio SDK)
# ==========================================

import gradio as gr

from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings, HuggingFacePipeline
from langchain_community.vectorstores import FAISS

from transformers import pipeline
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser


# --------- Load PDF ---------
PDF_PATH = "india_ai_impact_summit_2026.pdf"

loader = PyPDFLoader(PDF_PATH)
documents = loader.load()


# --------- Split Text ---------
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=80
)
docs = text_splitter.split_documents(documents)


# --------- Embeddings ---------
embedding_model = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2"
)


# --------- Vector Store ---------
vectorstore = FAISS.from_documents(docs, embedding_model)
retriever = vectorstore.as_retriever(search_kwargs={"k": 5})


# --------- LLM (FLAN-T5) ---------
pipe = pipeline(
    "text2text-generation",
    model="google/flan-t5-base",
    max_new_tokens=300
)
llm = HuggingFacePipeline(pipeline=pipe)


# --------- Prompt ---------
prompt = ChatPromptTemplate.from_template(
    """
    You are a helpful assistant for the India AI Impact Summit 2026.
    Answer the question using ONLY the context below.
    If the answer is not in the context, say "I don't know".

    Context:
    {context}

    Question:
    {question}
    """
)


# --------- RAG Chain ---------
rag_chain = (
    {
        "context": retriever,
        "question": RunnablePassthrough()
    }
    | prompt
    | llm
    | StrOutputParser()
)


# --------- Gradio UI ---------
def chat(question):
    return rag_chain.invoke(question)


demo = gr.Interface(
    fn=chat,
    inputs=gr.Textbox(lines=2, placeholder="Ask about the India AI Impact Summit 2026..."),
    outputs="text",
    title="🇮🇳 India AI Impact Summit 2026 Chatbot",
    description="Ask questions about the India AI Impact Summit 2026 – powered by RAG",
    examples=[
        ["What is the India AI Impact Summit 2026?"],
        ["What is BharatGPT-2?"],
        ["What is the National AI Mission 2.0?"],
        ["What is the AI Kisan programme?"],
        ["Who were the keynote speakers?"],
        ["What is the New Delhi AI Declaration?"],
        ["How is AI used in Indian healthcare?"],
    ]
)

demo.launch()