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

from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.output_parsers import JsonOutputParser
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, can be used by LLM) ----------
@tool
def search_web(query: str) -> str:
    """Search the web for the given query and return a short summary."""
    # Placeholder implementation - in real use you could call an API.
    return f"Search result for '{query}' (mock)."

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

# ---------- 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 under the key "plan".
Task: {state['task']}"""
    response = deep_agent.invoke(
        {"messages": [HumanMessage(content=prompt)]},
        {"configurable": {"thread_id": f"planner-{state['task'][:10]}"}},
    )
    content = response["messages"][-1].content
    parser = JsonOutputParser()
    try:
        plan = parser.parse(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()]
    return {
        "messages": state["messages"],
        "task": state["task"],
        "plan": plan,
        "current_step": 0,
        "results": [],
    }

# ---------- Executor Node ----------
def executor_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.
Step description: {step_instruction}
Provide a concise answer for this step."""
    response = deep_agent.invoke(
        {"messages": [HumanMessage(content=prompt)]},
        {"configurable": {"thread_id": f"executor-{state['task'][:10]}"}},
    )
    result = response["messages"][-1].content
    new_results = state["results"] + [f"[Step {step_idx + 1}] {result}"]
    return {
        "messages": state["messages"],
        "task": state["task"],
        "plan": state["plan"],
        "current_step": step_idx + 1,
        "results": new_results,
    }

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

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

graph.add_node("planning", planner_node)
graph.add_node("execution", executor_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")
app = graph.compile()

# ---------- Demo ----------
async def run_demo(task: str):
    # Initialize empty state
    init_state: PlanningState = {
        "messages": [],
        "task": task,
        "plan": None,
        "current_step": 0,
        "results": [],
    }
    async for event in app.astream(
        init_state,
        {"configurable": {"thread_id": "demo-session"}},
    ):
        # We only care about final state
        pass
    final_state = event
    print(f"Задача: {task}\n")
    print("План:")
    for i, step in enumerate(final_state["plan"], 1):
        print(f"{i}. {step}")
    print()
    for res in final_state["results"]:
        print(res)
    print("\nИтог:")
    summary_prompt = f"""Based on the following step results, provide a concise final summary.

Results:
{chr(10).join(final_state['results'])}"""
    summary_resp = deep_agent.invoke(
        {"messages": [HumanMessage(content=summary_prompt)]},
        {"configurable": {"thread_id": "summary"}},
    )
    print(summary_resp["messages"][-1].content)

if __name__ == "__main__":
    demo_task = "Сравни Python и JavaScript"
    asyncio.run(run_demo(demo_task))