from __future__ import annotations

from dataclasses import dataclass
from importlib import import_module
from typing import Any


@dataclass(frozen=True)
class HabitatSimulatorConfig:
    scene_path: str
    views_per_room: int = 6
    image_size: int = 512
    sensor_height: float = 1.5
    move_forward_step: float = 0.25
    enable_physics: bool = False
    sensor_uuid: str = "color_sensor"
    agent_id: int = 0
    depth_sensor_uuid: str | None = "depth_sensor"
    hfov_degrees: float = 90.0


def _camera_sensor_spec(
    habitat_sim_module: Any,
    *,
    uuid: str,
    sensor_type: Any,
    image_size: int,
    sensor_height: float,
    hfov_degrees: float,
) -> Any:
    spec = habitat_sim_module.CameraSensorSpec()
    spec.uuid = uuid
    spec.sensor_type = sensor_type
    spec.resolution = [image_size, image_size]
    spec.position = [0.0, sensor_height, 0.0]
    spec.hfov = float(hfov_degrees)
    return spec


def create_habitat_simulator(
    config: HabitatSimulatorConfig,
    habitat_sim_module: Any | None = None,
) -> tuple[Any, Any]:
    if config.views_per_room <= 0:
        raise ValueError("views_per_room must be greater than 0")

    if config.image_size <= 0:
        raise ValueError("image_size must be greater than 0")

    if config.move_forward_step <= 0:
        raise ValueError("move_forward_step must be greater than 0")

    if not (0 < config.hfov_degrees < 180):
        raise ValueError(
            f"hfov_degrees must be in (0, 180), got {config.hfov_degrees}"
        )

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

    sim_cfg = habitat_sim_module.SimulatorConfiguration()
    sim_cfg.scene_id = config.scene_path
    sim_cfg.enable_physics = config.enable_physics

    agent_cfg = habitat_sim_module.agent.AgentConfiguration()

    sensor_specs = [
        _camera_sensor_spec(
            habitat_sim_module,
            uuid=config.sensor_uuid,
            sensor_type=habitat_sim_module.SensorType.COLOR,
            image_size=config.image_size,
            sensor_height=config.sensor_height,
            hfov_degrees=config.hfov_degrees,
        ),
    ]

    if config.depth_sensor_uuid is not None:
        sensor_specs.append(
            _camera_sensor_spec(
                habitat_sim_module,
                uuid=config.depth_sensor_uuid,
                sensor_type=habitat_sim_module.SensorType.DEPTH,
                image_size=config.image_size,
                sensor_height=config.sensor_height,
                hfov_degrees=config.hfov_degrees,
            ),
        )

    agent_cfg.sensor_specifications = sensor_specs

    turn_angle = 360.0 / config.views_per_room
    agent_cfg.action_space = {
        "move_forward": habitat_sim_module.agent.ActionSpec(
            "move_forward",
            habitat_sim_module.agent.ActuationSpec(amount=config.move_forward_step),
        ),
        "turn_left": habitat_sim_module.agent.ActionSpec(
            "turn_left",
            habitat_sim_module.agent.ActuationSpec(amount=turn_angle),
        ),
        "turn_right": habitat_sim_module.agent.ActionSpec(
            "turn_right",
            habitat_sim_module.agent.ActuationSpec(amount=turn_angle),
        ),
    }

    simulator_cfg = habitat_sim_module.Configuration(sim_cfg, [agent_cfg])
    simulator = habitat_sim_module.Simulator(simulator_cfg)
    agent = simulator.initialize_agent(config.agent_id)
    return simulator, agent


def close_habitat_simulator(simulator: Any) -> None:
    close = getattr(simulator, "close", None)
    if callable(close):
        close()