core/agent.py

"""
core.agent
==========

Generic agent representation for WorldSmithAI.

Agents are domain-agnostic runtime entities. They do not contain species,
economy, civilization, or ecosystem logic. They hold state and delegate action
selection to policies and behavior execution to behaviors.
"""

from __future__ import annotations

import logging
from copy import deepcopy
from dataclasses import dataclass, field
from typing import Any, Mapping, MutableMapping, Protocol, Sequence, TypeAlias

import numpy as np
from numpy.typing import NDArray

logger = logging.getLogger(__name__)

PositionInput: TypeAlias = Sequence[float] | NDArray[np.float64] | None
Observation: TypeAlias = dict[str, Any]


class WorldProtocol(Protocol):
    """Structural protocol for world-like objects observed by Agent."""

    step_count: int

    def query(
        self,
        *,
        agent: Agent,
        position: NDArray[np.float64],
    ) -> Mapping[str, Any]:
        """Return a world observation for an agent."""
        ...


class BehaviorProtocol(Protocol):
    """Structural protocol for executable agent behaviors."""

    def check_preconditions(self, agent: Agent, world: WorldProtocol) -> bool:
        """Return whether this behavior can currently execute."""
        ...

    def execute(
        self,
        agent: Agent,
        world: WorldProtocol,
    ) -> BehaviorExecution | Mapping[str, Any] | None:
        """Execute the behavior for the given agent."""
        ...


class PolicyProtocol(Protocol):
    """Structural protocol for agent policies."""

    def choose_action(self, *args: Any, **kwargs: Any) -> str | BehaviorProtocol | None:
        """Choose the next behavior for an agent."""
        ...

    def update(self, *args: Any, **kwargs: Any) -> None:
        """Update the policy from a reward signal."""
        ...


