from typing import TypedDict, List

import streamlit as st
from langgraph.graph import StateGraph

from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.document_loaders import TextLoader

from langchain_text_splitters import RecursiveCharacterTextSplitter
from transformers import pipeline


# -------------------------
# STATE TYPE FOR LANGGRAPH
# -------------------------
class AgentState(TypedDict):
    question: str
    documents: List[str]
    answer: str
    reflection: str


# -------------------------
# LLM: HUGGINGFACE PIPELINE
# -------------------------

# Using a smaller model for faster loading on CPU and HF Spaces
llm_pipeline = pipeline(
    "text2text-generation",
    model="google/flan-t5-small",   # you can switch to flan-t5-base if you want
)

def generate_text(prompt: str) -> str:
    """Call the HF pipeline and return plain text."""
    out = llm_pipeline(prompt, max_new_tokens=256)
    if isinstance(out, list) and len(out) > 0 and "generated_text" in out[0]:
        return out[0]["generated_text"]
    return str(out)


# -------------------------
# EMBEDDINGS AND DOCUMENTS
# -------------------------

@st.cache_resource
def load_vectorstore():
    """Load documents, split, and build Chroma vector store once."""
    # Load plain text file; make sure knowledge.txt is in the same folder
    loader = TextLoader("knowledge.txt")
    docs = loader.load()

    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=300,
        chunk_overlap=50,
    )
    split_docs = text_splitter.split_documents(docs)

    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"
    )

    vectordb = Chroma.from_documents(
        split_docs,
        embedding=embeddings,
        persist_directory="./db"
    )
    return vectordb


vectordb = load_vectorstore()


# -------------------------
# LANGGRAPH NODES
# -------------------------

def plan_node(state: AgentState) -> AgentState:
    print("[PLAN] Understanding question...")
    # For now, always decide to retrieve. You could add logic here later.
    return state


def retrieve_node(state: AgentState) -> AgentState:
    print("[RETRIEVE] Searching knowledge base...")
    query = state["question"]

    # Newer LangChain versions: use similarity_search directly for stability
    docs = vectordb.similarity_search(query, k=4)

    state["documents"] = [doc.page_content for doc in docs]
    print(f"[RETRIEVE] Retrieved {len(state['documents'])} documents.")
    return state


def answer_node(state: AgentState) -> AgentState:
    print("[ANSWER] Generating answer from context...")

    context = "\n".join(state["documents"])

    prompt = f"""
You are a helpful assistant. Use only the context below to answer the question.

Context:
{context}

Question:
{state['question']}

Answer in 2-4 sentences, concise and clear.
"""

    answer = generate_text(prompt)
    state["answer"] = answer.strip()
    return state


def reflect_node(state: AgentState) -> AgentState:
    print("[REFLECT] Evaluating answer relevance...")

    reflection_prompt = f"""
Question: {state['question']}
Answer: {state['answer']}

Evaluate if the answer is relevant and complete based only on the question.
Reply in this format:
- Verdict: YES or NO
- Reason: one short sentence
"""

    reflection = generate_text(reflection_prompt)
    state["reflection"] = reflection.strip()
    return state


# -------------------------
# BUILD LANGGRAPH WORKFLOW
# -------------------------

builder = StateGraph(AgentState)

builder.add_node("plan", plan_node)
builder.add_node("retrieve", retrieve_node)
builder.add_node("answer", answer_node)
builder.add_node("reflect", reflect_node)

builder.set_entry_point("plan")
builder.add_edge("plan", "retrieve")
builder.add_edge("retrieve", "answer")
builder.add_edge("answer", "reflect")

agent = builder.compile()


# -------------------------
# STREAMLIT UI
# -------------------------

st.title("RAG Q&A Agent with LangGraph (Hugging Face Models)")
st.write(
    "Ask a question based on the knowledge stored in `knowledge.txt`. "
    "The agent will retrieve relevant context, answer, and then reflect on its own answer."
)

user_question = st.text_input("Enter your question:", value="What is renewable energy?")

if st.button("Ask"):
    if not user_question.strip():
        st.warning("Please enter a question.")
    else:
        # Initial state for LangGraph
        init_state: AgentState = {
            "question": user_question,
            "documents": [],
            "answer": "",
            "reflection": "",
        }

        with st.spinner("Running agent (plan → retrieve → answer → reflect)..."):
            result = agent.invoke(init_state)

        st.subheader("Final Answer")
        st.write(result["answer"])

        if result.get("documents"):
            st.subheader("Retrieved Context")
            for i, doc in enumerate(result["documents"], start=1):
                st.markdown(f"**Chunk {i}:**")
                st.write(doc)

        st.subheader("Reflection")
        st.write(result["reflection"])