"""Minimal DiagnosticAgent for tests.

Provides simple tools and a custom plan that writes to memory.
"""
from __future__ import annotations

from typing import Any, Dict, List

from evenage.core import Agent, AgentConfig, actor


def echo(text: str) -> str:
    """Echo back the provided text."""
    return text


def sum(a: int, b: int) -> int:  # noqa: A003 - test helper tool name
    """Return the sum of two integers."""
    return int(a) + int(b)


@actor(
    name="diagnostic",
    role="Diagnostic",
    goal="Test agent for unit tests",
    allow_delegation=False,
)
class DiagnosticAgent(Agent):
    """A tiny agent used in unit tests."""

    def __init__(self, config: AgentConfig, env_config=None):
        super().__init__(config, env_config)
        # Register two simple tools expected by tests
        self.tools.register(echo)
        self.tools.register(sum)

    def plan(self, task: Dict[str, Any]) -> List[Dict[str, Any]]:  # type: ignore[override]
        # Write the incoming message into memory, then finish
        return [
            {"action": "custom", "op": "write_memory", "key": "last_message"},
        ]

    async def _execute_custom_step(self, step: Dict[str, Any], context: Dict[str, Any]) -> Any:  # type: ignore[override]
        if step.get("op") == "write_memory":
            key = step.get("key", "last_message")
            msg = ""
            # The base Agent passes the original task in context["task"]
            task = context.get("task", {})
            if isinstance(task, dict):
                msg = str(task.get("message", ""))
            if self.memory is not None:
                self.memory.put(key, msg)
            return {"status": "ok"}
        return await super()._execute_custom_step(step, context)
