from typing import TypedDict, Annotated, List, Literal
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langgraph.prebuilt import ToolNode
from langgraph.checkpoint.memory import MemorySaver

# --- State ---
class ResearchState(TypedDict):
    messages: Annotated[List[BaseMessage], add_messages]
    topic: str
    research_plan: list[str]
    research_findings: list[str]
    draft: str
    critique: str
    quality_score: float
    iteration: int
    phase: str

# --- Tools ---
@tool
def web_search(query: str) -> str:
    """Search the web for information."""
    return f"Search results for '{query}': [Simulated results about {query}]"

@tool
def read_document(url: str) -> str:
    """Read content from a URL."""
    return f"Content from {url}: [Simulated document content]"

tools = [web_search, read_document]

# --- LLM ---
llm = ChatOpenAI(model="gpt-4o", temperature=0)

# --- Agent Nodes ---
def planner_agent(state: ResearchState):
    """Create a research plan."""
    prompt = [
        SystemMessage(content="""You are a research planner. Given a topic, create a step-by-step research plan.
Return ONLY a numbered list of research steps."""),
        HumanMessage(content=f"Create a research plan for: {state['topic']}")
    ]
    response = llm.invoke(prompt)
    plan = [line.strip() for line in response.content.split("\n") if line.strip()]
    return {"research_plan": plan, "phase": "planning"}

def researcher_agent(state: ResearchState):
    """Execute research using tools."""
    tools_with_search = ToolNode(tools)
    
    prompt = [
        SystemMessage(content="""You are a researcher. Search for information on the topic.
Use the web_search tool to find relevant information.
Return your findings as bullet points."""),
        HumanMessage(content=f"Research topic: {state['topic']}\nPlan: {state['research_plan']}")
    ]
    
    response = llm.bind_tools(tools).invoke(prompt)
    
    # Simulate findings
    findings = [
        f"Key finding 1 about {state['topic']}",
        f"Key finding 2 about {state['topic']}",
        f"Key finding 3 about {state['topic']}",
    ]
    
    return {"research_findings": findings, "phase": "research"}

def writer_agent(state: ResearchState):
    """Write a comprehensive report."""
    prompt = [
        SystemMessage(content="""You are a technical writer. Write a comprehensive report based on the research findings.
Include: introduction, key findings, analysis, and conclusion."""),
        HumanMessage(content=f"""Write a report on: {state['topic']}

Research findings:
{chr(10).join(state['research_findings'])}

Previous draft (if any): {state.get('draft', 'None')}
Critique (if any): {state.get('critique', 'None')}""")
    ]
    response = llm.invoke(prompt)
    return {"draft": response.content, "phase": "writing"}

def critic_agent(state: ResearchState):
    """Review the draft and provide feedback."""
    prompt = [
        SystemMessage(content="""You are a quality critic. Review the draft and provide:
1. Quality score (0.0 to 1.0)
2. Specific feedback for improvement
3. Whether the draft is ready (score >= 0.8 = ready)"""),
        HumanMessage(content=f"""Review this draft:

{state['draft']}""")
    ]
    response = llm.invoke(prompt)
    
    # Extract score (simplified)
    score = 0.9 if "excellent" in response.content.lower() else 0.6
    
    return {
        "critique": response.content,
        "quality_score": score,
        "iteration": state.get("iteration", 0) + 1,
        "phase": "review"
    }

def supervisor_router(state: ResearchState) -> Literal["planner", "researcher", "writer", "critic", "end"]:
    """Route to the right agent based on current phase."""
    phase = state.get("phase", "")
    iteration = state.get("iteration", 0)
    
    if phase == "" or phase == "start":
        return "planner"
    elif phase == "planning":
        return "researcher"
    elif phase == "research":
        return "writer"
    elif phase == "writing":
        return "critic"
    elif phase == "review":
        if state.get("quality_score", 0) >= 0.8 or iteration >= 3:
            return "end"
        return "writer"  # Revise
    return "end"

# --- Build Graph ---
workflow = StateGraph(ResearchState)

# Add all agent nodes
workflow.add_node("planner", planner_agent)
workflow.add_node("researcher", researcher_agent)
workflow.add_node("writer", writer_agent)
workflow.add_node("critic", critic_agent)

# Entry point
workflow.set_entry_point("planner")

# Add conditional routing
workflow.add_conditional_edges(
    "planner",
    lambda s: "researcher" if s.get("research_plan") else "end",
    {"researcher": "researcher", "end": END}
)
workflow.add_conditional_edges(
    "researcher",
    lambda s: "writer" if s.get("research_findings") else "end",
    {"writer": "writer", "end": END}
)
workflow.add_conditional_edges(
    "writer",
    lambda s: "critic",
    {"critic": "critic"}
)
workflow.add_conditional_edges(
    "critic",
    lambda s: "end" if s.get("quality_score", 0) >= 0.8 or s.get("iteration", 0) >= 3 else "writer",
    {"end": END, "writer": "writer"}
)

# Compile with checkpointing
memory = MemorySaver()
app = workflow.compile(checkpointer=memory)

# --- Run ---
config = {"configurable": {"thread_id": "research-1"}}
result = app.invoke(
    {"messages": [HumanMessage(content="Research the impact of AI on healthcare")], "topic": "AI in Healthcare", "phase": "start"},
    config
)

print("Final Report:")
print(result["draft"])
print(f"\nQuality Score: {result['quality_score']}")
print(f"Iterations: {result['iteration']}")