policies/base_policy.py

"""
Abstract policy interface for WorldSmithAI.

Policies are responsible for selecting which behavior an agent should execute.
This module intentionally contains no domain-specific decision logic. It does
not know about sheep, wolves, scientists, merchants, dragons, vehicles, cities,
markets, farms, or civilizations.

Concrete policies should inherit from ``BasePolicy`` and implement
``choose_action``. The policy may be rule-based, random-like but deterministic,
priority-based, contextual-bandit-based, planner-driven, or learned.

Example:
    class FirstValidPolicy(BasePolicy):
        name = "first_valid"

        def choose_action(self, agent, world, behaviors=None):
            valid_behaviors = self.available_behaviors(agent, world, behaviors)
            return valid_behaviors[0] if valid_behaviors else None

    selected = policy.choose_action(agent, world, agent.behaviors)
    if selected is not None:
        selected.execute(agent, world)

Future extensibility:
    - Add planning-aware policies.
    - Add contextual-bandit policies.
    - Add hierarchical policies.
    - Add ensemble policies.
    - Add policy serialization for Spaces demos.
    - Add policy diagnostics for dashboards and narrators.
"""

from __future__ import annotations

import copy
import logging
from abc import ABC, abstractmethod
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass, field
from enum import Enum
from numbers import Real
from typing import TYPE_CHECKING, Any, ClassVar, cast

from core.behavior import Behavior

if TYPE_CHECKING:
    from core.agent import Agent
    from core.world import World


logger = logging.getLogger(__name__)

BehaviorCollection = Sequence[Behavior] | Mapping[str, Behavior]


class PolicyDecisionStatus(str, Enum):
    """Possible outcomes of a policy decision."""

    SELECTED = "selected"
    DISABLED = "disabled"
    NO_BEHAVIORS = "no_behaviors"
    NO_VALID_BEHAVIORS = "no_valid_behaviors"
    NO_ACTION_SELECTED = "no_action_selected"
    ERROR = "error"


@dataclass(frozen=True)
class BehaviorCandidate:
    """Normalized behavior candidate considered by a policy.

    Attributes:
        behavior: Runtime behavior object.
        name: Stable behavior name.
        index: Deterministic index in the candidate list.
        available: Whether the behavior passed validation.
        reason: Optional explanation when unavailable.
        metadata: Optional behavior metadata copied from the behavior object.
    """

    behavior: Behavior = field(repr=False, compare=False)
    name: str
    index: int
    available: bool
    reason: str | None = None
    metadata: Mapping[str, Any] = field(default_factory=dict)

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

        return {
            "name": self.name,
            "index": self.index,
            "available": self.available,
            "reason": self.reason,
            "metadata": copy.deepcopy(dict(self.metadata)),
        }


@dataclass(frozen=True)
class PolicyDecision:
    """Structured policy decision result.

    This object is useful for debugging, metrics, logging, visualization,
    narration, and reward-learning policies. The selected behavior object is
    retained for runtime use, while ``to_dict`` emits only serializable fields.
    """

    policy_name: str
    agent_id: str
    status: PolicyDecisionStatus
    selected_behavior: Behavior | None = field(default=None, repr=False, compare=False)
    selected_behavior_name: str | None = None
    score: float | None = None
    reason: str | None = None
    candidate_count: int = 0
    valid_candidate_count: int = 0
    candidates: tuple[BehaviorCandidate, ...] = field(default_factory=tuple)
    metadata: Mapping[str, Any] = field(default_factory=dict)

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

        return {
            "policy_name": self.policy_name,
            "agent_id": self.agent_id,
            "status": self.status.value,
            "selected_behavior_name": self.selected_behavior_name,
            "score": self.score,
            "reason": self.reason,
            "candidate_count": self.candidate_count,
            "valid_candidate_count": self.valid_candidate_count,
            "candidates": [candidate.to_dict() for candidate in self.candidates],
            "metadata": copy.deepcopy(dict(self.metadata)),
        }


def _agent_id(agent: Agent) -> str:
    """Return a stable string identifier for an agent."""

    return str(getattr(agent, "id", "unknown_agent"))


def _world_step(world: World) -> int | None:
    """Return the current world step if available."""

    value = getattr(world, "step_count", None)
    if isinstance(value, Real) and not isinstance(value, bool):
        return int(value)

    return None


