"""
# main.py
# LangGraph agent with reflection and rewrite loop
# Author: ChatGPT
# Requirements: langgraph, langchain-openai, deepagents

import os
import asyncio
from typing import TypedDict

from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END
from deepagents import create_deep_agent
from deepagents.backends import CompositeBackend, LocalShellBackend

# LLM setup
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 = CompositeBackend(
    default=LocalShellBackend(root_dir="./workspace", virtual_mode=True, inherit_env=True),
    routes={},
)
agent = create_deep_agent(
    model=llm,
    tools=[],
    backend=backend,
    system_prompt="You are a helpful assistant.",
)

class ReflectState(TypedDict):
    question: str
    draft: str
    critique: str
    verdict: str
    round: int
    max_rounds: int

async def draft_answer(state: ReflectState) -> ReflectState:
    prompt = f"Write a concise answer (5–10 sentences) to the following question:\n\n{state['question']}"
    response = await agent.ainvoke({"messages": ["Human: " + prompt]}, {"configurable": {"thread_id": "draft"}})
    draft = response["messages"][-1].content
    state["draft"] = draft
    return state

async def reflect(state: ReflectState) -> ReflectState:
    prompt = (
        "You are a critical reviewer.\n"
        "Evaluate the following draft answer for completeness, specificity, and lack of filler.\n"
        "Provide a verdict: 'ok' if the answer is satisfactory, otherwise 'needs_revision'.\n"
        "If revision is needed, give 2–3 concrete points for improvement.\n"
        "Respond in JSON with keys 'verdict' and 'critique'.\n"
        f"Draft: {state['draft']}"
    )
    response = await agent.ainvoke({"messages": ["Human: " + prompt]}, {"configurable": {"thread_id": "reflect"}})
    import json
    try:
        data = json.loads(response["messages"][-1].content)
        verdict = data.get("verdict", "needs_revision")
        critique = data.get("critique", "")
    except Exception:
        verdict = "needs_revision"
        critique = "Could not parse critique."
    state["verdict"] = verdict
    state["critique"] = critique
    return state

async def rewrite(state: ReflectState) -> ReflectState:
    prompt = (
        "You are revising the following draft answer based on the critique.\n"
        "Make the answer clearer, more specific, and remove any filler.\n"
        "Do not add new information beyond what is already in the draft.\n"
        f"Draft: {state['draft']}\n"
        f"Critique: {state['critique']}"
    )
    response = await agent.ainvoke({"messages": ["Human: " + prompt]}, {"configurable": {"thread_id": "rewrite"}})
    new_draft = response["messages"][-1].content
    state["draft"] = new_draft
    state["round"] += 1
    return state


def build_graph() -> StateGraph[ReflectState]:
    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",
        lambda state: state["verdict"],
        {"ok": END, "needs_revision": "rewrite"},
    )
    graph.add_conditional_edges(
        "rewrite",
        lambda state: "rewrite" if state["round"] < state["max_rounds"] else END,
        {"rewrite": "reflect", END: END},
    )
    return graph

async def main():
    question = "Объясни студенту разницу между tool и resource в MCP"
    initial_state: ReflectState = {
        "question": question,
        "draft": "",
        "critique": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }
    graph = build_graph()
    result = await graph.ainvoke(initial_state)
    print("\n--- Final Answer ---")
    print(result["draft"])
    print("\n--- Final Critique ---")
    print(result["critique"])

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