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

Generic agent representation for WorldSmithAI.

This module defines the domain-agnostic Agent runtime object used by the
simulation engine. Agents do not contain species-specific, economy-specific,
civilization-specific, or ecosystem-specific logic. Instead, agents hold state
and delegate behavior execution to Behavior objects and action selection to
Policy objects.

Examples
--------
A DSL-generated world may create agents such as:

    Agent(id="wolf_1", type="wolf", ...)
    Agent(id="scientist_1", type="scientist", ...)
    Agent(id="merchant_1", type="merchant", ...)
    Agent(id="dragon_1", type="dragon", ...)

The Agent class does not know what a wolf, scientist, merchant, or dragon is.
Those meanings emerge from state variables, behaviors, policies, resources,
events, and world rules.

Minimal usage example
---------------------

    import numpy as np

    class MinimalWorld:
        step_count = 0

        def query(self, *, agent, position):
            return {
                "nearby_agents": [],
                "nearby_resources": [],
            }

    class RestBehavior:
        def check_preconditions(self, agent, world):
            return agent.alive

        def execute(self, agent, world):
            return BehaviorExecution(
                success=True,
                reward=0.1,
                state_updates={"energy": agent.state.get("energy", 1.0) + 0.05},
                message="Agent rested.",
            )

    agent = Agent(
        id="agent_1",
        type="generic_entity",
        position=np.array([0.0, 0.0]),
        state={"energy": 1.0},
        behaviors={"rest": RestBehavior()},
    )

    result = agent.step(MinimalWorld())
    print(result.success)
