"""Plan-and-execute orchestration primitives."""
from __future__ import annotations

import asyncio
from dataclasses import asdict, dataclass, field
import inspect
import json
import re
import time
from typing import Any, AsyncIterator, Callable

from pydantic import BaseModel, Field, ValidationError

from logos_adk.runnable import Runnable
from logos_adk.types import Response, StreamDone, StreamChunk, TextChunk, Usage


class PlanStepSchema(BaseModel):
    """Structured planner output for a single step."""

    id: str = Field(min_length=2, pattern=r"^E\d+$")
    instruction: str = Field(min_length=1)
    depends_on: list[str] = Field(default_factory=list)


class ExecutionPlanSchema(BaseModel):
    """Structured planner output."""

    steps: list[PlanStepSchema] = Field(min_length=1)


class ReWOOPlanStepSchema(BaseModel):
    """Formal ReWOO-style planner output for a single step."""

    id: str = Field(min_length=2, pattern=r"^E\d+$")
    instruction: str = Field(min_length=1)
    depends_on: list[str] = Field(default_factory=list)
    tool: str | None = Field(default=None, min_length=1)
    input_template: str | None = None
    tool_input_schema: dict[str, Any] | None = None
    arg_mapping: dict[str, Any] | None = None
    output_key: str | None = Field(
        default=None,
        pattern=r"^[A-Za-z_][A-Za-z0-9_]*$",
    )


class ReWOOExecutionPlanSchema(BaseModel):
    """Formal ReWOO-style structured planner output."""

    steps: list[ReWOOPlanStepSchema] = Field(min_length=1)


@dataclass
class PlanStep:
    """Executable step inside a plan."""

    id: str
    instruction: str
    depends_on: list[str] = field(default_factory=list)
    tool: str | None = None
    input_template: str | None = None
    tool_input_schema: dict[str, Any] | None = None
    arg_mapping: dict[str, Any] | None = None
    output_key: str | None = None


@dataclass
class ExecutionPlan:
    """Parsed execution plan."""

    steps: list[PlanStep]
    format: str = "structured"


@dataclass
class ExecutedPlanStep:
    """Recorded result of one executed step."""

    id: str
    instruction: str
    resolved_instruction: str
    tool: str | None
    input_template: str | None
    resolved_input: str | None
    tool_input_schema: dict[str, Any] | None
    arg_mapping: dict[str, Any] | None
    resolved_arguments: dict[str, Any] | None
    output_key: str | None
    depends_on: list[str]
    output: str
    duration_ms: float
    executor_name: str


PlanComponent = Runnable | Callable[..., Any]
_STEP_ID_RE = re.compile(r"^E(\d+)$")


