import os
import tempfile
from typing import TypedDict, List

import gradio as gr
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document
from langgraph.graph import StateGraph, END

# Optional readers
from pypdf import PdfReader
import docx

EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

llm = ChatGoogleGenerativeAI(
    model="gemini-2.5-flash",
    google_api_key=os.getenv("GOOGLE_API_KEY"),
    temperature=0.2,
)
embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)

VECTORSTORE = None
RETRIEVER = None

class AgentState(TypedDict):
    question: str
    context: str
    answer: str
    route: str


def read_file(path: str) -> str:
    ext = os.path.splitext(path)[1].lower()
    if ext == ".pdf":
        reader = PdfReader(path)
        return "\n".join(page.extract_text() or "" for page in reader.pages)
    if ext == ".docx":
        d = docx.Document(path)
        return "\n".join(p.text for p in d.paragraphs)
    if ext in [".txt", ".md", ".csv"]:
        with open(path, "r", encoding="utf-8", errors="ignore") as f:
            return f.read()
    return ""


def build_rag(files):
    global VECTORSTORE, RETRIEVER
    if not files:
        return "Please upload at least one PDF, DOCX, TXT, MD, or CSV file."

    docs: List[Document] = []
    for f in files:
        text = read_file(f.name)
        if text.strip():
            docs.append(Document(page_content=text, metadata={"source": os.path.basename(f.name)}))

    if not docs:
        return "No readable text found. Please upload PDF, DOCX, TXT, MD, or CSV files."

    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    chunks = splitter.split_documents(docs)

    persist_dir = tempfile.mkdtemp(prefix="chroma_")
    VECTORSTORE = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=persist_dir,
        collection_name="rag_collection",
    )
    RETRIEVER = VECTORSTORE.as_retriever(search_kwargs={"k": 4})
    return f"Knowledge base ready. Loaded {len(docs)} file(s) and created {len(chunks)} chunks."


def retrieve_node(state: AgentState):
    global RETRIEVER
    if RETRIEVER is None:
        return {"context": "", "route": "general"}
    docs = RETRIEVER.invoke(state["question"])
    context = "\n\n".join(
        [f"Source: {d.metadata.get('source', 'uploaded file')}\n{d.page_content}" for d in docs]
    )
    return {"context": context, "route": "rag" if context.strip() else "general"}


def router_node(state: AgentState):
    return state.get("route", "general")


def rag_answer_node(state: AgentState):
    prompt = f"""
You are a helpful RAG assistant. Answer using ONLY the context below.
If the answer is not in the context, say: "I cannot find this in the uploaded documents."

Context:
{state['context']}

Question:
{state['question']}
"""
    return {"answer": llm.invoke(prompt).content}


def general_answer_node(state: AgentState):
    prompt = f"""
The user has not built a usable knowledge base yet, or no relevant context was found.
Answer briefly and remind the user to upload documents for RAG-based answers.

Question: {state['question']}
"""
    return {"answer": llm.invoke(prompt).content}


graph = StateGraph(AgentState)
graph.add_node("retrieve", retrieve_node)
graph.add_node("rag_answer", rag_answer_node)
graph.add_node("general_answer", general_answer_node)
graph.set_entry_point("retrieve")
graph.add_conditional_edges(
    "retrieve",
    router_node,
    {"rag": "rag_answer", "general": "general_answer"},
)
graph.add_edge("rag_answer", END)
graph.add_edge("general_answer", END)
app_graph = graph.compile()


def ask_rag(question):
    if not question or not question.strip():
        return "Please enter a question."
    result = app_graph.invoke({"question": question, "context": "", "answer": "", "route": ""})
    return result["answer"]

with gr.Blocks(title="Gemini LangGraph RAG") as demo:
    gr.Markdown("# Gemini LangGraph RAG on Hugging Face")
    gr.Markdown("Upload documents, build a temporary Chroma knowledge base, then ask questions.")

    with gr.Row():
        files = gr.File(
            label="Upload documents",
            file_count="multiple",
            file_types=[".pdf", ".docx", ".txt", ".md", ".csv"],
        )
        status = gr.Textbox(label="Status")

    build_btn = gr.Button("Build Knowledge Base")
    build_btn.click(build_rag, inputs=files, outputs=status)

    question = gr.Textbox(label="Question", placeholder="Ask about your uploaded documents")
    answer = gr.Textbox(label="Answer", lines=10)
    ask_btn = gr.Button("Ask RAG")
    ask_btn.click(ask_rag, inputs=question, outputs=answer)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)