"""Update entities for Shelly devices."""

from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, Final, cast, override

from aioshelly.const import RPC_GENERATIONS
from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError, RpcCallError
from awesomeversion import AwesomeVersion, AwesomeVersionStrategy

from homeassistant.components.update import (
    UpdateDeviceClass,
    UpdateEntity,
    UpdateEntityDescription,
    UpdateEntityFeature,
    UpdateEntityStateAttribute,
)
from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
from homeassistant.helpers.restore_state import RestoreEntity

from .const import (
    CONF_SLEEP_PERIOD,
    DOMAIN,
    LOGGER,
    OTA_BEGIN,
    OTA_ERROR,
    OTA_PROGRESS,
    OTA_SUCCESS,
)
from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
from .entity import (
    RestEntityDescription,
    RpcEntityDescription,
    ShellyRestAttributeEntity,
    ShellyRpcAttributeEntity,
    ShellySleepingRpcAttributeEntity,
    async_setup_entry_rest,
    async_setup_entry_rpc,
)
from .utils import get_device_entry_gen, get_release_url

PARALLEL_UPDATES = 0


@dataclass(frozen=True, kw_only=True)
class RpcUpdateDescription(RpcEntityDescription, UpdateEntityDescription):
    """Class to describe a RPC update."""

    latest_version: Callable[[dict], Any]
    beta: bool


@dataclass(frozen=True, kw_only=True)
class RestUpdateDescription(RestEntityDescription, UpdateEntityDescription):
    """Class to describe a REST update."""

    latest_version: Callable[[dict], Any]
    beta: bool


REST_UPDATES: Final = {
    "fwupdate": RestUpdateDescription(
        key="fwupdate",
        latest_version=lambda status: status["update"]["new_version"],
        beta=False,
        device_class=UpdateDeviceClass.FIRMWARE,
        entity_category=EntityCategory.CONFIG,
        entity_registry_enabled_default=False,
    ),
    "fwupdate_beta": RestUpdateDescription(
        key="fwupdate",
        translation_key="beta_firmware",
        latest_version=lambda status: status["update"].get("beta_version"),
        beta=True,
        device_class=UpdateDeviceClass.FIRMWARE,
        entity_category=EntityCategory.CONFIG,
        entity_registry_enabled_default=False,
    ),
}


class RpcLoraAddOnUpdateEntity(ShellyRpcAttributeEntity, UpdateEntity):
    """Represent a RPC LoRa add-on update entity."""

    _attr_supported_features = (
        UpdateEntityFeature.INSTALL | UpdateEntityFeature.PROGRESS
    )
    entity_description: RpcUpdateDescription

    @property
    def _update_status(self) -> dict[str, Any]:
        """Status of the LoRa add-on update, reported by the lora component."""
        return self.status.get("update") or {}

    @property
    @override
    def installed_version(self) -> str | None:
        """Version currently in use."""
        return cast(str | None, self.status.get("fw_version"))

    @property
    @override
    def latest_version(self) -> str | None:
        """Latest version available for install."""
        status = self.entity_description.latest_version(self.status)
        # After firmware update, available_updates can be None, for a brief moment.
        new_version = (status.get("available_updates") or {}).get(
            "stable", {"version": ""}
        )["version"]
        if new_version:
            return cast(str, new_version)

        return self.installed_version

    @property
    @override
    def in_progress(self) -> bool:
        """Update installation in progress."""
        return bool(self._update_status.get("state") in ("started", "updating"))

    @property
    @override
    def update_percentage(self) -> int | None:
        """Update installation progress."""
        return cast(int | None, self._update_status.get("progress"))

    @override
    async def async_install(
        self, version: str | None, backup: bool, **kwargs: Any
    ) -> None:
        """Install the latest firmware version."""
        update_data = self.status["available_updates"]
        LOGGER.debug("LoRa add-on OTA update service - update_data: %s", update_data)

        new_version = update_data.get("stable", {"version": ""})["version"]

        LOGGER.info(
            "Starting OTA update of LoRa add-on on device %s from '%s' to '%s'",
            self.coordinator.name,
            self.installed_version,
            new_version,
        )
        try:
            await self.coordinator.device.trigger_add_on_ota_update()
        except DeviceConnectionError as err:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="ota_update_connection_error",
                translation_placeholders={"device": self.coordinator.name},
            ) from err
        except RpcCallError as err:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="ota_update_rpc_error",
                translation_placeholders={"device": self.coordinator.name},
            ) from err
        except InvalidAuthError:
            await self.coordinator.async_shutdown_device_and_start_reauth()
        else:
            LOGGER.debug(
                "LoRa add-on OTA update call for %s successful", self.coordinator.name
            )


