from __future__ import annotations

from importlib import import_module
from typing import Any, Callable, Iterable, Sequence

from rich.progress import track

RoomViewsByRoom = dict[str, list[Any]]
ProgressTrack = Callable[[Iterable[Any], str], Iterable[Any]]


def collect_room_views_by_room(
    agent: Any,
    sim: Any,
    room_nodes: Sequence[Any],
    views_per_room: int,
    *,
    habitat_sim_module: Any | None = None,
    sensor_uuid: str = "color_sensor",
    turn_action: str = "turn_left",
    progress_description: str = "Collecting room views",
    progress_track: ProgressTrack = track,
) -> RoomViewsByRoom:
    if views_per_room <= 0:
        raise ValueError("views_per_room must be greater than 0")

    if habitat_sim_module is None:
        habitat_sim_module = import_module("habitat_sim")

    all_room_views: RoomViewsByRoom = {}
    for room_node in progress_track(room_nodes, progress_description):
        agent_state = habitat_sim_module.AgentState()
        agent_state.position = room_node.center.copy()
        agent.set_state(agent_state)

        room_views = []
        for _ in range(views_per_room):
            observations = sim.get_sensor_observations()
            room_views.append(observations[sensor_uuid])
            sim.step(turn_action)

        all_room_views[room_node.room_id] = room_views

    return all_room_views