class PlanAndExecute(Runnable):
    """Run a planner, execute the plan step-by-step, then synthesize a final answer."""

    def __init__(
        self,
        name: str,
        *,
        planner: PlanComponent,
        executor: PlanComponent,
        solver: PlanComponent | None = None,
        max_steps: int = 8,
        parallel_execution: bool = False,
        max_parallel_steps: int | None = None,
        planner_output_schema: type[BaseModel] | None = None,
    ) -> None:
        self._name = name
        self._planner = planner
        self._executor = executor
        self._solver = solver or executor
        self._max_steps = max(1, int(max_steps))
        self._parallel_execution = bool(parallel_execution)
        self._max_parallel_steps = (
            max(1, int(max_parallel_steps))
            if max_parallel_steps is not None
            else None
        )
        self._planner_output_schema = planner_output_schema or ReWOOExecutionPlanSchema

    @property
    def name(self) -> str:
        return self._name

    async def run(self, prompt: str, *, session_id: str | None = None, **kwargs: Any) -> Response:
        planner_response = await _call_component(
            self._planner,
            self._build_planner_prompt(prompt),
            session_id=_child_session_id(session_id, "planner"),
            output_schema=self._planner_output_schema,
            **kwargs,
        )
        plan = self._parse_plan(planner_response)
        executed_steps = await self._execute_plan(
            prompt,
            plan,
            session_id=session_id,
            **kwargs,
        )
        solver_response = await _call_component(
            self._solver,
            self._build_solver_prompt(prompt, plan, executed_steps),
            session_id=_child_session_id(session_id, "solver"),
            **kwargs,
        )
        raw = dict(solver_response.raw)
        raw["plan_and_execute"] = {
            "task": prompt,
            "planner": _component_name(self._planner),
            "executor": _component_name(self._executor),
            "solver": _component_name(self._solver),
            "parallel_execution": self._parallel_execution,
            "max_parallel_steps": self._max_parallel_steps,
            "planner_output_schema": self._planner_output_schema.__name__,
            "plan": {
                "format": plan.format,
                "steps": [asdict(step) for step in plan.steps],
            },
            "planner_response": planner_response.content,
            "executed_steps": [asdict(step) for step in executed_steps],
        }
        return Response(
            content=solver_response.content,
            tool_calls=solver_response.tool_calls,
            usage=_aggregate_usage([planner_response, *self._execution_responses(executed_steps), solver_response]),
            raw=raw,
            reasoning=solver_response.reasoning,
        )

    async def stream(self, prompt: str, **kwargs: Any) -> AsyncIterator[StreamChunk]:
        response = await self.run(prompt, session_id=kwargs.pop("session_id", None), **kwargs)
        if response.content:
            yield TextChunk(text=response.content, reasoning=response.reasoning)
        yield StreamDone(response=response)

    def _parse_plan(self, response: Response) -> ExecutionPlan:
        if hasattr(response, "parsed") and response.parsed is not None:
            parsed = response.parsed
            if isinstance(parsed, (ExecutionPlanSchema, ReWOOExecutionPlanSchema)):
                return self._validate_plan(self._build_structured_plan(parsed.steps))

        content = _strip_code_fences(response.content)
        try:
            payload = json.loads(content)
        except json.JSONDecodeError:
            return self._validate_plan(self._parse_fallback_plan(content))

        try:
            structured = self._planner_output_schema.model_validate(payload)
        except ValidationError:
            if isinstance(payload, dict) and "steps" in payload:
                steps_payload = payload.get("steps")
                if not isinstance(steps_payload, list) or not steps_payload:
                    raise ValueError("Execution plan JSON must contain a non-empty 'steps' list")
                return self._validate_plan(
                    ExecutionPlan(
                        steps=[self._coerce_plan_step(step) for step in steps_payload if isinstance(step, dict)],
                        format="structured",
                    )
                )
            return self._validate_plan(self._parse_fallback_plan(content))

        return self._validate_plan(self._build_structured_plan(structured.steps))

    def _build_structured_plan(self, steps: list[Any]) -> ExecutionPlan:
        return ExecutionPlan(
            steps=[
                PlanStep(
                    id=step.id.strip(),
                    instruction=step.instruction.strip(),
                    depends_on=[item.strip() for item in step.depends_on if item.strip()],
                    tool=getattr(step, "tool", None),
                    input_template=getattr(step, "input_template", None),
                    tool_input_schema=getattr(step, "tool_input_schema", None),
                    arg_mapping=getattr(step, "arg_mapping", None),
                    output_key=getattr(step, "output_key", None),
                )
                for step in steps
            ],
            format="structured",
        )

    def _coerce_plan_step(self, payload: dict[str, Any]) -> PlanStep:
        return PlanStep(
            id=str(payload.get("id", "")).strip(),
            instruction=str(payload.get("instruction", "")).strip(),
            depends_on=[
                str(item).strip()
                for item in (payload.get("depends_on") or [])
                if str(item).strip()
            ],
            tool=str(payload.get("tool", "")).strip() or None,
            input_template=(
                str(payload.get("input_template")).strip()
                if payload.get("input_template") is not None
                else None
            ),
            tool_input_schema=(
                payload.get("tool_input_schema")
                if isinstance(payload.get("tool_input_schema"), dict)
                else None
            ),
            arg_mapping=(
                payload.get("arg_mapping")
                if isinstance(payload.get("arg_mapping"), dict)
                else None
            ),
            output_key=str(payload.get("output_key", "")).strip() or None,
        )

    def _parse_fallback_plan(self, content: str) -> ExecutionPlan:
        lines = [line.strip() for line in content.splitlines() if line.strip()]
        steps: list[PlanStep] = []
        for line in lines:
            match = re.match(r"^(?:step\s*)?(?:\d+|e\d+)[\)\.\:\-]\s*(.+)$", line, flags=re.IGNORECASE)
            if match:
                instruction = match.group(1).strip()
            elif line.lower().startswith(("plan", "steps")):
                continue
            else:
                instruction = line.lstrip("-* ").strip()
            if instruction:
                steps.append(PlanStep(id=f"E{len(steps) + 1}", instruction=instruction))
        if not steps:
            raise ValueError("Planner did not return any executable steps")
        return ExecutionPlan(steps=steps, format="fallback")

    def _validate_plan(self, plan: ExecutionPlan) -> ExecutionPlan:
        if not plan.steps:
            raise ValueError("Execution plan must contain at least one step")
        if len(plan.steps) > self._max_steps:
            raise ValueError(
                f"Execution plan exceeds max_steps={self._max_steps}: got {len(plan.steps)}"
            )
        seen: set[str] = set()
        seen_output_keys: set[str] = set()
        for index, step in enumerate(plan.steps, start=1):
            if not step.id:
                raise ValueError("Execution plan contains a step with an empty id")
            match = _STEP_ID_RE.match(step.id)
            if match is None:
                raise ValueError(
                    f"Execution plan step ids must use explicit ReWOO-style variables E1..En; got '{step.id}'"
                )
            if int(match.group(1)) != index:
                raise ValueError(
                    f"Execution plan step ids must be sequential E1..E{len(plan.steps)}; got '{step.id}' at position {index}"
                )
            if step.id in seen:
                raise ValueError(f"Execution plan contains duplicate step id '{step.id}'")
            missing = [dependency for dependency in step.depends_on if dependency not in seen]
            if missing:
                raise ValueError(
                    f"Step '{step.id}' depends on unknown or future steps: {', '.join(missing)}"
                )
            if (step.arg_mapping or step.tool_input_schema) and not step.tool:
                raise ValueError(
                    f"Step '{step.id}' defines tool_input_schema/arg_mapping without selecting a tool"
                )
            if step.output_key:
                if not re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", step.output_key):
                    raise ValueError(
                        f"Step '{step.id}' defines invalid output_key '{step.output_key}'"
                    )
                if step.output_key in seen_output_keys:
                    raise ValueError(
                        f"Execution plan contains duplicate output_key '{step.output_key}'"
                    )
                seen_output_keys.add(step.output_key)
            seen.add(step.id)
        return plan

    async def _execute_plan(
        self,
        task: str,
        plan: ExecutionPlan,
        *,
        session_id: str | None,
        **kwargs: Any,
    ) -> list[ExecutedPlanStep]:
        executed_by_id: dict[str, ExecutedPlanStep] = {}
        step_outputs: dict[str, str] = {}
        named_outputs: dict[str, str] = {}
        self._last_execution_responses_by_id: dict[str, Response] = {}
        pending = {step.id: step for step in plan.steps}
        order = {step.id: index for index, step in enumerate(plan.steps)}

        while pending:
            ready = [
                step
                for step in plan.steps
                if step.id in pending and all(dependency in step_outputs for dependency in step.depends_on)
            ]
            if not ready:
                unresolved = ", ".join(sorted(pending))
                raise ValueError(f"Execution plan contains unresolved dependencies for: {unresolved}")

            if not self._parallel_execution or len(ready) == 1:
                batch = ready[:1]
            else:
                batch_limit = self._max_parallel_steps or len(ready)
                batch = ready[:batch_limit]

            results = await asyncio.gather(
                *[
                    self._execute_single_step(
                        task=task,
                        step=step,
                        step_outputs=dict(step_outputs),
                        named_outputs=dict(named_outputs),
                        session_id=session_id,
                        **kwargs,
                    )
                    for step in batch
                ]
            )

            for executed_step, response in sorted(results, key=lambda item: order[item[0].id]):
                step_outputs[executed_step.id] = executed_step.output
                if executed_step.output_key:
                    named_outputs[executed_step.output_key] = executed_step.output
                executed_by_id[executed_step.id] = executed_step
                self._last_execution_responses_by_id[executed_step.id] = response
                pending.pop(executed_step.id, None)

        return [executed_by_id[step.id] for step in plan.steps]

    async def _execute_single_step(
        self,
        *,
        task: str,
        step: PlanStep,
        step_outputs: dict[str, str],
        named_outputs: dict[str, str],
        session_id: str | None,
        **kwargs: Any,
    ) -> tuple[ExecutedPlanStep, Response]:
        resolved_instruction = _render_rewoo_text(
            step.instruction,
            task=task,
            step_outputs=step_outputs,
            named_outputs=named_outputs,
        )
        resolved_input = self._render_step_input(
            task=task,
            step=step,
            resolved_instruction=resolved_instruction,
            step_outputs=step_outputs,
            named_outputs=named_outputs,
        )
        resolved_arguments = self._render_step_arguments(
            task=task,
            step=step,
            step_outputs=step_outputs,
            named_outputs=named_outputs,
        )
        started = time.monotonic()
        if step.tool:
            response = await self._execute_tool_step(
                step=step,
                resolved_input=resolved_input or resolved_instruction,
                resolved_arguments=resolved_arguments,
            )
        else:
            step_prompt = self._build_step_prompt(
                task=task,
                step=step,
                resolved_instruction=resolved_instruction,
                resolved_input=resolved_input,
                step_outputs=step_outputs,
                named_outputs=named_outputs,
            )
            response = await _call_component(
                self._executor,
                step_prompt,
                session_id=_child_session_id(session_id, step.id),
                **kwargs,
            )
        duration_ms = (time.monotonic() - started) * 1000
        output = response.content.strip()
        return (
            ExecutedPlanStep(
                id=step.id,
                instruction=step.instruction,
                resolved_instruction=resolved_instruction,
                tool=step.tool,
                input_template=step.input_template,
                resolved_input=resolved_input,
                tool_input_schema=step.tool_input_schema,
                arg_mapping=step.arg_mapping,
                resolved_arguments=resolved_arguments,
                output_key=step.output_key,
                depends_on=list(step.depends_on),
                output=output,
                duration_ms=duration_ms,
                executor_name=_component_name(self._executor),
            ),
            response,
        )

    async def _execute_tool_step(
        self,
        *,
        step: PlanStep,
        resolved_input: str,
        resolved_arguments: dict[str, Any] | None,
    ) -> Response:
        tool = _find_tool(self._executor, step.tool or "")
        if tool is None:
            raise ValueError(
                f"Step '{step.id}' requested tool '{step.tool}', but the executor does not expose it"
            )
        arguments = _build_tool_arguments(
            tool,
            rendered_input=resolved_input,
            rendered_arguments=resolved_arguments,
            tool_input_schema=step.tool_input_schema,
        )
        result = await tool(**arguments)
        return Response(
            content="" if result is None else str(result),
            tool_calls=[],
            usage=Usage(input_tokens=0, output_tokens=0),
            raw={
                "plan_and_execute_tool_step": {
                    "step_id": step.id,
                    "tool": tool.name,
                    "arguments": arguments,
                }
            },
        )

    def _render_step_input(
        self,
        *,
        task: str,
        step: PlanStep,
        resolved_instruction: str,
        step_outputs: dict[str, str],
        named_outputs: dict[str, str],
    ) -> str | None:
        template = step.input_template
        if template is None:
            return resolved_instruction if step.tool else None
        return _render_rewoo_text(
            template,
            task=task,
            step_outputs=step_outputs,
            named_outputs=named_outputs,
        )

    def _render_step_arguments(
        self,
        *,
        task: str,
        step: PlanStep,
        step_outputs: dict[str, str],
        named_outputs: dict[str, str],
    ) -> dict[str, Any] | None:
        if step.arg_mapping is None:
            return None
        rendered = _render_rewoo_value(
            step.arg_mapping,
            task=task,
            step_outputs=step_outputs,
            named_outputs=named_outputs,
        )
        if not isinstance(rendered, dict):
            raise ValueError(f"Step '{step.id}' arg_mapping must resolve to an object")
        return rendered

    def _execution_responses(self, executed_steps: list[ExecutedPlanStep]) -> list[Response]:
        responses = getattr(self, "_last_execution_responses_by_id", {})
        if responses and len(responses) == len(executed_steps):
            return [responses[step.id] for step in executed_steps]
        return [
            Response(content=step.output, tool_calls=[], usage=Usage(0, 0), raw={})
            for step in executed_steps
        ]

    def _build_planner_prompt(self, task: str) -> str:
        return (
            "You are a planning agent. Break the task into a short ordered ReWOO-style plan.\n"
            "Return JSON only.\n"
            f"Use the schema {self._planner_output_schema.__name__} with shape "
            "{\"steps\": [{\"id\": \"E1\", \"instruction\": \"...\", \"depends_on\": [], "
            "\"tool\": null, \"input_template\": null, \"tool_input_schema\": null, "
            "\"arg_mapping\": null, \"output_key\": null}]}.\n"
            f"Use between 1 and {self._max_steps} steps.\n"
            "Dependencies must only refer to earlier step ids.\n\n"
            "Rules:\n"
            "- Use explicit step ids E1..En.\n"
            "- Set tool when a specific executor tool should be called directly.\n"
            "- Use input_template to define the exact execution input. It may be plain text or a JSON object string for tool kwargs.\n"
            "- Use arg_mapping for strict structured tool arguments, especially for multi-argument tools.\n"
            "- Use tool_input_schema to describe/validate the expected tool argument object.\n"
            "- Use output_key when the step output should be referenced later by semantic name like {facts}.\n"
            "- Later steps may reference prior outputs with #E1 or {output_key}.\n\n"
            f"Task:\n{task}"
        )

    def _build_step_prompt(
        self,
        *,
        task: str,
        step: PlanStep,
        resolved_instruction: str,
        resolved_input: str | None,
        step_outputs: dict[str, str],
        named_outputs: dict[str, str],
    ) -> str:
        dependency_lines = [
            f"- {dependency}: {step_outputs[dependency]}"
            for dependency in step.depends_on
            if dependency in step_outputs
        ]
        previous_lines = [
            f"- {step_id}: {output}"
            for step_id, output in step_outputs.items()
        ]
        named_output_lines = [
            f"- {name}: {output}"
            for name, output in named_outputs.items()
        ]
        tool_lines = _tool_lines(self._executor)
        sections = [
            "You are executing one step in a larger ReWOO-style plan.",
            f"Original task:\n{task}",
            f"Current step ({step.id}):\n{resolved_instruction}",
            f"Write the result as the value assigned to variable {step.id}.",
        ]
        if step.output_key:
            sections.append(f"This step should also populate semantic output key '{step.output_key}'.")
        if dependency_lines:
            sections.append("Required context from dependencies:\n" + "\n".join(dependency_lines))
        elif previous_lines:
            sections.append("Completed prior steps:\n" + "\n".join(previous_lines))
        if named_output_lines:
            sections.append("Named outputs available from prior steps:\n" + "\n".join(named_output_lines))
        if resolved_input is not None:
            sections.append(f"Execution input:\n{resolved_input}")
        if tool_lines:
            sections.append(
                "Available tools for this executor:\n"
                + "\n".join(tool_lines)
                + "\nUse them when they materially improve the step result."
            )
        sections.append("Return only the result for this step.")
        return "\n\n".join(sections)

    def _build_solver_prompt(
        self,
        task: str,
        plan: ExecutionPlan,
        executed_steps: list[ExecutedPlanStep],
    ) -> str:
        plan_lines = []
        for step in plan.steps:
            plan_line = f"- {step.id}: {step.instruction}"
            if step.tool:
                plan_line += f" [tool={step.tool}]"
            if step.output_key:
                plan_line += f" [output_key={step.output_key}]"
            plan_lines.append(plan_line)
        result_lines = []
        for step in executed_steps:
            result_line = f"- {step.id} = {step.output}"
            if step.output_key:
                result_line += f" ({step.output_key})"
            result_lines.append(result_line)
        return (
            "You are the solver for a plan-and-execute workflow.\n\n"
            f"Original task:\n{task}\n\n"
            f"Plan:\n{chr(10).join(plan_lines)}\n\n"
            f"Step outputs:\n{chr(10).join(result_lines)}\n\n"
            "Produce the final answer for the user. Do not expose internal chain-of-thought."
        )


