import os
import asyncio
from typing import TypedDict

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

# 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,
)

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

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

def draft_answer(state: ReflectState) -> ReflectState:
    prompt = f"Write a short answer (5-10 sentences) to the following question: {state['question']}"
    msg = HumanMessage(content=prompt)
    response = llm.invoke([msg])
    state["draft"] = response.content
    state["round"] = 0
    return state

def reflect(state: ReflectState) -> ReflectState:
    prompt = (
        f"You are a critic. Evaluate the following draft answer:\n{state['draft']}\n\n"
        "Provide verdict 'ok' or 'needs_revision' and 2-3 points of critique."
    )
    msg = HumanMessage(content=prompt)
    response = llm.invoke([msg])
    text = response.content.strip()
    verdict = "ok"
    critique = ""
    if "needs_revision" in text.lower():
        verdict = "needs_revision"
    # Extract critique after the word 'Critique:' if present
    lower_text = text.lower()
    if "critique:" in lower_text:
        idx = lower_text.find("critique:")
        critique = text[idx + len("critique:") :].strip()
    else:
        parts = text.split("\n")
        if len(parts) > 1:
            critique = "\n".join(parts[1:]).strip()
    state["verdict"] = verdict
    state["critique"] = critique
    return state

def rewrite(state: ReflectState) -> ReflectState:
    prompt = (
        f"Rewrite the draft answer to address the following critique:\n{state['critique']}\n\n"
        "Keep the answer short (5-10 sentences)."
    )
    msg = HumanMessage(content=prompt)
    response = llm.invoke([msg])
    state["draft"] = response.content
    state["round"] += 1
    return state

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 "stop")
    ),
    {"ok": END, "rewrite": "rewrite", "stop": END},
)
graph.add_edge("rewrite", "reflect")
compiled_graph = graph.compile()

def run_reflection(question: str) -> str:
    state: ReflectState = {
        "question": question,
        "draft": "",
        "critique": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }
    final_state = compiled_graph.invoke(state)
    output = (
        f"Draft:\n{final_state['draft']}\n\n"
        f"Critique:\n{final_state['critique']}\n\n"
        f"Verdict: {final_state['verdict']}\n\n"
        f"Final answer:\n{final_state['draft']}\n"
    )
    return output

# ---------- DeepAgents tool ----------
@tool
def answer_with_reflection(question: str) -> str:
    """Generate a short answer with self-reflection and rewrite if needed."""
    return run_reflection(question)

# ---------- DeepAgents agent ----------
agent = create_deep_agent(
    model=llm,
    tools=[answer_with_reflection],
    backend=backend,
    system_prompt=(
        "You are a helpful agent that writes short answers and self-reflects. "
        "Use the tool 'answer_with_reflection' to answer questions."
    ),
)

# ---------- Demo ----------
async def main():
    question = (
        "Explain to a student the difference between tool and resource in 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())