from typing import TypedDict, List, Optional
from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from langchain_community.tools import DuckDuckGoSearchRun
from langchain.agents import tool

search_tool = DuckDuckGoSearchRun()

@tool
def search_web(query: str) -> str:
    """Search the web using DuckDuckGo and return the top result."""
    return search_tool.run(query)

class AgentState(TypedDict):
    question: str
    thoughts: List[str]
    tool_results: List[str]
    answer: Optional[str]

llm = ChatOpenAI(model="gpt-4", temperature=0.5)

def plan(state: AgentState) -> AgentState:
    prompt = f"""
You are an intelligent assistant. Here is the question:
{state['question']}

So far, these are your thoughts:
{state['thoughts']}

What is the next best step? Should you:
- Search the web (if more info is needed),
- Answer the question (if confident), or
- Think more before deciding?

Provide a clear next step starting with one of these prefixes exactly:
'Search:', 'Answer:', or 'Think:'
Then explain your reasoning.
"""
    thought = llm.invoke(prompt).content.strip()
    state['thoughts'].append(thought)
    return state

def act(state: AgentState) -> AgentState:
    latest = state['thoughts'][-1].strip()
    lower = latest.lower()
    if lower.startswith("search:"):
        query = latest[len("Search:"):].strip()
        if query:
            result = search_web.run(query)
            state['tool_results'].append(result)
    elif lower.startswith("answer:"):
        # Agent thinks it has the answer
        state['tool_results'].append("Answer ready")
    else:
        # Thinking or other - no tool used
        state['tool_results'].append("No tool used")
    return state

def observe(state: AgentState) -> AgentState:
    obs = state['tool_results'][-1] if state['tool_results'] else "Nothing found"
    state['thoughts'].append(f"Observed: {obs}")
    return state

def decide(state: AgentState) -> str:
    # Allow more thinking steps (e.g. 5) before summarizing
    return END if len(state["thoughts"]) >= 5 else "plan"

def summarize(state: AgentState) -> AgentState:
    prompt = f"""
You have gathered information and thoughts:
{state['thoughts']}

Based on all this, give a clear, concise, and final answer to the question:
{state['question']}
"""
    answer = llm.invoke(prompt).content.strip()
    state["answer"] = answer
    return state

workflow = StateGraph(AgentState)
workflow.add_node("plan", plan)
workflow.add_node("act", act)
workflow.add_node("observe", observe)
workflow.add_node("decide", decide)
workflow.add_node("summarize", summarize)

workflow.set_entry_point("plan")
workflow.add_edge("plan", "act")
workflow.add_edge("act", "observe")
workflow.add_edge("observe", "decide")
def route_decision(state: AgentState) -> str:
    return "plan" if len(state["thoughts"]) < 3 else "summarize"

workflow.add_conditional_edges("decide", {
    "plan": route_decision
})

workflow.add_edge("summarize", END)

agent = workflow.compile()