class ReWOOAgent(PlanAndExecute):
    """Compatibility alias for users who want a ReWOO-style planner/executor wrapper."""


async def _call_component(
    component: PlanComponent,
    prompt: str,
    *,
    session_id: str | None = None,
    **kwargs: Any,
) -> Response:
    if isinstance(component, Runnable):
        return await component.run(prompt, session_id=session_id, **kwargs)

    try:
        result = component(prompt, session_id=session_id, **kwargs)
    except TypeError:
        result = component(prompt)

    if inspect.isawaitable(result):
        result = await result
    if isinstance(result, Response):
        return result
    return Response(
        content=str(result),
        tool_calls=[],
        usage=Usage(input_tokens=0, output_tokens=0),
        raw={},
    )


def _component_name(component: PlanComponent) -> str:
    return getattr(component, "name", None) or getattr(component, "__name__", None) or type(component).__name__


def _child_session_id(session_id: str | None, suffix: str) -> str | None:
    if session_id is None:
        return None
    return f"{session_id}:{suffix}"


def _strip_code_fences(content: str) -> str:
    text = content.strip()
    if text.startswith("```"):
        lines = text.splitlines()
        if len(lines) >= 3 and lines[-1].strip() == "```":
            return "\n".join(lines[1:-1]).strip()
    return text


