"""World model: causal state transitions for AGI."""

from typing import Any, Protocol

from fusionagi.schemas.plan import Plan
from fusionagi.schemas.world_model import StateTransition, UncertaintyInfo


class WorldModel(Protocol):
    """Protocol for causal model of environment: how actions change state."""

    def predict(self, state: dict[str, Any], action: str, action_args: dict[str, Any]) -> StateTransition:
        """Predict result of action in state."""
        ...

    def uncertainty(self, state: dict[str, Any], action: str) -> UncertaintyInfo:
        """Return uncertainty/risk for action in state."""
        ...


class SimpleWorldModel:
    """
    Minimal world model: state is a dict; actions are recorded but
    prediction returns placeholder. Replace with real causal model.
    """

    def __init__(self) -> None:
        self._transitions: list[StateTransition] = []

    def predict(self, state: dict[str, Any], action: str, action_args: dict[str, Any]) -> StateTransition:
        """Return placeholder transition (state unchanged)."""
        return StateTransition(
            from_state=dict(state),
            action=action,
            action_args=dict(action_args),
            to_state=dict(state),
            confidence=0.5,
        )

    def uncertainty(self, state: dict[str, Any], action: str) -> UncertaintyInfo:
        return UncertaintyInfo(confidence=0.5, risk_level="medium", rationale="SimpleWorldModel placeholder")