import os
import asyncio
import json
from typing import TypedDict

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langgraph.graph import StateGraph, START, END

from deepagents import create_deep_agent, tool
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend

# LLM configuration - OpenRouter
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,
)

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

# Node: draft_answer
def draft_answer(state: ReflectState) -> ReflectState:
    prompt = f"Write a concise answer (5-10 sentences) to the following question:\n\n{state['question']}"
    response = llm.invoke([HumanMessage(content=prompt)])
    state["draft"] = response.content.strip()
    return state

# Node: reflect
def reflect(state: ReflectState) -> ReflectState:
    prompt = f"""You are a critic evaluating the following draft answer. Provide a verdict ('ok' or 'needs_revision') and 2-3 specific points of improvement. Do not provide the revised answer. Use JSON format:
{{
  "verdict": "ok" | "needs_revision",
  "critique": "list of points"
}}
Draft:
{state['draft']}"""
    response = llm.invoke([HumanMessage(content=prompt)])
    try:
        data = json.loads(response.content)
    except Exception:
        data = {"verdict": "needs_revision", "critique": "Could not parse critique"}
    state["critique"] = data.get("critique", "")
    state["verdict"] = data.get("verdict", "needs_revision")
    return state

# Node: rewrite
def rewrite(state: ReflectState) -> ReflectState:
    prompt = f"""You are revising the draft answer based on the following critique. Produce a revised answer (5-10 sentences). Do not include the critique. Use the critique points to improve clarity, specificity, and remove filler. Draft:\n{state['draft']}\nCritique:\n{state['critique']}"""
    response = llm.invoke([HumanMessage(content=prompt)])
    state["draft"] = response.content.strip()
    state["round"] = state.get("round", 0) + 1
    return state

# Build the graph
def build_graph() -> StateGraph:
    graph = StateGraph(ReflectState)
    graph.add_node("draft_answer", draft_answer)
    graph.add_node("reflect", reflect)
    graph.add_node("rewrite", rewrite)
    graph.set_entry_point("draft_answer")
    graph.add_edge("draft_answer", "reflect")
    graph.add_conditional_edges(
        "reflect",
        lambda s: "ok" if s["verdict"] == "ok" else ("rewrite" if s["round"] < s["max_rounds"] else "END"),
        {
            "ok": END,
            "rewrite": "rewrite",
            "END": END,
        },
    )
    graph.add_edge("rewrite", "reflect")
    return graph

# Tool that runs the graph
def answer_question_tool(question: str, max_rounds: int = 2) -> str:
    graph = build_graph()
    initial_state: ReflectState = {
        "question": question,
        "draft": "",
        "critique": "",
        "verdict": "",
        "round": 0,
        "max_rounds": max_rounds,
    }
    final_state = graph.invoke(initial_state)
    return final_state["draft"]

# DeepAgent tool
@tool
def answer_question(query: str) -> str:
    """Answer a question using a self-reflective process."""
    return answer_question_tool(query)

# Backend for DeepAgent
backend = CompositeBackend(
    [
        LocalShellBackend(workspace_dir="./workspace"),
        FilesystemBackend(),
    ]
)

# Create the DeepAgent
agent = create_deep_agent(
    model=llm,
    tools=[answer_question],
    backend=backend,
    system_prompt="You are an assistant that answers questions using a self-reflective process. Use the tool 'answer_question' to answer the question.",
)

# CLI demo
async def main():
    question = "Объясни студенту разницу между tool и resource в MCP"
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=question)]},
        {"configurable": {"thread_id": "session-1"}},
    )
    print(result["messages"][-1].content)

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