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

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

# ---------- Nodes ----------
async def planning_node(state: PlanningState) -> PlanningState:
    task = state["messages"][-1].content if state["messages"] else ""
    prompt = f"""
You are a planning assistant. Given the following task, break it into 3-6 concrete steps. Return the steps as a JSON array of strings.

Task: {task}

JSON output only:
"""
    response = await llm.ainvoke([HumanMessage(content=prompt)])
    try:
        import json
        plan = json.loads(response.content)
        if not isinstance(plan, list):
            raise ValueError
    except Exception:
        # fallback: simple split by lines
        plan = [line.strip('-•* ') for line in response.content.splitlines() if line.strip()]
    return {
        "plan": plan,
        "current_step": 0,
        "results": [],
    }

async def execution_node(state: PlanningState) -> PlanningState:
    step = state["plan"][state["current_step"]]
    prompt = f"""
You are an executor. Perform the following step and return the result as plain text.

Step: {step}
"""
    response = await llm.ainvoke([HumanMessage(content=prompt)])
    result = response.content.strip()
    new_results = state["results"].copy()
    new_results.append(result)
    return {
        "results": new_results,
        "current_step": state["current_step"] + 1,
    }

# ---------- Graph ----------
graph = StateGraph(PlanningState)
graph.add_node("planning", planning_node)
graph.add_node("execution", execution_node)

# Condition to decide next step
def should_continue(state: PlanningState) -> str:
    if state["current_step"] >= len(state["plan"]):
        return "finish"
    return "execute"

graph.set_conditional_entry_point("planning", should_continue, {
    "execute": "execution",
    "finish": END,
})

# After finish, combine results
async def final_node(state: PlanningState) -> PlanningState:
    summary = "\n".join(state["results"])
    return {"messages": [HumanMessage(content=summary)]}

graph.add_node("final", final_node)
graph.set_entry_point("planning")
graph.add_edge("execution", "should_continue")
graph.add_edge("should_continue", "final", label="finish")

planner = graph.compile()

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

@tool
async def run_planning(task: str) -> str:
    """Run the planning graph for the given task and return the final summary."""
    # Initialize state with the task as a message
    state = {"messages": [HumanMessage(content=task)], "plan": None, "current_step": 0, "results": []}
    result = await planner.ainvoke(state)
    # result contains 'messages' with final summary
    return result["messages"][-1].content

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

async def main():
    task = "Сравни Python и JavaScript"
    response = await agent.ainvoke(
        {"messages": [HumanMessage(content=task)]},
        {"configurable": {"thread_id": "session-1"}},
    )
    print("\n".join(msg.content for msg in response["messages"]))

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