from __future__ import annotations

from dovla_cil.sim.base import ActionChunk, SimState
from dovla_cil.sim.registry import get_simulator_backend
from dovla_cil.sim.toy_backend import ToyBackend
from dovla_cil.tasks.library import ToyTaskLibrary
from dovla_cil.tasks.predicates import evaluate_task_success


def test_toy_backend_can_reset_task() -> None:
    task = ToyTaskLibrary().get_by_id("toy_put_red_mug_in_blue_bowl")
    sim = ToyBackend()
    sim.seed(123)
    state = sim.reset_task(task)
    observation = sim.render_observation()

    assert isinstance(state, SimState)
    assert state.task_id == task.task_id
    assert "red_mug" in observation["objects"]
    assert observation["robot"]["gripper"] == "open"


def test_serialize_restore_gives_identical_symbolic_state() -> None:
    task = ToyTaskLibrary().get_by_id("toy_put_red_mug_in_blue_bowl")
    sim = ToyBackend()
    sim.seed(123)
    sim.reset_task(task)
    state_blob = sim.serialize_state()
    expected = sim.get_symbolic_state()

    sim.execute_action_chunk(ActionChunk.single("push", object="red_mug", dx=0.2, dy=0.0))
    sim.restore_state(state_blob)

    assert sim.get_symbolic_state() == expected


def test_same_action_from_same_restored_state_is_deterministic() -> None:
    task = ToyTaskLibrary().get_by_id("toy_push_cube_to_target_zone")
    sim = ToyBackend()
    sim.seed(17)
    sim.reset_task(task)
    state_blob = sim.serialize_state()
    action = ActionChunk.single("push", object="cube", dx=0.1, dy=0.2)

    first = sim.execute_action_chunk(action)
    sim.restore_state(state_blob)
    second = sim.execute_action_chunk(action)

    assert first.after_state == second.after_state
    assert first.observation == second.observation
    assert first.contacts == second.contacts
    assert first.reward == second.reward


def test_predicate_success_after_toy_actions() -> None:
    task = ToyTaskLibrary().get_by_id("toy_put_red_mug_in_blue_bowl")
    sim = ToyBackend()
    sim.seed(4)
    sim.reset_task(task)

    result = sim.execute_action_chunk(
        ActionChunk(
            [
                {"command": "move_to", "object": "red_mug"},
                {"command": "grasp", "object": "red_mug"},
                {
                    "command": "place_at",
                    "object": "red_mug",
                    "container": "blue_bowl",
                    "relation": "inside",
                },
            ]
        )
    )

    assert evaluate_task_success(task, sim.get_symbolic_state())
    assert result.info["success"] is True
    assert result.reward == 1.0


def test_registry_toy_backend_works() -> None:
    assert isinstance(get_simulator_backend("toy"), ToyBackend)