import warnings
warnings.filterwarnings("ignore")

import uuid
import gradio as gr
import os

from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.document_loaders import PyPDFLoader
from langchain.chains import RetrievalQA
from langchain_community.embeddings import HuggingFaceEmbeddings


# -------------------------
# LLM (Hugging Face Router API)
# -------------------------

repo_id = "Qwen/Qwen2.5-7B-Instruct"

llm_endpoint = HuggingFaceEndpoint(
    repo_id=repo_id,
    task="text-generation",
    max_new_tokens=512,
    temperature=0.1,
    huggingfacehub_api_token=os.getenv("HF_TOKEN"),
    provider="hf-inference",   # forces router.huggingface.co
    model_kwargs={"wait_for_model": True},
)

llm = ChatHuggingFace(llm=llm_endpoint)


# -------------------------
# Document loader
# -------------------------

def document_loader(file_path: str):
    loader = PyPDFLoader(file_path)
    return loader.load()


# -------------------------
# Text splitter
# -------------------------

def text_splitter(data):
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=50,
        length_function=len,
    )
    return splitter.split_documents(data)


# -------------------------
# Embedding model
# -------------------------

def embedding_model():
    return HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"
    )


# -------------------------
# Vector DB
# -------------------------

def vector_database(chunks):
    chunks = [
        c for c in chunks
        if c.page_content and c.page_content.strip() and len(c.page_content.strip()) > 10
    ]

    vectordb = Chroma.from_documents(
        documents=chunks,
        embedding=embedding_model(),
        collection_name=f"pdf_chat_{uuid.uuid4().hex}",
    )
    return vectordb


# -------------------------
# Retriever
# -------------------------

def build_retriever(file_path: str):
    docs = document_loader(file_path)
    chunks = text_splitter(docs)
    vectordb = vector_database(chunks)
    return vectordb.as_retriever()


# -------------------------
# QA Chain
# -------------------------

def retriever_qa(file, query):
    retr = build_retriever(file)

    qa = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retr,
        return_source_documents=False,
    )

    result = qa.invoke({"query": query})
    return result["result"]


# -------------------------
# Gradio interface
# -------------------------

rag_application = gr.Interface(
    fn=retriever_qa,
    allow_flagging="never",
    inputs=[
        gr.File(
            label="Upload PDF File",
            file_count="single",
            file_types=[".pdf"],
            type="filepath",
        ),
        gr.Textbox(
            label="Input Query",
            lines=2,
            placeholder="Type your question here...",
        ),
    ],
    outputs=gr.Textbox(label="Output"),
    title="PDF QA Chatbot",
    description=(
        "Upload a PDF document and ask any question. "
        "The chatbot will try to answer using the provided document."
    ),
)

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