from pydantic import BaseModel
from typing import Literal

class PlanStep(BaseModel):
    step_id: str
    description: str
    tool: str
    tool_args: dict
    depends_on: list[str]
    status: Literal["pending", "running", "done", "failed"] = "pending"

class Plan(BaseModel):
    goal: str
    steps: list[PlanStep]

def generate_plan(goal: str, available_tools: list[str], llm_client) -> Plan:
    """Phase 1: ask the LLM to decompose the goal into ordered steps."""
    system = (
        "You are a planning agent. Given a goal and available tools, "
        "produce a JSON execution plan with ordered steps. "
        "Each step must specify exactly which tool to call and with what arguments."
    )
    response = llm_client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": f"Goal: {goal}\nTools: {available_tools}"},
        ],
        response_format={"type": "json_object"},
    )
    return Plan.model_validate_json(response.choices[0].message.content)


def execute_plan(plan: Plan, tool_registry: dict, llm_client) -> str:
    """Phase 2: execute each step in dependency order; replan on failure."""
    results: dict[str, str] = {}

    for step in plan.steps:
        # Check dependencies are complete
        if not all(results.get(dep) for dep in step.depends_on):
            continue   # will be retried after dependencies complete

        step.status = "running"
        result = dispatch_tool_call(tool_registry, step.tool, step.tool_args)

        if '"error"' in result:
            step.status = "failed"
            plan = replan(plan, step, result, llm_client)   # dynamic replanning
        else:
            step.status = "done"
            results[step.step_id] = result

    return results.get(plan.steps[-1].step_id, "Plan did not complete")