"""
# main.py
# Implementation of the LangGraph reflective agent using deepagents
# Author: Auto-generated for the assignment
# Requires: deepagents, langchain-openai, langgraph, langchain
"""
import os
import asyncio
from typing import TypedDict, Annotated

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

# ---------- 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 implementations ----------
async def draft_answer(state: ReflectState) -> ReflectState:
    """Generate an initial draft answer (5–10 sentences)."""
    prompt = (
        "You are an expert tutor. Answer the following question in 5–10 concise sentences. "
        "Avoid filler and keep it clear.
"
        f"Question: {state['question']}"
    )
    response = await llm.ainvoke([SystemMessage(content="You are a helpful tutor."), HumanMessage(content=prompt)])
    state["draft"] = response.content.strip()
    state["round"] = 1
    return state

async def reflect(state: ReflectState) -> ReflectState:
    """Critique the draft and decide if revision is needed."""
    prompt = (
        "You are a critical reviewer. Evaluate the following answer for completeness, specificity, and lack of filler. "
        "Respond with a verdict of either "ok" or "needs_revision", followed by 2–3 bullet points of constructive feedback.
"
        f"Answer draft:\n{state['draft']}"
    )
    response = await llm.ainvoke([SystemMessage(content="You are a critical reviewer."), HumanMessage(content=prompt)])
    text = response.content.strip()
    # Parse verdict and critique
    if "needs_revision" in text.lower():
        verdict = "needs_revision"
    else:
        verdict = "ok"
    # Extract critique lines after the verdict
    lines = text.splitlines()
    critique_lines = [line for line in lines if line.strip() and line.strip().lower() not in {"ok", "needs_revision"}]
    critique = "\n".join(critique_lines).strip()
    state["critique"] = critique
    state["verdict"] = verdict
    return state

async def rewrite(state: ReflectState) -> ReflectState:
    """Rewrite the draft incorporating the critique."""
    prompt = (
        "You are revising an answer based on the following critique. Produce a new version that addresses the points and remains 5–10 sentences.
"
        f"Original draft:\n{state['draft']}\n\nCritique:\n{state['critique']}"
    )
    response = await llm.ainvoke([SystemMessage(content="You are a revising tutor."), HumanMessage(content=prompt)])
    state["draft"] = response.content.strip()
    state["round"] += 1
    return state

# ---------- Graph construction ----------
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)

    # Entry point
    graph.set_entry_point("draft_answer")

    # Transitions
    graph.add_conditional_edges(
        "draft_answer",
        lambda _: "reflect",
    )
    graph.add_conditional_edges(
        "reflect",
        lambda state: "END" if state["verdict"] == "ok" else "rewrite",
    )
    graph.add_conditional_edges(
        "rewrite",
        lambda state: "END" if state["round"] > state["max_rounds"] else "reflect",
    )

    return graph

# ---------- Tool that runs the graph ----------
@tool
async def answer_question(query: str) -> str:
    """Run the reflective LangGraph to answer a question."""
    graph = build_graph()
    # Initialize state
    state: ReflectState = {
        "question": query,
        "draft": "",
        "critique": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }
    # Run graph
    final_state = await graph.ainvoke(state)
    return final_state["draft"]

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

agent = create_deep_agent(
    model=llm,
    tools=[answer_question],
    backend=backend,
    system_prompt="You are a helpful educational agent. Use the provided tools to answer questions.",
)

# ---------- CLI entry point ----------
async def main():
    question = "Объясни студенту разницу между tool и resource в MCP"
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=question)]},
        {"configurable": {"thread_id": "session-1"}},
    )
    # The tool returns the final answer as the last message content
    print("\nFinal answer:\n", result["messages"][-1].content)

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