import os
import json
import asyncio
from typing import TypedDict

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend
from langgraph.graph import StateGraph, START, END

# ---------- LLM ----------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.0,
)

# ---------- Backend & Agent ----------
backend = FilesystemBackend()
agent = create_deep_agent(
    model=llm,
    tools=[],
    backend=backend,
    system_prompt="You are a helpful assistant.",
)

# ---------- State ----------
class ReflectState(TypedDict):
    question: str
    draft: str
    critique: str
    verdict: str  # ok | needs_revision
    round: int
    max_rounds: int

# ---------- Nodes ----------
async def draft_answer(state: ReflectState) -> ReflectState:
    prompt = f"Write a brief answer (5-10 sentences) to the following question: {state['question']}"
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=prompt)], "configurable": {"thread_id": "draft"}},
        {},
    )
    state["draft"] = result["messages"][-1].content
    return state

async def reflect(state: ReflectState) -> ReflectState:
    prompt = (
        f"Critique the following answer. Provide verdict ok or needs_revision and 2-3 points of critique in JSON format with keys verdict and critique.\nAnswer: {state['draft']}"
    )
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=prompt)], "configurable": {"thread_id": "reflect"}},
        {},
    )
    content = result["messages"][-1].content
    try:
        data = json.loads(content)
        state["verdict"] = data.get("verdict", "").lower()
        state["critique"] = data.get("critique", "")
    except json.JSONDecodeError:
        state["verdict"] = "needs_revision"
        state["critique"] = content
    return state

async def rewrite(state: ReflectState) -> ReflectState:
    prompt = (
        f"Rewrite the answer to improve it based on the critique: {state['critique']}\nPrevious draft: {state['draft']}\nProvide the improved answer."
    )
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=prompt)], "configurable": {"thread_id": "rewrite"}},
        {},
    )
    state["draft"] = result["messages"][-1].content
    state["round"] += 1
    return state

# ---------- Conditional Edge ----------
def reflect_cond(state: ReflectState):
    if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"]:
        return "rewrite"
    return "end"

# ---------- Graph ----------

graph = StateGraph(ReflectState)
graph.add_node("draft_answer", draft_answer)
graph.add_node("reflect", reflect)
graph.add_node("rewrite", rewrite)

graph.add_edge(START, "draft_answer")
graph.add_edge("draft_answer", "reflect")
graph.add_conditional_edges("reflect", reflect_cond, {"rewrite": "rewrite", "end": END})
graph.add_edge("rewrite", "reflect")

graph.set_entry_point("draft_answer")
graph.set_finish_point(END)

executor = graph.compile()

# ---------- CLI ----------
async def main():
    question = input("Enter a question: ")
    initial_state: ReflectState = {
        "question": question,
        "draft": "",
        "critique": "",
        "verdict": "",
        "round": 1,
        "max_rounds": 2,
    }
    final_state = await executor(initial_state)
    print("\nFinal answer:\n")
    print(final_state["draft"])

if __name__ == "__main__":
    asyncio.run(main())