@dataclass
class BasePolicy(ABC):
    """Abstract base class for all WorldSmithAI policies.

    A policy decides which behavior an agent should execute next. The base class
    provides utility methods for behavior normalization, precondition checking,
    structured decision records, state serialization, and reward-update hooks.

    Subclasses must implement ``choose_action``.

    Design constraints:
        - No domain-specific assumptions.
        - No hardcoded behavior names.
        - No global mutable state.
        - Deterministic by default.
        - Safe failure handling around behavior preconditions.
    """

    enabled: bool = True
    metadata: Mapping[str, Any] = field(default_factory=dict)

    name: ClassVar[str] = "base_policy"

    @abstractmethod
    def choose_action(
        self,
        agent: Agent,
        world: World,
        behaviors: BehaviorCollection | None = None,
    ) -> Behavior | None:
        """Choose one behavior for an agent to execute.

        Args:
            agent: Agent requesting a behavior.
            world: Current world container.
            behaviors: Optional behavior collection. If omitted, the policy
                reads ``agent.behaviors``.

        Returns:
            The selected behavior, or ``None`` if no behavior should execute.

        Notes:
            Implementations should usually call ``self.available_behaviors`` to
            avoid choosing behaviors whose preconditions are not satisfied.
        """

    @property
    def policy_name(self) -> str:
        """Return the stable policy name."""

        return str(getattr(self, "name", self.__class__.__name__))

    def decide(
        self,
        agent: Agent,
        world: World,
        behaviors: BehaviorCollection | None = None,
    ) -> PolicyDecision:
        """Return a structured decision record.

        This method wraps ``choose_action`` with candidate normalization,
        availability checks, exception handling, and serializable diagnostics.
        It is useful for dashboards, tests, metrics, and narration.

        Args:
            agent: Agent requesting a behavior.
            world: Current world container.
            behaviors: Optional behavior collection.

        Returns:
            A ``PolicyDecision`` object.
        """

        candidates = self.candidate_behaviors(agent, world, behaviors)
        valid_candidates = tuple(candidate for candidate in candidates if candidate.available)

        if not self.enabled:
            return self._decision(
                agent=agent,
                status=PolicyDecisionStatus.DISABLED,
                reason="policy_disabled",
                candidates=candidates,
            )

        if not candidates:
            return self._decision(
                agent=agent,
                status=PolicyDecisionStatus.NO_BEHAVIORS,
                reason="no_behaviors_available",
                candidates=candidates,
            )

        if not valid_candidates:
            return self._decision(
                agent=agent,
                status=PolicyDecisionStatus.NO_VALID_BEHAVIORS,
                reason="no_behaviors_passed_preconditions",
                candidates=candidates,
            )

        try:
            selected_behavior = self.choose_action(
                agent,
                world,
                tuple(candidate.behavior for candidate in valid_candidates),
            )
        except Exception as exc:  # pragma: no cover - defensive guard for subclasses
            logger.exception(
                "Policy %s failed while choosing an action for agent %s",
                self.policy_name,
                _agent_id(agent),
            )
            return self._decision(
                agent=agent,
                status=PolicyDecisionStatus.ERROR,
                reason=str(exc),
                candidates=candidates,
                metadata={"exception_type": exc.__class__.__name__},
            )

        if selected_behavior is None:
            return self._decision(
                agent=agent,
                status=PolicyDecisionStatus.NO_ACTION_SELECTED,
                reason="policy_returned_none",
                candidates=candidates,
            )

        selected_name = self.behavior_name(selected_behavior)

        if not any(candidate.behavior is selected_behavior for candidate in valid_candidates):
            available, reason = self.behavior_is_available(selected_behavior, agent, world)
            if not available:
                return self._decision(
                    agent=agent,
                    status=PolicyDecisionStatus.NO_ACTION_SELECTED,
                    reason=f"selected_behavior_unavailable:{reason}",
                    candidates=candidates,
                    metadata={"selected_behavior_name": selected_name},
                )

        return self._decision(
            agent=agent,
            status=PolicyDecisionStatus.SELECTED,
            selected_behavior=selected_behavior,
            selected_behavior_name=selected_name,
            score=self.score_behavior(agent, world, selected_behavior),
            reason="selected",
            candidates=candidates,
        )

    def available_behaviors(
        self,
        agent: Agent,
        world: World,
        behaviors: BehaviorCollection | None = None,
    ) -> tuple[Behavior, ...]:
        """Return behaviors whose preconditions are currently satisfied."""

        return tuple(
            candidate.behavior
            for candidate in self.candidate_behaviors(agent, world, behaviors)
            if candidate.available
        )

    def candidate_behaviors(
        self,
        agent: Agent,
        world: World,
        behaviors: BehaviorCollection | None = None,
        *,
        validate: bool = True,
    ) -> tuple[BehaviorCandidate, ...]:
        """Return normalized behavior candidates.

        Args:
            agent: Agent requesting a behavior.
            world: Current world container.
            behaviors: Optional behavior collection.
            validate: Whether to call behavior preconditions.

        Returns:
            Tuple of normalized behavior candidates.
        """

        normalized_behaviors = self.normalize_behaviors(agent, behaviors)
        candidates: list[BehaviorCandidate] = []

        for index, behavior in enumerate(normalized_behaviors):
            behavior_name = self.behavior_name(behavior)
            metadata = self.behavior_metadata(behavior)

            if validate:
                available, reason = self.behavior_is_available(behavior, agent, world)
            else:
                available, reason = True, None

            candidates.append(
                BehaviorCandidate(
                    behavior=behavior,
                    name=behavior_name,
                    index=index,
                    available=available,
                    reason=reason,
                    metadata=metadata,
                )
            )

        return tuple(candidates)

    def normalize_behaviors(
        self,
        agent: Agent,
        behaviors: BehaviorCollection | None = None,
    ) -> tuple[Behavior, ...]:
        """Normalize behavior storage into a deterministic tuple.

        The framework allows ``agent.behaviors`` to be represented as either a
        sequence or a mapping. This helper makes policy implementations
        independent of that storage detail.
        """

        raw_behaviors: Any = behaviors
        if raw_behaviors is None:
            raw_behaviors = getattr(agent, "behaviors", ())

        if raw_behaviors is None:
            return ()

        if isinstance(raw_behaviors, Mapping):
            iterable: Iterable[Any] = raw_behaviors.values()
        elif isinstance(raw_behaviors, Iterable) and not isinstance(raw_behaviors, (str, bytes)):
            iterable = raw_behaviors
        else:
            iterable = (raw_behaviors,)

        normalized: list[Behavior] = []
        for item in iterable:
            if item is None:
                continue

            if isinstance(item, Behavior):
                normalized.append(item)
                continue

            if callable(getattr(item, "execute", None)):
                normalized.append(cast(Behavior, item))
                continue

            logger.debug(
                "Skipping non-behavior candidate on agent %s: %r",
                _agent_id(agent),
                item,
            )

        return tuple(normalized)

    def behavior_is_available(
        self,
        behavior: Behavior,
        agent: Agent,
        world: World,
    ) -> tuple[bool, str | None]:
        """Safely check whether a behavior can execute.

        Args:
            behavior: Behavior candidate.
            agent: Agent that may execute the behavior.
            world: Current world container.

        Returns:
            ``(True, None)`` if available, otherwise ``(False, reason)``.
        """

        if not callable(getattr(behavior, "execute", None)):
            return False, "missing_execute_method"

        check_preconditions = getattr(behavior, "check_preconditions", None)
        if not callable(check_preconditions):
            return True, None

        try:
            result = check_preconditions(agent, world)
        except Exception as exc:  # pragma: no cover - defensive guard for arbitrary behaviors
            logger.exception(
                "Behavior %s failed precondition check for agent %s",
                self.behavior_name(behavior),
                _agent_id(agent),
            )
            return False, f"precondition_error:{exc.__class__.__name__}"

        if bool(result):
            return True, None

        return False, "preconditions_not_met"

    def score_behavior(
        self,
        agent: Agent,
        world: World,
        behavior: Behavior,
    ) -> float | None:
        """Return an optional score for a behavior.

        The base policy does not score behaviors, so this returns ``None``.
        Rule-based and learning policies may override it.
        """

        del agent, world, behavior
        return None

    def selection_context(
        self,
        agent: Agent,
        world: World,
        behaviors: BehaviorCollection | None = None,
    ) -> dict[str, Any]:
        """Return lightweight context useful for debugging and logging."""

        candidates = self.candidate_behaviors(agent, world, behaviors)
        return {
            "policy_name": self.policy_name,
            "agent_id": _agent_id(agent),
            "world_step": _world_step(world),
            "candidate_names": [candidate.name for candidate in candidates],
            "available_candidate_names": [
                candidate.name for candidate in candidates if candidate.available
            ],
            "candidate_count": len(candidates),
            "valid_candidate_count": sum(1 for candidate in candidates if candidate.available),
        }

    def update(
        self,
        agent: Agent,
        world: World,
        behavior: Behavior | str | None,
        reward: float,
        context: Sequence[float] | Mapping[str, Any] | None = None,
        metadata: Mapping[str, Any] | None = None,
    ) -> None:
        """Update the policy from reward feedback.

        The base implementation is intentionally a no-op. Learning policies,
        such as contextual bandits, should override this method.

        Args:
            agent: Agent receiving reward feedback.
            world: Current world container.
            behavior: Behavior object or behavior name associated with reward.
            reward: Numeric reward signal.
            context: Optional context vector or context mapping.
            metadata: Optional additional reward metadata.
        """

        del agent, world, behavior, reward, context, metadata

    def receive_reward(
        self,
        reward: float,
        *,
        behavior: Behavior | str | None = None,
        context: Sequence[float] | Mapping[str, Any] | None = None,
        metadata: Mapping[str, Any] | None = None,
    ) -> None:
        """Receive reward feedback when agent/world are not available.

        This no-op hook allows simple agents to forward rewards to a policy
        without requiring every policy to be adaptive.
        """

        del reward, behavior, context, metadata

    def reset(self) -> None:
        """Reset policy-local state.

        Stateless policies do not need to do anything. Adaptive policies should
        override this method to clear learned parameters.
        """

    def state_dict(self) -> dict[str, Any]:
        """Return a serializable snapshot of policy configuration/state."""

        return {
            "name": self.policy_name,
            "enabled": bool(self.enabled),
            "metadata": copy.deepcopy(dict(self.metadata)),
        }

    def load_state_dict(self, state: Mapping[str, Any]) -> None:
        """Load policy configuration/state from a mapping.

        Adaptive policies can override this method to restore learned
        parameters. The base implementation restores only common fields.
        """

        if "enabled" in state:
            self.enabled = bool(state["enabled"])

        metadata = state.get("metadata")
        if isinstance(metadata, Mapping):
            self.metadata = copy.deepcopy(dict(metadata))

    def __call__(
        self,
        agent: Agent,
        world: World,
        behaviors: BehaviorCollection | None = None,
    ) -> Behavior | None:
        """Choose an action, allowing policy objects to be called directly."""

        return self.choose_action(agent, world, behaviors)

    def _decision(
        self,
        *,
        agent: Agent,
        status: PolicyDecisionStatus,
        selected_behavior: Behavior | None = None,
        selected_behavior_name: str | None = None,
        score: float | None = None,
        reason: str | None = None,
        candidates: Sequence[BehaviorCandidate] = (),
        metadata: Mapping[str, Any] | None = None,
    ) -> PolicyDecision:
        """Build a structured policy decision."""

        return PolicyDecision(
            policy_name=self.policy_name,
            agent_id=_agent_id(agent),
            status=status,
            selected_behavior=selected_behavior,
            selected_behavior_name=selected_behavior_name,
            score=score,
            reason=reason,
            candidate_count=len(candidates),
            valid_candidate_count=sum(1 for candidate in candidates if candidate.available),
            candidates=tuple(candidates),
            metadata=metadata or {},
        )

    @staticmethod
    def behavior_name(behavior: Behavior) -> str:
        """Return a stable display name for a behavior."""

        raw_name = getattr(behavior, "name", None)
        if raw_name is not None:
            return str(raw_name)

        return behavior.__class__.__name__

    @staticmethod
    def behavior_metadata(behavior: Behavior) -> Mapping[str, Any]:
        """Return optional behavior metadata as a safe mapping."""

        raw_metadata = getattr(behavior, "metadata", None)
        if isinstance(raw_metadata, Mapping):
            return copy.deepcopy(dict(raw_metadata))

        return {}


Policy = BasePolicy


__all__ = [
    "BasePolicy",
    "BehaviorCandidate",
    "BehaviorCollection",
    "Policy",
    "PolicyDecision",
    "PolicyDecisionStatus",
]