def _render_rewoo_text(
    text: str,
    *,
    task: str,
    step_outputs: dict[str, str],
    named_outputs: dict[str, str],
) -> str:
    resolved = text.replace("{task}", task)
    for step_id in sorted(step_outputs, key=len, reverse=True):
        pattern = rf"(?<![A-Za-z0-9_])#{re.escape(step_id)}(?![A-Za-z0-9_])"
        resolved = re.sub(pattern, step_outputs[step_id], resolved)
    for output_key in sorted(named_outputs, key=len, reverse=True):
        pattern = rf"\{{{re.escape(output_key)}}}"
        resolved = re.sub(pattern, named_outputs[output_key], resolved)
    return resolved


def _render_rewoo_value(
    value: Any,
    *,
    task: str,
    step_outputs: dict[str, str],
    named_outputs: dict[str, str],
) -> Any:
    if isinstance(value, str):
        return _render_rewoo_text(
            value,
            task=task,
            step_outputs=step_outputs,
            named_outputs=named_outputs,
        )
    if isinstance(value, list):
        return [
            _render_rewoo_value(
                item,
                task=task,
                step_outputs=step_outputs,
                named_outputs=named_outputs,
            )
            for item in value
        ]
    if isinstance(value, dict):
        return {
            key: _render_rewoo_value(
                item,
                task=task,
                step_outputs=step_outputs,
                named_outputs=named_outputs,
            )
            for key, item in value.items()
        }
    return value


