import os
import asyncio
from typing import TypedDict, Annotated

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

# ----------------------------------------------------------------------
# LLM configuration (OpenRouter, required by the course)
# ----------------------------------------------------------------------
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 (required by the framework)
# ----------------------------------------------------------------------
backend = CompositeBackend(
    [
        LocalShellBackend(workspace_dir="./workspace"),
        FilesystemBackend(),
    ]
)

# ----------------------------------------------------------------------
# State definition for the reflection loop
# ----------------------------------------------------------------------
class ReflectState(TypedDict):
    question: str
    draft: str
    critique: str
    verdict: str  # "ok" or "needs_revision"
    round: int
    max_rounds: int

# ----------------------------------------------------------------------
# Helper function to build a simple LLM chain for a given prompt
# ----------------------------------------------------------------------
def llm_call(prompt: str, state: ReflectState) -> str:
    """Invoke the LLM with a system prompt and the current state."""
    messages = [
        HumanMessage(content=prompt.format(**state))
    ]
    response = llm.invoke(messages)
    return response.content.strip()

# ----------------------------------------------------------------------
# Node: draft_answer - produce the first answer
# ----------------------------------------------------------------------
def draft_answer(state: ReflectState) -> ReflectState:
    prompt = (
        "You are an expert educator. Answer the following question in 5-10 sentences, "
        "clear and concise, without unnecessary filler. Question: {question}"
    )
    draft = llm_call(prompt, state)
    state["draft"] = draft
    state["round"] = 0
    return state

# ----------------------------------------------------------------------
# Node: reflect - LLM critic evaluates the draft
# ----------------------------------------------------------------------
def reflect(state: ReflectState) -> ReflectState:
    critique_prompt = (
        "You are a reviewer. Evaluate the draft answer provided below. "
        "Assess completeness, concreteness and absence of filler. "
        "Return a verdict ('ok' or 'needs_revision') and list 2-3 short remarks. "
        "Format exactly as:\n"
        "Verdict: <verdict>\n"
        "Critique:\n"
        "- <remark 1>\n"
        "- <remark 2>\n"
        "Draft:\n{draft}"
    )
    critique_raw = llm_call(critique_prompt, state)
    # Parse the structured response
    lines = critique_raw.splitlines()
    verdict_line = next((l for l in lines if l.lower().startswith("verdict:")), "")
    verdict = verdict_line.split(":", 1)[1].strip().lower()
    critique_start = lines.index("Critique:") + 1 if "Critique:" in lines else 0
    critique_items = [l.lstrip("- ").strip() for l in lines[critique_start:] if l.startswith("-")]
    state["verdict"] = verdict
    state["critique"] = "\n".join(critique_items)
    return state

# ----------------------------------------------------------------------
# Node: rewrite - improve the draft based on critique
# ----------------------------------------------------------------------
def rewrite(state: ReflectState) -> ReflectState:
    rewrite_prompt = (
        "You are a writer. Improve the previous draft according to the following critique points. "
        "Make the answer clearer, more concrete and remove any filler. Keep the length 5-10 sentences.\n"
        "Critique:\n{critique}\n\nCurrent draft:\n{draft}"
    )
    new_draft = llm_call(rewrite_prompt, state)
    state["draft"] = new_draft
    state["round"] += 1
    return state

# ----------------------------------------------------------------------
# DESIGN DECISION: Use a pure LangGraph state machine for reflection.
# NECESSITY: The assignment explicitly requires a separate reflect node and
#           iteration via rewrite → reflect, not a try/except retry loop.
# OPTIMALITY: Graph representation makes the flow declarative, guarantees
#           max_rounds enforcement, and isolates each responsibility.
# ALTERNATIVES CONSIDERED: A manual while-loop with try/except was removed
#           because it mixes error handling with logical revision, violating
#           the task specification.
# ----------------------------------------------------------------------
def build_graph() -> StateGraph:
    graph = StateGraph(ReflectState)

    graph.add_node("draft_answer", draft_answer)
    graph.add_node("reflect", reflect)
    graph.add_node("rewrite", rewrite)

    # START → draft_answer
    graph.add_edge(START, "draft_answer")
    # draft_answer → reflect
    graph.add_edge("draft_answer", "reflect")

    # reflect → END if ok
    graph.add_conditional_edges(
        "reflect",
        lambda s: END if s["verdict"] == "ok" else "rewrite",
    )
    # rewrite → reflect (if rounds left)
    def rewrite_condition(s: ReflectState):
        if s["round"] < s["max_rounds"]:
            return "reflect"
        return END

    graph.add_edge("rewrite", rewrite_condition)

    graph.set_entry_point("draft_answer")
    return graph

# ----------------------------------------------------------------------
# DeepAgent wrapper - required by the course
# ----------------------------------------------------------------------
agent = create_deep_agent(
    model=llm,
    tools=[],  # No external tools needed for this assignment
    backend=backend,
    system_prompt="You are a reflective assistant that writes concise answers and improves them based on critique.",
)

# ----------------------------------------------------------------------
# Demo execution
# ----------------------------------------------------------------------
async def main():
    question = "Объясни студенту разницу между tool и resource в MCP"
    initial_state: ReflectState = {
        "question": question,
        "draft": "",
        "critique": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }

    graph = build_graph()
    # Run the graph synchronously (LangGraph supports async, but our nodes are sync)
    final_state = await graph.ainvoke(initial_state, config={"configurable": {"thread_id": "demo-1"}})

    print("=== Final Answer ===")
    print(final_state["draft"])
    print("\n=== Verdict ===")
    print(final_state["verdict"])
    if final_state["critique"]:
        print("\n=== Critique ===")
        print(final_state["critique"])

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