from types import SimpleNamespace

import pytest

from simulator import collect_room_views_by_room


class _FakeAgent:
    def __init__(self):
        self.positions = []

    def set_state(self, state):
        self.positions.append(state.position)


class _FakeSimulator:
    def __init__(self):
        self._frame_index = 0
        self.actions = []

    def get_sensor_observations(self):
        observations = {
            "color_sensor": f"frame_{self._frame_index}",
            "depth_sensor": f"depth_{self._frame_index}",
        }
        self._frame_index += 1
        return observations

    def step(self, action_name):
        self.actions.append(action_name)


class _FakeAgentState:
    def __init__(self):
        self.position = None


def test_collect_room_views_by_room_collects_grouped_frames_with_single_outer_progress():
    track_calls = []

    def fake_track(iterable, description):
        track_calls.append(description)
        return iterable

    agent = _FakeAgent()
    sim = _FakeSimulator()
    room_nodes = [
        SimpleNamespace(room_id="room_00", center=[1.0, 2.0, 3.0]),
        SimpleNamespace(room_id="room_01", center=[4.0, 5.0, 6.0]),
    ]
    fake_habitat = SimpleNamespace(AgentState=_FakeAgentState)

    room_views = collect_room_views_by_room(
        agent=agent,
        sim=sim,
        room_nodes=room_nodes,
        views_per_room=3,
        habitat_sim_module=fake_habitat,
        progress_track=fake_track,
    )

    assert track_calls == ["Collecting room views"]
    assert room_views == {
        "room_00": ["frame_0", "frame_1", "frame_2"],
        "room_01": ["frame_3", "frame_4", "frame_5"],
    }
    assert sim.actions == [
        "turn_left",
        "turn_left",
        "turn_left",
        "turn_left",
        "turn_left",
        "turn_left",
    ]
    assert agent.positions == [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]


def test_collect_room_views_by_room_uses_custom_sensor_and_turn_action():
    agent = _FakeAgent()
    sim = _FakeSimulator()
    room_nodes = [SimpleNamespace(room_id="room_00", center=[0.0, 1.0, 0.0])]
    fake_habitat = SimpleNamespace(AgentState=_FakeAgentState)

    room_views = collect_room_views_by_room(
        agent=agent,
        sim=sim,
        room_nodes=room_nodes,
        views_per_room=2,
        habitat_sim_module=fake_habitat,
        sensor_uuid="depth_sensor",
        turn_action="turn_right",
        progress_track=lambda iterable, description: iterable,
    )

    assert room_views == {"room_00": ["depth_0", "depth_1"]}
    assert sim.actions == ["turn_right", "turn_right"]


def test_collect_room_views_by_room_validates_views_per_room():
    with pytest.raises(ValueError, match="views_per_room"):
        collect_room_views_by_room(
            agent=_FakeAgent(),
            sim=_FakeSimulator(),
            room_nodes=[SimpleNamespace(room_id="room_00", center=[0.0, 1.0, 0.0])],
            views_per_room=0,
            habitat_sim_module=SimpleNamespace(AgentState=_FakeAgentState),
            progress_track=lambda iterable, description: iterable,
        )