import os
import asyncio
from typing import TypedDict, Annotated

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 CompositeBackend, LocalShellBackend, FilesystemBackend

from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

# ---------- LLM ----------
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 ----------
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

# ---------- Planning State ----------
class PlanningState(TypedDict):
    messages: Annotated[list, add_messages]
    plan: list[str]
    current_step: int
    results: list[str]

# ---------- Planner Node ----------
async def planner_node(state: PlanningState):
    task = state["messages"][-1].content
    prompt = (
        "You are a task planner.\n"
        "Task: {task}\n"
        "Break the task into 3–6 concrete steps.\n"
        "Return the steps as a numbered list or JSON array.\n"
        "Do not include any other text."
    ).format(task=task)
    plan_text = await llm.ainvoke([HumanMessage(content=prompt)])
    plan_str = plan_text.content.strip()
    # Try to parse JSON
    plan = []
    try:
        import json
        plan = json.loads(plan_str)
        if not isinstance(plan, list):
            raise ValueError
    except Exception:
        # Fallback to numbered list parsing
        import re
        plan = [line.strip() for line in plan_str.splitlines() if re.match(r"^\s*\d+\.", line)]
    return {
        "plan": plan,
        "current_step": 0,
        "results": [],
    }

# ---------- Executor Node ----------
async def executor_node(state: PlanningState):
    step = state["plan"][state["current_step"]]
    prompt = (
        "You are an executor.\n"
        "Step: {step}\n"
        "Provide a concise result for this step."
    ).format(step=step)
    result_text = await llm.ainvoke([HumanMessage(content=prompt)])
    result = result_text.content.strip()
    new_results = state["results"] + [result]
    return {
        "results": new_results,
        "current_step": state["current_step"] + 1,
    }

# ---------- Should Continue ----------
def should_continue(state: PlanningState):
    if state["current_step"] >= len(state["plan"]):
        return "finish"
    return "execute"

# ---------- Graph ----------
graph = StateGraph(PlanningState)
graph.add_node("planner", planner_node)
graph.add_node("executor", executor_node)
graph.add_conditional_edges("planner", lambda _: "execute")
graph.add_conditional_edges("executor", should_continue)
graph.set_entry_point("planner")
graph.add_edge("execute", "executor")
graph.add_edge("finish", END)
planner_chain = graph.compile()

# ---------- Tool that runs the planner graph ----------
@tool
def run_planner(task: str) -> str:
    """Run the planning graph on the given task and return the final summary."""
    # Initialize state with the task as a message
    init_state = {"messages": [HumanMessage(content=task)], "plan": [], "current_step": 0, "results": []}
    final_state = planner_chain.invoke(init_state)
    # Build final output
    plan_lines = [f"{i+1}. {step}" for i, step in enumerate(final_state["plan"])]
    results = final_state["results"]
    summary = "\n".join(results)
    return (
        f"План:\n" + "\n".join(plan_lines) + "\n\n[Шаги]" + "\n".join([f"[Шаг {i+1}] {r}" for i, r in enumerate(results)]) + "\n\nИтог: " + summary
    )

# ---------- Deep Agent ----------
agent = create_deep_agent(
    model=llm,
    tools=[run_planner],
    backend=backend,
    system_prompt="You are a helpful agent that can plan and execute tasks.",
)

# ---------- Demo ----------
async def main():
    task = "Сравни Python и JavaScript"
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=task)]},
        {"configurable": {"thread_id": "session-1"}},
    )
    print(result["messages"][-1].content)

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