from __future__ import annotations

from enum import Enum

from pydantic import Field, ValidationInfo, field_validator

from common.message import Message, Pose
from common.session.views.articulation import ArticulationConfig
from common.session.views.camera import CameraConfig
from common.session.views.imu import ImuConfig
from common.session.views.pir_sensor import PirSensorConfig
from common.session.views.radar import RadarConfig


class HardwareType(str, Enum):
    """
    Types of hardware that can be attached to an instance.
    """

    CAMERA = "camera"
    IMU = "imu"
    PIR = "pir"
    RADAR = "radar"
    ACTUATOR_GROUP = "actuator_group"


class ActuatorGroupHardware(Message):
    """
    Actuator group hardware that controls multiple joints.
    """

    module: str
    name: str
    config: ArticulationConfig


class ImuHardware(Message):
    """
    IMU sensor hardware attached to a link.
    """

    module: str
    name: str
    path: str
    pose: Pose = Field(default_factory=Pose)
    parent_link: str | None = None
    config: ImuConfig


class PirHardware(Message):
    """
    PIR (Passive Infrared) sensor hardware attached to a link.
    """

    module: str
    name: str
    path: str
    pose: Pose = Field(default_factory=Pose)
    parent_link: str | None = None
    config: PirSensorConfig


class RadarHardware(Message):
    """
    Radar sensor hardware attached to a link.
    """

    module: str
    name: str
    path: str
    pose: Pose = Field(default_factory=Pose)
    parent_link: str | None = None
    config: RadarConfig


class CameraHardware(Message):
    """
    Camera sensor hardware attached to a link.

    Used for both RGB and depth cameras - the mode is specified in the config.
    """

    module: str
    name: str
    path: str
    pose: Pose = Field(default_factory=Pose)
    parent_link: str | None = None
    config: CameraConfig


class Hardware(Message):
    """
    Union of all hardware types for an Ark.
    """

    type: HardwareType
    config: ActuatorGroupHardware | ImuHardware | PirHardware | RadarHardware | CameraHardware

    @field_validator("config", mode="before")
    @classmethod
    def validate_config(cls, config: dict, info: ValidationInfo):
        """
        Validate the hardware config based on the type field.
        """

        hw_type = info.data.get("type")
        if hw_type is None:
            raise ValueError("Hardware type is required")
        match hw_type:
            case HardwareType.ACTUATOR_GROUP:
                return ActuatorGroupHardware.model_validate(config)
            case HardwareType.IMU:
                return ImuHardware.model_validate(config)
            case HardwareType.PIR:
                return PirHardware.model_validate(config)
            case HardwareType.RADAR:
                return RadarHardware.model_validate(config)
            case HardwareType.CAMERA:
                return CameraHardware.model_validate(config)
            case _:
                raise ValueError(f"Unknown hardware type: {hw_type}")