def _find_tool(component: PlanComponent, name: str) -> Any | None:
    if not name:
        return None
    finder = getattr(component, "_find_tool", None)
    if callable(finder):
        try:
            tool = finder(name)
        except Exception:
            tool = None
        if tool is not None:
            return tool
    config = getattr(component, "_config", None)
    tools = getattr(config, "tools", None) if config is not None else None
    if tools is None and hasattr(component, "get_tools") and callable(component.get_tools):
        try:
            tools = component.get_tools()
        except Exception:
            tools = None
    if not tools:
        return None
    for tool in tools:
        if getattr(tool, "name", None) == name:
            return tool
    return None


def _build_tool_arguments(
    tool: Any,
    *,
    rendered_input: str,
    rendered_arguments: dict[str, Any] | None,
    tool_input_schema: dict[str, Any] | None,
) -> dict[str, Any]:
    if rendered_arguments is not None:
        if tool_input_schema is not None:
            _validate_schema_value(
                rendered_arguments,
                tool_input_schema,
                path="tool_input",
            )
        return rendered_arguments

    payload = rendered_input.strip()
    if not payload:
        return {}
    try:
        decoded = json.loads(payload)
    except json.JSONDecodeError:
        decoded = None
    if isinstance(decoded, dict):
        if tool_input_schema is not None:
            _validate_schema_value(decoded, tool_input_schema, path="tool_input")
        return decoded

    schema = getattr(tool, "schema", {}) or {}
    parameters = schema.get("parameters", {})
    properties = parameters.get("properties", {}) or {}
    required = parameters.get("required", []) or []
    candidates = [
        name
        for name in [required[0] if len(required) == 1 else None, "input", "query", "text"]
        if name is not None and name in properties
    ]
    if not candidates and len(properties) == 1:
        candidates = [next(iter(properties))]
    if not candidates:
        raise ValueError(
            f"Tool '{getattr(tool, 'name', '<unknown>')}' requires a JSON object input_template"
        )
    value = decoded if decoded is not None else payload
    arguments = {candidates[0]: value}
    if tool_input_schema is not None:
        _validate_schema_value(arguments, tool_input_schema, path="tool_input")
    return arguments


