import os
import asyncio
import json
from typing import TypedDict, Annotated, List

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage
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,
)

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

# ---------- Tools (optional example) ----------
@tool
def echo_tool(text: str) -> str:
    """Return the given text unchanged."""
    return text

# ---------- Deep Agent ----------
deep_agent = create_deep_agent(
    model=llm,
    tools=[echo_tool],
    backend=backend,
    system_prompt="You are a helpful planning agent.",
)

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

# ---------- Planner Node ----------
def planner_node(state: PlanningState) -> PlanningState:
    prompt = f"""You are given a task. Break it into 3-6 concrete steps.
Return the plan as a JSON array of strings, e.g. ["step 1", "step 2", ...].
Task: {state['task']}"""
    response = deep_agent.invoke(
        {"messages": [HumanMessage(content=prompt)]},
        {"configurable": {"thread_id": "planner"}},
    )
    content = response["messages"][-1].content
    try:
        plan = json.loads(content)
        if not isinstance(plan, list):
            raise ValueError
    except Exception:
        # Fallback: try to extract lines starting with numbers
        lines = [line.strip() for line in content.splitlines() if line.strip()]
        plan = [line.split(".", 1)[-1].strip() for line in lines if line[0].isdigit()]
    state["plan"] = plan
    state["current_step"] = 0
    state["results"] = []
    return state

# ---------- Execution Node ----------
def execution_node(state: PlanningState) -> PlanningState:
    step_idx = state["current_step"]
    step_instruction = state["plan"][step_idx]
    prompt = f"""You are executing step {step_idx + 1} of a plan.
Task: {state['task']}
Step: {step_instruction}
Provide a concise answer for this step."""
    response = deep_agent.invoke(
        {"messages": [HumanMessage(content=prompt)]},
        {"configurable": {"thread_id": f"exec-{step_idx}"}},
    )
    result = response["messages"][-1].content
    state["results"].append(f"[Step {step_idx + 1}] {result}")
    state["current_step"] += 1
    return state

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

# ---------- Build Graph ----------
graph = StateGraph(PlanningState)

graph.add_node("planning", planner_node)
graph.add_node("execution", execution_node)

graph.add_edge(START, "planning")
graph.add_edge("planning", "execution")
graph.add_conditional_edges(
    "execution",
    should_continue,
    {"execute": "execution", "finish": END},
)

graph.set_entry_point("planning")
graph = graph.compile()

# ---------- Runner ----------
async def run_planning_agent(task: str) -> str:
    # Initialise state
    state: PlanningState = {
        "messages": [],
        "task": task,
        "plan": None,
        "current_step": 0,
        "results": [],
    }
    # Run graph
    async for event in graph.astream(state):
        # We only need final state
        pass
    final_state = event
    # Build final answer
    plan_text = "\n".join(f"{i+1}. {step}" for i, step in enumerate(final_state["plan"]))
    steps_text = "\n".join(final_state["results"])
    summary_prompt = f"""You have completed all steps of the following task.
Task: {task}
Plan:
{plan_text}
Steps results:
{steps_text}
Provide a concise final summary."""
    summary_resp = deep_agent.invoke(
        {"messages": [HumanMessage(content=summary_prompt)]},
        {"configurable": {"thread_id": "summary"}},
    )
    summary = summary_resp["messages"][-1].content
    output = f"""Задача: {task}

План:
{plan_text}

{steps_text}
Итог: {summary}
"""
    return output

# ---------- Main ----------
if __name__ == "__main__":
    example_task = "Сравни Python и JavaScript"
    result = asyncio.run(run_planning_agent(example_task))
    print(result)