from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Any, Callable

logger = logging.getLogger(__name__)

StepHandler = Callable[[str, dict[str, Any], int], dict[str, Any]]


@dataclass
class AgentLoopRun:
    loop_id: str
    status: str
    message: str
    steps: list[dict[str, Any]] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "loop_id": self.loop_id,
            "status": self.status,
            "message": self.message,
            "steps": self.steps,
        }


def _default_step_handler(agent: str, context: dict[str, Any], index: int) -> dict[str, Any]:
    logger.info("agent loop step %s agent=%s context_keys=%s", index, agent, sorted(context.keys()))
    return {"message": f"step {index + 1}: {agent} ok"}


class AgentLoopRunner:
    """Sequential agent loop runner — v1 log + status only, no LLM."""

    def __init__(
        self,
        handlers: dict[str, StepHandler] | None = None,
        default_handler: StepHandler | None = None,
    ) -> None:
        self._handlers = handlers or {}
        self._default = default_handler or _default_step_handler

    def run(self, loop_id: str, cfg: dict[str, Any], *, context: dict[str, Any] | None = None) -> AgentLoopRun:
        ctx = dict(context or {})
        raw_steps = cfg.get("steps") or []
        executed: list[dict[str, Any]] = []
        for index, step in enumerate(raw_steps):
            agent = str(step.get("agent") if isinstance(step, dict) else step)
            handler = self._handlers.get(agent, self._default)
            detail = handler(agent, ctx, index)
            executed.append({"index": index, "agent": agent, "status": "done", **detail})
        return AgentLoopRun(
            loop_id=loop_id,
            status="done",
            message=f"completed {len(executed)} step(s)",
            steps=executed,
        )