def _validate_schema_value(value: Any, schema: dict[str, Any], *, path: str) -> None:
    expected_type = schema.get("type")
    if expected_type is None and "properties" in schema:
        expected_type = "object"

    if expected_type == "object":
        if not isinstance(value, dict):
            raise ValueError(f"{path} must be an object")
        properties = schema.get("properties", {}) or {}
        required = schema.get("required", []) or []
        missing = [name for name in required if name not in value]
        if missing:
            raise ValueError(f"{path} is missing required fields: {', '.join(missing)}")
        for name, sub_schema in properties.items():
            if name in value and isinstance(sub_schema, dict):
                _validate_schema_value(value[name], sub_schema, path=f"{path}.{name}")
        return

    if expected_type == "array":
        if not isinstance(value, list):
            raise ValueError(f"{path} must be an array")
        item_schema = schema.get("items")
        if isinstance(item_schema, dict):
            for index, item in enumerate(value):
                _validate_schema_value(item, item_schema, path=f"{path}[{index}]")
        return

    validators = {
        "string": lambda item: isinstance(item, str),
        "integer": lambda item: isinstance(item, int) and not isinstance(item, bool),
        "number": lambda item: isinstance(item, (int, float)) and not isinstance(item, bool),
        "boolean": lambda item: isinstance(item, bool),
        "null": lambda item: item is None,
    }
    if expected_type in validators and not validators[expected_type](value):
        raise ValueError(f"{path} must be of type {expected_type}")


def _tool_lines(component: PlanComponent) -> list[str]:
    tools = None
    config = getattr(component, "_config", None)
    if config is not None:
        tools = getattr(config, "tools", None)
    if tools is None and hasattr(component, "get_tools") and callable(component.get_tools):
        try:
            tools = component.get_tools()
        except Exception:
            tools = None
    if not tools:
        return []
    lines: list[str] = []
    for tool in tools:
        name = getattr(tool, "name", None)
        description = getattr(tool, "description", None) or ""
        if name:
            lines.append(f"- {name}: {description}".rstrip())
    return lines


def _aggregate_usage(responses: list[Response]) -> Usage:
    return Usage(
        input_tokens=sum(response.usage.input_tokens for response in responses),
        output_tokens=sum(response.usage.output_tokens for response in responses),
    )


__all__ = [
    "ExecutedPlanStep",
    "ExecutionPlan",
    "ExecutionPlanSchema",
    "PlanAndExecute",
    "PlanStep",
    "PlanStepSchema",
    "ReWOOExecutionPlanSchema",
    "ReWOOPlanStepSchema",
    "ReWOOAgent",
]