RPC_UPDATES: Final = {
    "fwupdate": RpcUpdateDescription(
        key="sys",
        sub_key="available_updates",
        latest_version=lambda status: status.get("stable", {"version": ""})["version"],
        beta=False,
        device_class=UpdateDeviceClass.FIRMWARE,
        entity_category=EntityCategory.CONFIG,
    ),
    "fwupdate_beta": RpcUpdateDescription(
        key="sys",
        sub_key="available_updates",
        translation_key="beta_firmware",
        latest_version=lambda status: status.get("beta", {"version": ""})["version"],
        beta=True,
        device_class=UpdateDeviceClass.FIRMWARE,
        entity_category=EntityCategory.CONFIG,
        entity_registry_enabled_default=False,
    ),
    "loraupdate": RpcUpdateDescription(
        key="lora",
        translation_key="lora_firmware",
        # Right after firmware update, lora status can be None
        latest_version=lambda status: status or {},
        beta=False,
        device_class=UpdateDeviceClass.FIRMWARE,
        entity_category=EntityCategory.CONFIG,
        entity_class=RpcLoraAddOnUpdateEntity,
    ),
}


async def async_setup_entry(
    hass: HomeAssistant,
    config_entry: ShellyConfigEntry,
    async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
    """Set up update entities."""
    if get_device_entry_gen(config_entry) in RPC_GENERATIONS:
        return _async_setup_rpc_entry(hass, config_entry, async_add_entities)

    return _async_setup_block_entry(hass, config_entry, async_add_entities)


@callback
def _async_setup_block_entry(
    hass: HomeAssistant,
    config_entry: ShellyConfigEntry,
    async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
    """Set up entities for BLOCK device."""
    if not config_entry.data[CONF_SLEEP_PERIOD]:
        async_setup_entry_rest(
            hass,
            config_entry,
            async_add_entities,
            REST_UPDATES,
            RestUpdateEntity,
        )


@callback
def _async_setup_rpc_entry(
    hass: HomeAssistant,
    config_entry: ShellyConfigEntry,
    async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
    """Set up entities for RPC device."""
    if config_entry.data[CONF_SLEEP_PERIOD]:
        async_setup_entry_rpc(
            hass,
            config_entry,
            async_add_entities,
            RPC_UPDATES,
            RpcSleepingUpdateEntity,
        )
    else:
        async_setup_entry_rpc(
            hass, config_entry, async_add_entities, RPC_UPDATES, RpcUpdateEntity
        )


class RestUpdateEntity(ShellyRestAttributeEntity, UpdateEntity):
    """Represent a REST update entity."""

    _attr_supported_features = (
        UpdateEntityFeature.INSTALL | UpdateEntityFeature.PROGRESS
    )
    entity_description: RestUpdateDescription

    def __init__(
        self,
        block_coordinator: ShellyBlockCoordinator,
        attribute: str,
        description: RestUpdateDescription,
    ) -> None:
        """Initialize update entity."""
        super().__init__(block_coordinator, attribute, description)
        self._attr_release_url = get_release_url(
            block_coordinator.device.gen,
            block_coordinator.model,
            description.beta,
        )
        self._in_progress_old_version: str | None = None

    @property
    @override
    def installed_version(self) -> str | None:
        """Version currently in use."""
        return cast(str, self.block_coordinator.device.status["update"]["old_version"])

    @property
    @override
    def latest_version(self) -> str | None:
        """Latest version available for install."""
        new_version = self.entity_description.latest_version(
            self.block_coordinator.device.status,
        )
        if new_version:
            return cast(str, new_version)

        return self.installed_version

    @property
    @override
    def in_progress(self) -> bool:
        """Update installation in progress."""
        return self._in_progress_old_version == self.installed_version

    @override
    async def async_install(
        self, version: str | None, backup: bool, **kwargs: Any
    ) -> None:
        """Install the latest firmware version."""
        self._in_progress_old_version = self.installed_version
        beta = self.entity_description.beta
        update_data = self.coordinator.device.status["update"]
        LOGGER.debug("OTA update service - update_data: %s", update_data)

        new_version = update_data["new_version"]
        if beta:
            new_version = update_data["beta_version"]

        LOGGER.info(
            "Starting OTA update of device %s from '%s' to '%s'",
            self.name,
            self.coordinator.device.firmware_version,
            new_version,
        )
        try:
            result = await self.coordinator.device.trigger_ota_update(beta=beta)
        except DeviceConnectionError as err:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="ota_update_connection_error",
                translation_placeholders={"device": self.coordinator.name},
            ) from err
        except InvalidAuthError:
            await self.coordinator.async_shutdown_device_and_start_reauth()
        else:
            LOGGER.debug("Result of OTA update call: %s", result)

    @override
    def version_is_newer(self, latest_version: str, installed_version: str) -> bool:
        """Return True if available version is newer then installed version.

        Default strategy generate an exception with Shelly firmware format
        thus making the entity state always true.
        """
        return AwesomeVersion(
            latest_version,
            find_first_match=True,
            ensure_strategy=[AwesomeVersionStrategy.SEMVER],
        ) > AwesomeVersion(
            installed_version,
            find_first_match=True,
            ensure_strategy=[AwesomeVersionStrategy.SEMVER],
        )


class RpcUpdateEntity(ShellyRpcAttributeEntity, UpdateEntity):
    """Represent a RPC update entity."""

    _attr_supported_features = (
        UpdateEntityFeature.INSTALL | UpdateEntityFeature.PROGRESS
    )
    entity_description: RpcUpdateDescription

    def __init__(
        self,
        coordinator: ShellyRpcCoordinator,
        key: str,
        attribute: str,
        description: RpcUpdateDescription,
    ) -> None:
        """Initialize update entity."""
        super().__init__(coordinator, key, attribute, description)
        self._ota_in_progress = False
        self._ota_progress_percentage: int | None = None
        self._attr_release_url = get_release_url(
            coordinator.device.gen, coordinator.model, description.beta
        )

    @override
    async def async_added_to_hass(self) -> None:
        """When entity is added to hass."""
        await super().async_added_to_hass()
        self.async_on_remove(
            self.coordinator.async_subscribe_ota_events(self._ota_progress_callback)
        )

    @callback
    def _ota_progress_callback(self, event: dict[str, Any]) -> None:
        """Handle device OTA progress."""
        if self.in_progress is not False:
            event_type = event["event"]
            if event_type == OTA_BEGIN:
                self._ota_progress_percentage = 0
            elif event_type == OTA_PROGRESS:
                self._ota_progress_percentage = event["progress_percent"]
            elif event_type in (OTA_ERROR, OTA_SUCCESS):
                self._ota_in_progress = False
                self._ota_progress_percentage = None
            self.async_write_ha_state()

    @property
    @override
    def installed_version(self) -> str | None:
        """Version currently in use."""
        return cast(str, self.coordinator.device.shelly["ver"])

    @property
    @override
    def latest_version(self) -> str | None:
        """Latest version available for install."""
        new_version = self.entity_description.latest_version(self.sub_status)
        if new_version:
            return cast(str, new_version)

        return self.installed_version

    @property
    @override
    def in_progress(self) -> bool:
        """Update installation in progress."""
        return self._ota_in_progress

    @property
    @override
    def update_percentage(self) -> int | None:
        """Update installation progress."""
        return self._ota_progress_percentage

    @override
    async def async_install(
        self, version: str | None, backup: bool, **kwargs: Any
    ) -> None:
        """Install the latest firmware version."""
        beta = self.entity_description.beta
        update_data = self.coordinator.device.status["sys"]["available_updates"]
        LOGGER.debug("OTA update service - update_data: %s", update_data)

        new_version = update_data.get("stable", {"version": ""})["version"]
        if beta:
            new_version = update_data.get("beta", {"version": ""})["version"]

        LOGGER.info(
            "Starting OTA update of device %s from '%s' to '%s'",
            self.coordinator.name,
            self.coordinator.device.firmware_version,
            new_version,
        )
        try:
            await self.coordinator.device.trigger_ota_update(beta=beta)
        except DeviceConnectionError as err:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="ota_update_connection_error",
                translation_placeholders={"device": self.coordinator.name},
            ) from err
        except RpcCallError as err:
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="ota_update_rpc_error",
                translation_placeholders={"device": self.coordinator.name},
            ) from err
        except InvalidAuthError:
            await self.coordinator.async_shutdown_device_and_start_reauth()
        else:
            self._ota_in_progress = True
            self._ota_progress_percentage = None
            LOGGER.debug("OTA update call for %s successful", self.coordinator.name)


