from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Protocol

from dovla_cil.tasks.schema import SceneSpec, TaskSpec

Observation = dict[str, Any]


@dataclass(frozen=True)
class SimState:
    task_id: str
    scene_id: str | None
    seed: int | None
    symbolic_state: dict[str, Any]
    metadata: dict[str, Any] = field(default_factory=dict)


class ActionChunk:
    """Semantic simulator action chunk.

    Examples:
    - `ActionChunk.single("move_to", object="red_mug")`
    - `ActionChunk([{"command": "grasp", "object": "red_mug"}])`
    """

    def __init__(
        self,
        actions: dict[str, Any] | list[dict[str, Any]] | tuple[dict[str, Any], ...] | None = None,
        **single_action: Any,
    ) -> None:
        if actions is None:
            if not single_action:
                raise ValueError("ActionChunk requires at least one action")
            actions = single_action
        if isinstance(actions, dict):
            action_list = [actions]
        else:
            action_list = list(actions)
        if not action_list:
            raise ValueError("ActionChunk requires at least one action")
        normalized: list[dict[str, Any]] = []
        for action in action_list:
            if not isinstance(action, dict):
                raise TypeError("Each semantic action must be a dict")
            if not _action_command(action):
                raise ValueError("Each semantic action requires 'command', 'type', 'op', or 'name'")
            normalized.append(dict(action))
        self.actions = tuple(normalized)

    @classmethod
    def single(cls, command: str, **params: Any) -> ActionChunk:
        return cls({"command": command, **params})

    def to_dict(self) -> dict[str, Any]:
        return {"actions": [dict(action) for action in self.actions]}

    def __iter__(self):
        return iter(self.actions)

    def __eq__(self, other: Any) -> bool:
        return isinstance(other, ActionChunk) and self.actions == other.actions

    def __repr__(self) -> str:
        return f"ActionChunk(actions={self.actions!r})"


@dataclass(frozen=True)
class RolloutResult:
    observation: Observation
    reward: float
    done: bool
    info: dict[str, Any] = field(default_factory=dict)
    trajectory: list[dict[str, Any]] = field(default_factory=list)
    contacts: list[dict[str, Any]] = field(default_factory=list)
    before_state: dict[str, Any] = field(default_factory=dict)
    after_state: dict[str, Any] = field(default_factory=dict)

    @property
    def next_observation(self) -> Observation:
        """Compatibility alias for the earlier scaffold."""

        return self.observation


SimulatorTransition = RolloutResult


class SimulatorBackend(Protocol):
    name: str

    def seed(self, seed: int) -> None:
        """Set deterministic simulator seed."""

    def reset_task(self, task: TaskSpec, scene: SceneSpec | None = None) -> SimState:
        """Reset simulator to a task/scene and return the initial symbolic state."""

    def serialize_state(self) -> bytes:
        """Serialize the exact simulator state."""

    def restore_state(self, state_blob: bytes) -> None:
        """Restore an exact serialized simulator state."""

    def render_observation(self) -> Observation:
        """Render or expose the current observation."""

    def get_symbolic_state(self) -> dict[str, Any]:
        """Return a symbolic state dictionary."""

    def execute_action_chunk(self, action: ActionChunk) -> RolloutResult:
        """Execute an action chunk and return rollout metadata."""

    def close(self) -> None:
        """Release resources."""


def _action_command(action: dict[str, Any]) -> str:
    return str(
        action.get("command") or action.get("type") or action.get("op") or action.get("name") or ""
    )