"""

from __future__ import annotations

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

import numpy as np
from numpy.typing import NDArray

logger = logging.getLogger(__name__)

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


class WorldProtocol(Protocol):
    """
    Structural protocol for objects that can be observed by an Agent.

    The concrete World class will be implemented in core/world.py. This protocol
    allows Agent to type-check against a world-like object without importing the
    concrete World class and creating circular dependencies.
    """

    step_count: int

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

        Parameters
        ----------
        agent:
            The observing agent.
        position:
            The current position of the observing agent.

        Returns
        -------
        Mapping[str, Any]
            A mapping containing world information visible to the agent.
        """
        ...


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

    Concrete behavior classes will be implemented in core/behavior.py and the
    behaviors/ package. This protocol keeps Agent decoupled from concrete
    behavior inheritance hierarchies.
    """

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

        Parameters
        ----------
        agent:
            The agent attempting to perform the behavior.
        world:
            The world in which the behavior would execute.

        Returns
        -------
        bool
            True if the behavior can execute, otherwise False.
        """
        ...

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

        Parameters
        ----------
        agent:
            The agent performing the behavior.
        world:
            The world in which the behavior executes.

        Returns
        -------
        BehaviorExecution | Mapping[str, Any] | None
            The behavior outcome. Mapping returns are normalized into
            BehaviorExecution by Agent.
        """
        ...


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

    Concrete policies will be implemented in policies/base_policy.py,
    policies/rule_policy.py, and policies/contextual_bandit.py.
    """

    def choose_action(
        self,
        *,
        agent: Agent,
        world: WorldProtocol,
        observation: Mapping[str, Any],
    ) -> str | BehaviorProtocol | None:
        """
        Choose the next behavior for an agent.

        Parameters
        ----------
        agent:
            The agent requesting an action.
        world:
            The world in which the agent exists.
        observation:
            The agent's current observation.

        Returns
        -------
        str | BehaviorProtocol | None
            A behavior id, behavior object, or None if no action is selected.
        """
        ...

    def update(
        self,
        *,
        agent: Agent,
        reward: float,
        action_id: str | None,
        observation: Mapping[str, Any],
        metadata: Mapping[str, Any],
    ) -> None:
        """
        Update the policy from a reward signal.

        Parameters
        ----------
        agent:
            The agent receiving the reward.
        reward:
            The scalar reward value.
        action_id:
            The id of the behavior that produced the reward.
        observation:
            The observation used when selecting the action.
        metadata:
            Additional behavior-specific reward metadata.
        """
        ...


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

    Behaviors may return this class directly or return a mapping with compatible
    fields. Agent.step() normalizes behavior outputs into this structure.

    Attributes
    ----------
    success:
        Whether the behavior executed successfully.
    reward:
        Scalar reward associated with the behavior result.
    state_updates:
        State changes that should be applied to the agent.
    memory_updates:
        Additional memory payload to append to the agent's memory.
    message:
        Human-readable behavior outcome message.
    metadata:
        Additional structured behavior outcome data.
    """

    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.

        Raises
        ------
        ValueError
            If reward is not finite.
        TypeError
            If update payloads are not mappings.
        """
        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.

    Attributes
    ----------
    agent_id:
        Unique id of the agent that stepped.
    agent_type:
        Domain-level type string, such as "wolf", "merchant", or "scientist".
    step_count:
        Current world step count, if available.
    behavior_id:
        Identifier of the behavior selected for execution.
    success:
        Whether the step completed successfully.
    reward:
        Scalar reward received from the behavior.
    message:
        Human-readable step result.
    metadata:
        Additional structured step result data.
    """

    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.

        Returns
        -------
        dict[str, Any]
            Serializable representation of the step result.
        """
        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.

    Agent stores identity, position, state, goals, memory, available behaviors,
    policy, and lifecycle status. It contains no domain-specific logic and does
    not know what its type means. All domain behavior is delegated to behavior
    objects and policies.

    Attributes
    ----------
    id:
        Unique agent identifier.
    type:
        Domain-level agent type string generated from the DSL.
    position:
        Numeric position vector. The engine does not assume 2D only.
    state:
        Mutable dictionary of agent-specific state variables.
    memory:
        Ordered list of memory entries collected during simulation.
    goals:
        Ordered list of goal descriptors.
    behaviors:
        Mapping from behavior id to behavior object.
    policy:
        Optional policy used to choose behaviors.
    alive:
        Whether the agent is active in the simulation.
    memory_limit:
        Optional maximum number of memory entries retained by the agent.
    """

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

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

        Raises
        ------
        ValueError
            If id, type, position, or memory_limit are invalid.
        TypeError
            If state, memory, goals, or behaviors have invalid types.
        """
        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.")

        normalized_position = self._normalize_position(self.position)
        self.position = normalized_position

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

        if not isinstance(self.memory, list):
            raise TypeError("Agent.memory must be a list.")

        if not isinstance(self.goals, list):
            raise TypeError("Agent.goals must be a list.")

        if not isinstance(self.behaviors, dict):
            raise TypeError("Agent.behaviors 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.")

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

        The step flow is:

        1. Observe the world.
        2. Select a behavior through the policy, or deterministic fallback.
        3. Check behavior preconditions.
        4. Execute the behavior.
        5. Apply state updates.
        6. Store memory.
        7. Forward reward to the policy.

        Parameters
        ----------
        world:
            The simulation world.

        Returns
        -------
        AgentStepResult
            Structured result of this agent's 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,
        )

        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.

        The observation combines the agent's own snapshot with information
        returned by world.query(). If the world query fails, the exception is
        logged and the observation still includes the agent's own state.

        Parameters
        ----------
        world:
            The simulation world.

        Returns
        -------
        Observation
            Dictionary containing self and world observation data.
        """
        observation: Observation = {
            "self": self.snapshot(),
            "world": {},
        }

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

        try:
            world_view = query(agent=self, position=self.position.copy())
        except Exception as exc:
            logger.exception(
                "World query failed for agent_id=%s at position=%s",
                self.id,
                self.position.tolist(),
            )
            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.

        Parameters
        ----------
        updates:
            Mapping of state keys to new values.

        Raises
        ------
        TypeError
            If updates is not a mapping or contains non-string keys.
        """
        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] = 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,
    ) -> None:
        """
        Receive and record a reward signal.

        If the agent has a policy, the reward is forwarded to the policy's
        update method. This supports future reinforcement learning and
        contextual bandit policies without coupling Agent to a specific
        learning algorithm.

        Parameters
        ----------
        reward:
            Scalar reward value.
        action_id:
            Behavior id associated with the reward.
        observation:
            Observation used when selecting the behavior.
        metadata:
            Additional structured reward metadata.

        Raises
        ------
        ValueError
            If reward is not finite.
        """
        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

        try:
            self.policy.update(
                agent=self,
                reward=normalized_reward,
                action_id=action_id,
                observation=normalized_observation,
                metadata=normalized_metadata,
            )
        except Exception:
            logger.exception(
                "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.

        This method is intentionally generic. Movement behaviors can use it
        without Agent knowing what movement means in a specific world.

        Parameters
        ----------
        position:
            New numeric position vector.

        Raises
        ------
        ValueError
            If the position is invalid.
        """
        self.position = self._normalize_position(position)

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

        Parameters
        ----------
        reason:
            Optional lifecycle reason stored in memory.
        """
        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.

        This supports worlds where agents can be reactivated, restored,
        resurrected, restarted, or reintroduced.

        Parameters
        ----------
        reason:
            Optional lifecycle reason stored in memory.
        """
        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.

        Returns
        -------
        dict[str, Any]
            Snapshot of identity, type, position, state, goals, and lifecycle.
        """
        return {
            "id": self.id,
            "type": self.type,
            "position": self.position.tolist(),
            "state": deepcopy(self.state),
            "goals": list(self.goals),
            "alive": self.alive,
        }

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

        Parameters
        ----------
        world:
            The simulation world.
        observation:
            Current agent observation.

        Returns
        -------
        BehaviorProtocol | None
            Selected behavior object, or None.
        """
        if not self.behaviors:
            return None

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

        selected = self.policy.choose_action(
            agent=self,
            world=world,
            observation=observation,
        )

        return self._resolve_behavior(selected)

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

        This deterministic fallback is useful for simple DSL worlds that do not
        provide explicit policies.

        Parameters
        ----------
        world:
            The simulation world.

        Returns
        -------
        BehaviorProtocol | None
            First executable behavior, or None.
        """
        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 or behavior object.

        Parameters
        ----------
        selected:
            Behavior id, behavior object, or None.

        Returns
        -------
        BehaviorProtocol | None
            Resolved behavior object.
        """
        if selected is None:
            return None

        if isinstance(selected, str):
            behavior = self.behaviors.get(selected)
            if behavior is None:
                logger.warning(
                    "Policy selected unknown behavior_id=%s for agent_id=%s",
                    selected,
                    self.id,
                )
            return behavior

        return selected

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

        Parameters
        ----------
        behavior:
            Behavior object.

        Returns
        -------
        str
            Behavior id.
        """
        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.

        Parameters
        ----------
        raw_execution:
            Raw behavior result.

        Returns
        -------
        BehaviorExecution
            Normalized behavior execution object.

        Raises
        ------
        TypeError
            If raw_execution has an unsupported type.
        """
        if raw_execution is None:
            return BehaviorExecution()

        if isinstance(raw_execution, BehaviorExecution):
            return raw_execution

        if isinstance(raw_execution, Mapping):
            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", {})),
                memory_updates=dict(raw_execution.get("memory_updates", {})),
                message=str(raw_execution.get("message", "")),
                metadata=dict(raw_execution.get("metadata", {})),
            )

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

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

        Parameters
        ----------
        entry:
            Memory entry to append.

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

        self.memory.append(dict(entry))

        if self.memory_limit is None:
            return

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

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

        Parameters
        ----------
        position:
            Sequence or numpy array representing a position.

        Returns
        -------
        NDArray[np.float64]
            Normalized position vector.

        Raises
        ------
        ValueError
            If the position is empty, non-finite, or not one-dimensional.
        """
        normalized = np.asarray(position, dtype=np.float64)

        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

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

        Parameters
        ----------
        world:
            Simulation world.

        Returns
        -------
        int | None
            Step count if available.
        """
        step_count = getattr(world, "step_count", None)

        if isinstance(step_count, int):
            return step_count

        return None