class RpcSleepingUpdateEntity(
    ShellySleepingRpcAttributeEntity, UpdateEntity, RestoreEntity
):
    """Represent a RPC sleeping update entity."""

    entity_description: RpcUpdateDescription

    @override
    async def async_added_to_hass(self) -> None:
        """Handle entity which will be added."""
        await super().async_added_to_hass()
        self.last_state = await self.async_get_last_state()

    @property
    @override
    def installed_version(self) -> str | None:
        """Version currently in use."""
        if self.coordinator.device.initialized:
            return cast(str, self.coordinator.device.shelly["ver"])

        if self.last_state is None:
            return None

        return self.last_state.attributes.get(
            UpdateEntityStateAttribute.INSTALLED_VERSION
        )

    @property
    @override
    def latest_version(self) -> str | None:
        """Latest version available for install."""
        if self.coordinator.device.initialized:
            new_version = self.entity_description.latest_version(self.sub_status)
            if new_version:
                return cast(str, new_version)

            return self.installed_version

        if self.last_state is None:
            return None

        return self.last_state.attributes.get(UpdateEntityStateAttribute.LATEST_VERSION)

    @property
    @override
    def release_url(self) -> str | None:
        """URL to the full release notes."""
        if not self.coordinator.device.initialized:
            return None

        return get_release_url(
            self.coordinator.device.gen,
            self.coordinator.model,
            self.entity_description.beta,
        )
