"""TaskRunner — deterministic execution of a static `TaskList`. Zero LLM.

Executes tasks in dependency order, parallelizing each ready "wave" with
`asyncio.gather`. For each task it resolves `${t<id>}` placeholders from upstream
results, does an internal `validate_args`, invokes each tool via the `ToolInvoker`
seam, and records a `TaskResult`. On failure it **degrades and continues**: the
task is marked failed, its dependents are skipped, independent branches keep
running. There is no replanning and no mid-run LLM (INV-6).

`success_criteria` is *not* machine-evaluated here (it is free text); task status
is derived from tool execution outcomes and carried to the Assembler to report.

See AGENT_ARCHITECTURE_CONTEXT_new.md §7.4.
"""

from __future__ import annotations

import asyncio
from typing import Any

from src.middlewares.logging import get_logger

from ..planner.contracts import ToolOutput, ToolRegistry
from ..planner.schemas import PLACEHOLDER_RE, Task
from ..planner.schemas import TaskList as PlanTaskList
from .invoker import ToolInvoker
from .schemas import RunState, TaskResult, TaskStatus

logger = get_logger("task_runner")


class TaskRunner:
    """Runs a `TaskList` against a `ToolInvoker`, producing a `RunState`."""

    def __init__(self, invoker: ToolInvoker, registry: ToolRegistry) -> None:
        self._invoker = invoker
        self._registry = registry

    async def run(self, task_list: PlanTaskList, business_context_id: str) -> RunState:
        tasks_by_id: dict[str, Task] = {t.id: t for t in task_list.tasks}
        results: dict[str, TaskResult] = {}
        remaining: set[str] = set(tasks_by_id)

        while remaining:
            ready = [
                tid
                for tid in remaining
                if all(dep in results for dep in tasks_by_id[tid].depends_on)
            ]
            if not ready:
                # A dependency points outside the plan (or a cycle slipped past the
                # planner validator): nothing more can run. Fail the rest honestly.
                for tid in list(remaining):
                    results[tid] = TaskResult(
                        task_id=tid,
                        stage=tasks_by_id[tid].stage,
                        status="failure",
                        objective=tasks_by_id[tid].objective,
                        error="unresolved dependency; task could not run",
                    )
                    remaining.discard(tid)
                break

            # Skip any ready task whose dependency failed (degrade-and-continue).
            to_run: list[Task] = []
            for tid in ready:
                task = tasks_by_id[tid]
                failed = [d for d in task.depends_on if results[d].status == "failure"]
                if failed:
                    results[tid] = TaskResult(
                        task_id=tid,
                        stage=task.stage,
                        status="failure",
                        objective=task.objective,
                        error=f"skipped: upstream {failed} did not succeed",
                    )
                    remaining.discard(tid)
                else:
                    to_run.append(task)

            if not to_run:
                continue  # remaining dependents will be re-evaluated (and skipped)

            wave = await asyncio.gather(
                *(self._run_task(task, results) for task in to_run)
            )
            for tr in wave:
                results[tr.task_id] = tr
                remaining.discard(tr.task_id)

        return RunState(
            plan_id=task_list.plan_id,
            business_context_id=business_context_id,
            results=results,
            open_questions=list(task_list.open_questions),
        )

    async def _run_task(self, task: Task, results: dict[str, TaskResult]) -> TaskResult:
        outputs: list[ToolOutput] = []
        for call in task.tool_calls:
            resolved = self._resolve_args(call.args, results)
            arg_error = self._validate_args(call.tool, resolved)
            if arg_error is not None:
                outputs.append(ToolOutput(tool=call.tool, kind="error", error=arg_error))
                continue
            outputs.append(await self._safe_invoke(call.tool, resolved))

        status = _label(outputs)
        error: str | None = None
        if status == "failure":
            errs = [o.error for o in outputs if o.kind == "error" and o.error]
            error = errs[0] if errs else "all tool calls failed"
        return TaskResult(
            task_id=task.id,
            stage=task.stage,
            status=status,
            objective=task.objective,
            outputs=outputs,
            error=error,
        )

    def _resolve_args(
        self, args: dict[str, Any], results: dict[str, TaskResult]
    ) -> dict[str, Any]:
        return {k: self._resolve_value(v, results) for k, v in args.items()}

    @staticmethod
    def _resolve_value(value: Any, results: dict[str, TaskResult]) -> Any:
        # A data arg is exactly a "${t<id>}" placeholder (Pattern A); resolve it to
        # the referenced task's representative output (its last ToolOutput).
        # Materializing that envelope into a DataFrame is the invoker's job.
        if isinstance(value, str):
            match = PLACEHOLDER_RE.fullmatch(value.strip())
            if match:
                upstream = results.get(match.group(1))
                if upstream is None or not upstream.outputs:
                    return None
                return upstream.outputs[-1]
        return value

    def _validate_args(self, tool: str, resolved: dict[str, Any]) -> str | None:
        spec = self._registry.get(tool)
        if spec is None:
            return f"tool {tool!r} not in registry"
        required = spec.input_schema.get("required", [])
        missing = [r for r in required if resolved.get(r) is None]
        if missing:
            return f"missing required arg(s): {sorted(missing)}"
        return None

    async def _safe_invoke(self, tool: str, args: dict[str, Any]) -> ToolOutput:
        try:
            return await self._invoker.invoke(tool, args)
        except Exception as exc:  # noqa: BLE001 — backstop; the invoker is never-throw (§8.4)
            logger.warning("tool invoker raised", tool=tool, error=str(exc))
            return ToolOutput(tool=tool, kind="error", error=f"invoker raised: {exc}")


def _label(outputs: list[ToolOutput]) -> TaskStatus:
    if not outputs:
        return "failure"
    errors = sum(1 for o in outputs if o.kind == "error")
    if errors == 0:
        return "success"
    if errors == len(outputs):
        return "failure"
    return "partial"