"""Support updates for SLZB-06 ESP32 and Zigbee firmwares."""

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

from pysmlight.const import Events as SmEvents
from pysmlight.models import Firmware, Info
from pysmlight.sse import MessageEvent

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

from .const import DOMAIN, LOGGER, ZWAVE_TYPES
from .coordinator import SmConfigEntry, SmFirmwareUpdateCoordinator, SmFwData
from .entity import SmEntity

PARALLEL_UPDATES = 1


def zigbee_latest_version(data: SmFwData, idx: int) -> Firmware | None:
    """Get the latest Zigbee firmware version."""

    if idx < len(data.zb_firmware):
        firmware_list = data.zb_firmware[idx]
        if firmware_list:
            return firmware_list[0]
    return None


@dataclass(frozen=True, kw_only=True)
class SmUpdateEntityDescription(UpdateEntityDescription):
    """Describes SMLIGHT SLZB-06 update entity."""

    installed_version: Callable[[Info, int], str | None]
    latest_version: Callable[[SmFwData, int], Firmware | None]


CORE_UPDATE_ENTITY = SmUpdateEntityDescription(
    key="core_update",
    translation_key="core_update",
    installed_version=lambda x, idx: x.sw_version,
    latest_version=lambda x, idx: x.esp_firmware[0] if x.esp_firmware else None,
)

ZB_UPDATE_ENTITY = SmUpdateEntityDescription(
    key="zigbee_update",
    translation_key="zigbee_update",
    installed_version=lambda x, idx: x.radios[idx].zb_version,
    latest_version=zigbee_latest_version,
)


async def async_setup_entry(
    hass: HomeAssistant,
    entry: SmConfigEntry,
    async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
    """Set up the SMLIGHT update entities."""
    coordinator = entry.runtime_data.firmware

    # updates not available for legacy API, user will get repair to update externally
    if coordinator.legacy_api == 2:
        return

    entities = [SmUpdateEntity(coordinator, CORE_UPDATE_ENTITY)]
    radios = coordinator.data.info.radios

    entities.extend(
        SmUpdateEntity(coordinator, ZB_UPDATE_ENTITY, idx)
        for idx, radio in enumerate(radios)
        if radio.zb_type != -1
    )

    async_add_entities(entities)


class SmUpdateEntity(SmEntity, UpdateEntity):
    """Representation for SLZB-06 update entities."""

    coordinator: SmFirmwareUpdateCoordinator
    entity_description: SmUpdateEntityDescription
    _attr_entity_category = EntityCategory.CONFIG
    _attr_device_class = UpdateDeviceClass.FIRMWARE
    _attr_supported_features = (
        UpdateEntityFeature.INSTALL
        | UpdateEntityFeature.PROGRESS
        | UpdateEntityFeature.RELEASE_NOTES
    )

    def __init__(
        self,
        coordinator: SmFirmwareUpdateCoordinator,
        description: SmUpdateEntityDescription,
        idx: int = 0,
    ) -> None:
        """Initialize the entity."""
        super().__init__(coordinator)

        self.entity_description = description
        device = description.key + (f"_{idx}" if idx else "")
        self._attr_unique_id = f"{coordinator.unique_id}-{device}"

        self._finished_event = asyncio.Event()
        self._firmware: Firmware | None = None
        self._unload: list[Callable] = []
        self.idx = idx

        if (
            (data := coordinator.data)
            and idx < len(data.info.radios)
            and data.info.radios[idx].zb_type in ZWAVE_TYPES
        ):
            if description.key == "zigbee_update":
                self._attr_translation_key = "z_wave_update"

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

    @callback
    @override
    def _handle_coordinator_update(self) -> None:
        """Handle coordinator update callbacks."""
        self._firmware = self.entity_description.latest_version(
            self.coordinator.data, self.idx
        )
        if self._firmware:
            self.async_write_ha_state()

    @property
    @override
    def installed_version(self) -> str | None:
        """Version installed.."""
        data = self.coordinator.data

        return self.entity_description.installed_version(data.info, self.idx)

    @property
    @override
    def latest_version(self) -> str | None:
        """Latest version available for install."""

        return self._firmware.ver if self._firmware else None

    def register_callbacks(self) -> None:
        """Register callbacks for SSE update events."""
        self._unload.append(
            self.coordinator.client.sse.register_callback(
                SmEvents.ZB_FW_prgs, self._update_progress
            )
        )
        self._unload.append(
            self.coordinator.client.sse.register_callback(
                SmEvents.FW_UPD_done, self._update_finished
            )
        )
        if self.coordinator.legacy_api == 1:
            self._unload.append(
                self.coordinator.client.sse.register_callback(
                    SmEvents.ESP_UPD_done, self._update_finished
                )
            )
        self._unload.append(
            self.coordinator.client.sse.register_callback(
                SmEvents.ZB_FW_err, self._update_failed
            )
        )

    @override
    def release_notes(self) -> str | None:
        """Return release notes for firmware."""
        if "zigbee" in self.entity_description.key:
            radio_desc = "Zigbee"
            if (data := self.coordinator.data) and self.idx < len(data.info.radios):
                radio = data.info.radios[self.idx]
                if radio.zb_type in ZWAVE_TYPES:
                    radio_desc = "Z-Wave"
                elif radio.zb_hw and "EFR32" in radio.zb_hw:
                    radio_desc = "EZSP"
                elif radio.zb_hw and "CC2" in radio.zb_hw:
                    radio_desc = "ZNP"
            notes = f"### {radio_desc} Firmware\n\n"
        else:
            notes = "### Core Firmware\n\n"

        if self._firmware and self._firmware.notes:
            notes += self._firmware.notes
            return notes

        return None

    @callback
    def _update_progress(self, progress: MessageEvent) -> None:
        """Update install progress on event."""

        progress = int(progress.data)
        self._attr_update_percentage = progress
        self.async_write_ha_state()

    def _update_done(self) -> None:
        """Handle cleanup for update done."""
        self._finished_event.set()

        for remove_cb in self._unload:
            remove_cb()
        self._unload.clear()

        self._attr_in_progress = False
        self._attr_update_percentage = None
        self.async_write_ha_state()

    @callback
    def _update_finished(self, event: MessageEvent) -> None:
        """Handle event for update finished."""

        self._update_done()

    @callback
    def _update_failed(self, event: MessageEvent) -> None:
        self._update_done()
        self.coordinator.in_progress = False
        raise HomeAssistantError(
            translation_domain=DOMAIN,
            translation_key="firmware_update_failed",
            translation_placeholders={
                "device_name": str(self.name),
            },
        )

    @override
    async def async_install(
        self, version: str | None, backup: bool, **kwargs: Any
    ) -> None:
        """Install firmware update."""

        if not self.coordinator.in_progress and self._firmware:
            self.coordinator.in_progress = True
            self._attr_in_progress = True
            self._attr_update_percentage = None
            self.register_callbacks()

            await self.coordinator.client.fw_update(self._firmware, self.idx)

            # block until update finished event received
            await self._finished_event.wait()

            # allow time for SLZB-06 to reboot before updating coordinator data
            try:
                async with asyncio.timeout(180):
                    while (
                        self.coordinator.in_progress
                        and self.installed_version != self._firmware.ver
                    ):
                        await self.coordinator.async_refresh()
                        await asyncio.sleep(1)
            # pylint: disable-next=home-assistant-action-swallowed-exception
            except TimeoutError:
                LOGGER.warning(
                    "Timeout waiting for %s to reboot after update",
                    self.coordinator.data.info.hostname,
                )

            self.coordinator.in_progress = False
            self._finished_event.clear()
