from enum import Enum
from typing import Literal, Optional, Type, Union

from alitra import Pose, Position
from pydantic import BaseModel, Field

from robot_interface.models.exceptions.robot_exceptions import ErrorMessage
from robot_interface.models.inspection.inspection import (
    Audio,
    CO2Measurement,
    Image,
    Inspection,
    ThermalImage,
    ThermalVideo,
    Video,
)
from robot_interface.models.mission.status import TaskStatus
from robot_interface.utilities.uuid_string_factory import uuid4_string


class TaskTypes(str, Enum):
    ReturnToHome = "return_to_home"
    TakeImage = "take_image"
    TakeThermalImage = "take_thermal_image"
    TakeVideo = "take_video"
    TakeThermalVideo = "take_thermal_video"
    TakeCO2Measurement = "take_co2_measurement"
    RecordAudio = "record_audio"


class ZoomDescription(BaseModel):
    objectWidth: float
    objectHeight: float


class Task(BaseModel):
    status: TaskStatus = Field(default=TaskStatus.NotStarted)
    error_message: Optional[ErrorMessage] = Field(default=None)
    tag_id: Optional[str] = Field(default=None)
    id: str = Field(default_factory=uuid4_string, frozen=True)


class InspectionTask(Task):
    """
    Base class for all inspection tasks which produce results to be uploaded.
    """

    inspection_id: str = Field(default_factory=uuid4_string, frozen=True)
    robot_pose: Pose = Field(default=None, init=True)
    inspection_description: Optional[str] = Field(default=None)
    zoom: Optional[ZoomDescription] = Field(default=None)

    @staticmethod
    def get_inspection_type() -> Type[Inspection]:
        return Inspection


class ReturnToHome(Task):
    """
    Task which cases the robot to return home
    """

    type: Literal[TaskTypes.ReturnToHome] = TaskTypes.ReturnToHome


class TakeImage(InspectionTask):
    """
    Task which causes the robot to take an image towards the given target.
    """

    target: Position = Field(default=None)
    type: Literal[TaskTypes.TakeImage] = TaskTypes.TakeImage

    @staticmethod
    def get_inspection_type() -> Type[Inspection]:
        return Image


class TakeThermalImage(InspectionTask):
    """
    Task which causes the robot to take a thermal image towards the given target.
    """

    target: Position = Field(default=None)
    type: Literal[TaskTypes.TakeThermalImage] = TaskTypes.TakeThermalImage

    @staticmethod
    def get_inspection_type() -> Type[Inspection]:
        return ThermalImage


class TakeVideo(InspectionTask):
    """
    Task which causes the robot to take a video towards the given target.

    Duration of video is given in seconds.
    """

    target: Position = Field(default=None)
    duration: float = Field(default=None)
    type: Literal[TaskTypes.TakeVideo] = TaskTypes.TakeVideo

    @staticmethod
    def get_inspection_type() -> Type[Inspection]:
        return Video


class TakeThermalVideo(InspectionTask):
    """
    Task which causes the robot to record thermal video towards the given target

    Duration of video is given in seconds.
    """

    target: Position = Field(default=None)
    duration: float = Field(default=None)
    type: Literal[TaskTypes.TakeThermalVideo] = TaskTypes.TakeThermalVideo

    @staticmethod
    def get_inspection_type() -> Type[Inspection]:
        return ThermalVideo


class RecordAudio(InspectionTask):
    """
    Task which causes the robot to record a video at its position, facing the target.

    Duration of audio is given in seconds.
    """

    target: Position = Field(default=None)
    duration: float = Field(default=None)
    type: Literal[TaskTypes.RecordAudio] = TaskTypes.RecordAudio

    @staticmethod
    def get_inspection_type() -> Type[Inspection]:
        return Audio


class TakeCO2Measurement(InspectionTask):
    """
    Task which causes the robot to take a CO2 measurement at its position.
    """

    type: Literal[TaskTypes.TakeCO2Measurement] = TaskTypes.TakeCO2Measurement

    @staticmethod
    def get_inspection_type() -> Type[Inspection]:
        return CO2Measurement


TASKS = Union[
    ReturnToHome,
    TakeImage,
    TakeThermalImage,
    TakeVideo,
    TakeThermalVideo,
    TakeCO2Measurement,
    RecordAudio,
]