@dataclass(frozen=True, slots=True)
class BehaviorExecution:
    """Normalized result returned by behavior execution."""

    success: bool = True
    reward: float = 0.0
    state_updates: Mapping[str, Any] = field(default_factory=dict)
    memory_updates: Mapping[str, Any] = field(default_factory=dict)
    message: str = ""
    metadata: Mapping[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        """Validate and normalize behavior execution data."""

        if not np.isfinite(self.reward):
            raise ValueError("BehaviorExecution.reward must be finite.")

        if not isinstance(self.state_updates, Mapping):
            raise TypeError("BehaviorExecution.state_updates must be a mapping.")

        if not isinstance(self.memory_updates, Mapping):
            raise TypeError("BehaviorExecution.memory_updates must be a mapping.")

        if not isinstance(self.metadata, Mapping):
            raise TypeError("BehaviorExecution.metadata must be a mapping.")

        object.__setattr__(self, "reward", float(self.reward))
        object.__setattr__(self, "state_updates", dict(self.state_updates))
        object.__setattr__(self, "memory_updates", dict(self.memory_updates))
        object.__setattr__(self, "metadata", dict(self.metadata))


@dataclass(frozen=True, slots=True)
class AgentStepResult:
    """Result of one agent simulation step."""

    agent_id: str
    agent_type: str
    step_count: int | None
    behavior_id: str | None
    success: bool
    reward: float = 0.0
    message: str = ""
    metadata: Mapping[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        """Convert the step result into a JSON-friendly dictionary."""

        return {
            "agent_id": self.agent_id,
            "agent_type": self.agent_type,
            "step_count": self.step_count,
            "behavior_id": self.behavior_id,
            "success": self.success,
            "reward": self.reward,
            "message": self.message,
            "metadata": dict(self.metadata),
        }


@dataclass(slots=True)
class Agent:
    """Generic domain-agnostic agent.

    Compatibility note:
        WorldSmithAI uses mapping-style memory throughout the DSL and behavior
        system. Examples include ``memory.goals``, ``memory.known_options``,
        and ``memory.relationships``. Therefore memory is a mutable mapping,
        not a list. Older list-style memory is preserved under
        ``memory["entries"]``.
    """

    id: str
    type: str
    position: PositionInput = None
    state: dict[str, Any] = field(default_factory=dict)
    memory: MutableMapping[str, Any] = field(default_factory=dict)
    goals: Any = field(default_factory=list)
    behaviors: Mapping[str, BehaviorProtocol] | Sequence[BehaviorProtocol] = field(default_factory=dict)
    policy: PolicyProtocol | Any | None = None
    alive: bool = True
    memory_limit: int | None = None
    metadata: dict[str, Any] = field(default_factory=dict)
    dsl_spec: Any | None = None

    def __post_init__(self) -> None:
        """Validate and normalize agent fields after initialization."""

        if not self.id or not isinstance(self.id, str):
            raise ValueError("Agent.id must be a non-empty string.")

        if not self.type or not isinstance(self.type, str):
            raise ValueError("Agent.type must be a non-empty string.")

        self.id = self.id.strip()
        self.type = self.type.strip()
        self.position = self._normalize_position(self.position)

        if not isinstance(self.state, dict):
            if isinstance(self.state, Mapping):
                self.state = dict(self.state)
            else:
                raise TypeError("Agent.state must be a dictionary.")

        self.memory = self._normalize_memory(self.memory)
        self.goals = self._normalize_goals(self.goals)
        self.behaviors = self._normalize_behaviors(self.behaviors)

        if not isinstance(self.metadata, dict):
            if isinstance(self.metadata, Mapping):
                self.metadata = dict(self.metadata)
            else:
                raise TypeError("Agent.metadata must be a dictionary.")

        if self.memory_limit is not None and self.memory_limit <= 0:
            raise ValueError("Agent.memory_limit must be positive when provided.")

        self.alive = bool(self.alive)

    def step(self, world: WorldProtocol) -> AgentStepResult:
        """Execute one agent step."""

        step_count = self._get_world_step_count(world)

        if not self.alive:
            return AgentStepResult(
                agent_id=self.id,
                agent_type=self.type,
                step_count=step_count,
                behavior_id=None,
                success=False,
                reward=0.0,
                message="Agent is not alive.",
                metadata={"skipped": True},
            )

        observation = self.observe(world)
        behavior = self._select_behavior(world=world, observation=observation)

        if behavior is None:
            result = AgentStepResult(
                agent_id=self.id,
                agent_type=self.type,
                step_count=step_count,
                behavior_id=None,
                success=False,
                reward=0.0,
                message="No executable behavior selected.",
                metadata={"skipped": True},
            )
            self._remember(
                {
                    "kind": "step",
                    "step_count": step_count,
                    "result": result.to_dict(),
                }
            )
            return result

        behavior_id = self._behavior_identifier(behavior)

        try:
            can_execute = behavior.check_preconditions(self, world)
        except Exception as exc:
            logger.exception(
                "Behavior precondition check failed for agent_id=%s behavior_id=%s",
                self.id,
                behavior_id,
            )
            return AgentStepResult(
                agent_id=self.id,
                agent_type=self.type,
                step_count=step_count,
                behavior_id=behavior_id,
                success=False,
                reward=0.0,
                message="Behavior precondition check raised an exception.",
                metadata={"error": repr(exc)},
            )

        if not can_execute:
            result = AgentStepResult(
                agent_id=self.id,
                agent_type=self.type,
                step_count=step_count,
                behavior_id=behavior_id,
                success=False,
                reward=0.0,
                message="Behavior preconditions were not satisfied.",
                metadata={"skipped": True},
            )
            self._remember(
                {
                    "kind": "step",
                    "step_count": step_count,
                    "behavior_id": behavior_id,
                    "result": result.to_dict(),
                }
            )
            return result

        try:
            raw_execution = behavior.execute(self, world)
            execution = self._normalize_behavior_execution(raw_execution)
        except Exception as exc:
            logger.exception(
                "Behavior execution failed for agent_id=%s behavior_id=%s",
                self.id,
                behavior_id,
            )
            return AgentStepResult(
                agent_id=self.id,
                agent_type=self.type,
                step_count=step_count,
                behavior_id=behavior_id,
                success=False,
                reward=0.0,
                message="Behavior execution raised an exception.",
                metadata={"error": repr(exc)},
            )

        if execution.state_updates:
            self.update_state(execution.state_updates)

        if execution.memory_updates:
            self._remember(
                {
                    "kind": "behavior_memory",
                    "step_count": step_count,
                    "behavior_id": behavior_id,
                    "payload": dict(execution.memory_updates),
                }
            )

        self.receive_reward(
            execution.reward,
            action_id=behavior_id,
            observation=observation,
            metadata=execution.metadata,
            world=world,
            behavior=behavior,
        )

        result = AgentStepResult(
            agent_id=self.id,
            agent_type=self.type,
            step_count=step_count,
            behavior_id=behavior_id,
            success=execution.success,
            reward=execution.reward,
            message=execution.message,
            metadata=execution.metadata,
        )

        self._remember(
            {
                "kind": "step",
                "step_count": step_count,
                "behavior_id": behavior_id,
                "result": result.to_dict(),
            }
        )

        return result

    def observe(self, world: WorldProtocol) -> Observation:
        """Generate an observation for this agent."""

        observation: Observation = {
            "self": self.snapshot(),
            "world": {},
        }

        query = getattr(world, "query", None)
        if not callable(query):
            return observation

        try:
            position = self._position_array()
            world_view = query(agent=self, position=position.copy())
        except Exception as exc:
            logger.exception("World query failed for agent_id=%s", self.id)
            observation["world_query_error"] = repr(exc)
            return observation

        if isinstance(world_view, Mapping):
            observation["world"] = dict(world_view)
        else:
            observation["world"] = {"value": world_view}

        return observation

    def update_state(self, updates: Mapping[str, Any]) -> None:
        """Update this agent's mutable state."""

        if not isinstance(updates, Mapping):
            raise TypeError("updates must be a mapping.")

        normalized_updates: dict[str, Any] = {}

        for key, value in updates.items():
            if not isinstance(key, str):
                raise TypeError("Agent state keys must be strings.")
            normalized_updates[key] = deepcopy(value)

        self.state.update(normalized_updates)

    def receive_reward(
        self,
        reward: float,
        *,
        action_id: str | None = None,
        observation: Mapping[str, Any] | None = None,
        metadata: Mapping[str, Any] | None = None,
        world: WorldProtocol | None = None,
        behavior: BehaviorProtocol | str | None = None,
    ) -> None:
        """Receive and record a reward signal."""

        if not np.isfinite(reward):
            raise ValueError("reward must be finite.")

        normalized_reward = float(reward)
        normalized_observation = dict(observation or {})
        normalized_metadata = dict(metadata or {})

        self._remember(
            {
                "kind": "reward",
                "action_id": action_id,
                "reward": normalized_reward,
                "metadata": normalized_metadata,
            }
        )

        if self.policy is None:
            return

        update = getattr(self.policy, "update", None)
        if not callable(update):
            return

        behavior_for_update = behavior if behavior is not None else action_id

        try:
            update(
                agent=self,
                world=world,
                behavior=behavior_for_update,
                reward=normalized_reward,
                context=normalized_observation,
                metadata=normalized_metadata,
            )
            return
        except TypeError:
            pass
        except Exception:
            logger.exception(
                "Policy update failed for agent_id=%s action_id=%s",
                self.id,
                action_id,
            )
            return

        try:
            update(
                agent=self,
                reward=normalized_reward,
                action_id=action_id,
                observation=normalized_observation,
                metadata=normalized_metadata,
            )
        except Exception:
            logger.exception(
                "Legacy policy update failed for agent_id=%s action_id=%s",
                self.id,
                action_id,
            )

    def set_position(self, position: PositionInput) -> None:
        """Set the agent's position."""

        self.position = self._normalize_position(position)

    def mark_dead(self, *, reason: str | None = None) -> None:
        """Mark this agent as inactive/dead."""

        self.alive = False
        self._remember(
            {
                "kind": "lifecycle",
                "event": "marked_dead",
                "reason": reason,
            }
        )

    def revive(self, *, reason: str | None = None) -> None:
        """Mark this agent as active/alive again."""

        self.alive = True
        self._remember(
            {
                "kind": "lifecycle",
                "event": "revived",
                "reason": reason,
            }
        )

    def snapshot(self) -> dict[str, Any]:
        """Return a JSON-friendly snapshot of the agent."""

        return {
            "id": self.id,
            "type": self.type,
            "position": None if self.position is None else self._position_array().tolist(),
            "state": deepcopy(self.state),
            "memory_keys": sorted(str(key) for key in self.memory.keys()),
            "goals": deepcopy(self.goals),
            "alive": self.alive,
            "metadata": deepcopy(self.metadata),
        }

    def to_dict(self) -> dict[str, Any]:
        """Return a JSON-friendly dictionary representation."""

        return self.snapshot()

    def _select_behavior(
        self,
        *,
        world: WorldProtocol,
        observation: Mapping[str, Any],
    ) -> BehaviorProtocol | None:
        """Select a behavior through the policy or deterministic fallback."""

        if not self.behaviors:
            return None

        if self.policy is None:
            return self._first_available_behavior(world)

        choose_action = getattr(self.policy, "choose_action", None)
        if not callable(choose_action):
            return self._first_available_behavior(world)

        try:
            selected = choose_action(self, world, self.behaviors)
            return self._resolve_behavior(selected)
        except TypeError:
            pass

        try:
            selected = choose_action(agent=self, world=world, behaviors=self.behaviors)
            return self._resolve_behavior(selected)
        except TypeError:
            pass

        try:
            selected = choose_action(agent=self, world=world, observation=observation)
            return self._resolve_behavior(selected)
        except Exception:
            logger.exception("Policy action selection failed for agent_id=%s", self.id)
            return self._first_available_behavior(world)

    def _first_available_behavior(self, world: WorldProtocol) -> BehaviorProtocol | None:
        """Return the first behavior whose preconditions pass."""

        for behavior in self.behaviors.values():
            try:
                if behavior.check_preconditions(self, world):
                    return behavior
            except Exception:
                logger.exception(
                    "Fallback behavior precondition failed for agent_id=%s",
                    self.id,
                )

        return None

    def _resolve_behavior(
        self,
        selected: str | BehaviorProtocol | None,
    ) -> BehaviorProtocol | None:
        """Resolve a policy-selected behavior id, behavior name, or object."""

        if selected is None:
            return None

        if isinstance(selected, str):
            behavior = self.behaviors.get(selected)
            if behavior is not None:
                return behavior

            for candidate in self.behaviors.values():
                if getattr(candidate, "name", None) == selected:
                    return candidate
                if getattr(candidate, "id", None) == selected:
                    return candidate

            logger.warning(
                "Policy selected unknown behavior_id=%s for agent_id=%s",
                selected,
                self.id,
            )
            return None

        return selected

    def _behavior_identifier(self, behavior: BehaviorProtocol) -> str:
        """Determine a stable behavior identifier."""

        for behavior_id, candidate in self.behaviors.items():
            if candidate is behavior:
                return behavior_id

        explicit_id = getattr(behavior, "id", None)
        if isinstance(explicit_id, str) and explicit_id:
            return explicit_id

        explicit_name = getattr(behavior, "name", None)
        if isinstance(explicit_name, str) and explicit_name:
            return explicit_name

        return behavior.__class__.__name__

    def _normalize_behavior_execution(
        self,
        raw_execution: BehaviorExecution | Mapping[str, Any] | None,
    ) -> BehaviorExecution:
        """Normalize behavior execution output."""

        if raw_execution is None:
            return BehaviorExecution()

        if isinstance(raw_execution, BehaviorExecution):
            return raw_execution

        if isinstance(raw_execution, Mapping):
            metadata = dict(raw_execution.get("metadata") or {})

            if "details" in raw_execution and "details" not in metadata:
                metadata["details"] = deepcopy(raw_execution["details"])

            return BehaviorExecution(
                success=bool(raw_execution.get("success", True)),
                reward=float(raw_execution.get("reward", 0.0)),
                state_updates=dict(raw_execution.get("state_updates") or {}),
                memory_updates=dict(raw_execution.get("memory_updates") or {}),
                message=str(raw_execution.get("message", "")),
                metadata=metadata,
            )

        raise TypeError(
            "Behavior execution must return BehaviorExecution, mapping, or None."
        )

    def _remember(self, entry: Mapping[str, Any]) -> None:
        """Append an entry to mapping-backed agent memory."""

        if not isinstance(entry, Mapping):
            raise TypeError("Memory entry must be a mapping.")

        entries = self.memory.get("entries")
        if not isinstance(entries, list):
            entries = []
            self.memory["entries"] = entries

        entries.append(dict(entry))

        if self.memory_limit is None:
            return

        overflow = len(entries) - self.memory_limit
        if overflow > 0:
            del entries[:overflow]

    def _position_array(self) -> NDArray[np.float64]:
        """Return position as a NumPy array, repairing tuple/list assignment."""

        if self.position is None:
            self.position = np.zeros(2, dtype=np.float64)
            return self.position

        if isinstance(self.position, np.ndarray):
            return self.position

        self.position = self._normalize_position(self.position)
        return self.position

    def _query_position(self) -> NDArray[np.float64]:
        """Return a numeric position for world.query."""

        return self._position_array().copy()

    @staticmethod
    def _normalize_position(position: PositionInput) -> NDArray[np.float64] | None:
        """Normalize a position input into a one-dimensional numpy float array."""

        if position is None:
            return None

        normalized = np.asarray(position, dtype=np.float64).reshape(-1)

        if normalized.ndim != 1:
            raise ValueError("Agent.position must be a one-dimensional vector.")

        if normalized.size == 0:
            raise ValueError("Agent.position cannot be empty.")

        if not np.all(np.isfinite(normalized)):
            raise ValueError("Agent.position must contain only finite values.")

        return normalized.astype(np.float64, copy=True)

    @staticmethod
    def _normalize_memory(memory: Any) -> MutableMapping[str, Any]:
        """Normalize memory into a mutable mapping."""

        if memory is None:
            return {}

        if isinstance(memory, MutableMapping):
            return dict(memory)

        if isinstance(memory, Mapping):
            return dict(memory)

        if isinstance(memory, list):
            return {"entries": list(memory)}

        raise TypeError("Agent.memory must be a mapping or a list of entries.")

    @staticmethod
    def _normalize_goals(goals: Any) -> Any:
        """Normalize goals while preserving DSL flexibility."""

        if goals is None:
            return []

        if isinstance(goals, list):
            return deepcopy(goals)

        if isinstance(goals, tuple):
            return list(deepcopy(goals))

        if isinstance(goals, Mapping):
            return dict(deepcopy(goals))

        return deepcopy(goals)

    @staticmethod
    def _normalize_behaviors(
        behaviors: Mapping[str, BehaviorProtocol] | Sequence[BehaviorProtocol] | None,
    ) -> dict[str, BehaviorProtocol]:
        """Normalize behavior storage into a dictionary."""

        if behaviors is None:
            return {}

        if isinstance(behaviors, Mapping):
            return {
                str(key): behavior
                for key, behavior in behaviors.items()
                if behavior is not None
            }

        if isinstance(behaviors, Sequence) and not isinstance(behaviors, (str, bytes)):
            normalized: dict[str, BehaviorProtocol] = {}

            for index, behavior in enumerate(behaviors):
                if behavior is None:
                    continue

                key = Agent._behavior_key_for_index(behavior, index)

                if key in normalized:
                    key = f"{key}#{index}"

                normalized[key] = behavior

            return normalized

        raise TypeError("Agent.behaviors must be a mapping or a sequence of behaviors.")

    @staticmethod
    def _behavior_key_for_index(behavior: BehaviorProtocol, index: int) -> str:
        """Return a stable dictionary key for a behavior object."""

        explicit_id = getattr(behavior, "id", None)
        if isinstance(explicit_id, str) and explicit_id:
            return explicit_id

        explicit_name = getattr(behavior, "name", None)
        if isinstance(explicit_name, str) and explicit_name:
            return explicit_name

        return f"{behavior.__class__.__name__}_{index}"

    @staticmethod
    def _get_world_step_count(world: WorldProtocol) -> int | None:
        """Safely retrieve the world's current step count."""

        step_count = getattr(world, "step_count", None)

        if isinstance(step_count, int):
            return step_count

        return None


__all__ = [
    "Agent",
    "AgentStepResult",
    "BehaviorExecution",
    "BehaviorProtocol",
    "Observation",
    "PolicyProtocol",
    "PositionInput",
    "WorldProtocol",
]