import streamlit as st
from typing import TypedDict, List
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.tools import DuckDuckGoSearchRun
from langgraph.graph import StateGraph, END

# --- 1. STATE DEFINITION ---
class AgentState(TypedDict):
    question: str
    context: str
    answer: str
    trace: List[str]

# --- 2. DYNAMIC MODEL INITIALIZATION ---
def get_llm(token: str):
    """Initializes the LLM using the token provided in the UI."""
    try:
        llm = HuggingFaceEndpoint(
            repo_id="Qwen/Qwen2.5-7B-Instruct",
            task="text-generation",
            max_new_tokens=512,
            huggingfacehub_api_token=token,
        )
        return ChatHuggingFace(llm=llm, token=token)
    except Exception as e:
        st.error(f"Failed to initialize model: {str(e)}")
        return None

# --- 3. NODE LOGIC ---
def router_node(state: AgentState, token: str):
    llm = get_llm(token)
    prompt = ChatPromptTemplate.from_template(
        "Determine if this needs a web search for real-time info. "
        "Reply ONLY with 'SEARCH' or 'KNOWLEDGE'.\nQuestion: {question}"
    )
    chain = prompt | llm
    res = chain.invoke({"question": state["question"]})
    decision = "SEARCH" if "SEARCH" in res.content.upper() else "KNOWLEDGE"
    return {"context": decision, "trace": [f"🔍 Router: Selected {decision} path"]}

def search_node(state: AgentState):
    search = DuckDuckGoSearchRun()
    results = search.run(state["question"])
    return {"context": results, "trace": state["trace"] + ["🌐 Web Search: Results gathered"]}

def writer_node(state: AgentState, token: str):
    llm = get_llm(token)
    prompt = ChatPromptTemplate.from_template(
        "Context: {context}\n\nProvide a detailed answer to: {question}"
    )
    chain = prompt | llm
    res = chain.invoke({"context": state["context"], "question": state["question"]})
    return {"answer": res.content, "trace": state["trace"] + ["✍️ Writer: Generated final response"]}

# --- 4. APP UI ---
st.set_page_config(page_title="LangChain Advanced Practice", layout="wide")

st.title("🏗️ LangGraph Practice Lab")
st.markdown("Build and test advanced autonomous workflows using open-source models.")

# Sidebar for Token Management
with st.sidebar:
    st.header("Authentication")
    user_token = st.text_input("Enter Hugging Face Token (hf_...)", type="password")
    st.info("Get your token at: [hf.co/settings/tokens](https://huggingface.co/settings/tokens)")
    
    if user_token:
        st.success("Token provided!")
    else:
        st.warning("Please enter a token to unlock the agent.")

# Main Task Area
if user_token:
    query = st.text_input("Task: Ask a question that requires reasoning or search")
    
    if st.button("Execute Agentic Workflow"):
        if query:
            # Construct the Graph dynamically to inject the token
            workflow = StateGraph(AgentState)
            
            # Nodes
            workflow.add_node("router", lambda x: router_node(x, user_token))
            workflow.add_node("search", search_node)
            workflow.add_node("writer", lambda x: writer_node(x, user_token))
            
            # Edges
            workflow.set_entry_point("router")
            workflow.add_conditional_edges(
                "router",
                lambda x: "web" if x["context"] == "SEARCH" else "direct",
                {"web": "search", "direct": "writer"}
            )
            workflow.add_edge("search", "writer")
            workflow.add_edge("writer", END)
            
            app = workflow.compile()
            
            # Execution
            with st.status("Agent Orchestrating...", expanded=True) as status:
                final_state = app.invoke({"question": query, "trace": []})
                
                st.subheader("Process Trace")
                for step in final_state["trace"]:
                    st.write(step)
                
                st.subheader("Final Output")
                st.write(final_state["answer"])
                status.update(label="Task Finished!", state="complete")
        else:
            st.error("Enter a query first.")
else:
    st.info("👈 Enter your Hugging Face token in the sidebar to start practicing.")

# Footer Practice Guide
st.divider()
with st.expander("🎓 Advanced Level Learning Goals"):
    st.markdown("""
    - **LCEL Implementation**: Watch how the `prompt | llm` chain is constructed.
    - **State Control**: Observe how `AgentState` passes data between nodes.
    - **Conditional Routing**: The router node uses the LLM as a logic gate.
    - **Tool Integration**: DuckDuckGo search is used as a dynamic external tool.
    """)