"""Update platform for ESPHome."""

import asyncio
from typing import Any, override

from aioesphomeapi import (
    DeviceInfo as ESPHomeDeviceInfo,
    EntityInfo,
    UpdateCommand,
    UpdateInfo,
    UpdateState,
)

from homeassistant.components.update import (
    UpdateDeviceClass,
    UpdateEntity,
    UpdateEntityFeature,
)
from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers import device_registry as dr
from homeassistant.helpers.device_registry import DeviceInfo
from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from homeassistant.util.enum import try_parse_enum

from .const import DOMAIN
from .coordinator import ESPHomeDashboardCoordinator
from .dashboard import async_get_dashboard
from .entity import (
    EsphomeEntity,
    async_esphome_state_property,
    convert_api_error_ha_error,
    esphome_state_property,
    platform_async_setup_entry,
)
from .entry_data import ESPHomeConfigEntry, RuntimeEntryData

PARALLEL_UPDATES = 0

KEY_UPDATE_LOCK = "esphome_update_lock"

NO_FEATURES = UpdateEntityFeature(0)


async def async_setup_entry(
    hass: HomeAssistant,
    entry: ESPHomeConfigEntry,
    async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
    """Set up ESPHome update based on a config entry."""
    await platform_async_setup_entry(
        hass,
        entry,
        async_add_entities,
        info_type=UpdateInfo,
        entity_type=ESPHomeUpdateEntity,
        state_type=UpdateState,
    )

    if (dashboard := async_get_dashboard(hass)) is None:
        return
    entry_data = entry.runtime_data
    assert entry_data.device_info is not None
    device_name = entry_data.device_info.name
    unsubs: list[CALLBACK_TYPE] = []

    @callback
    def _async_setup_update_entity() -> None:
        """Set up the update entity."""
        assert dashboard is not None
        # Keep listening until device is available
        if not entry_data.available or not dashboard.last_update_success:
            return

        # Do not add Dashboard Entity if this device is not
        # known to the ESPHome dashboard.
        if dashboard.data is None or dashboard.data.get(device_name) is None:
            return

        for unsub in unsubs:
            unsub()
        unsubs.clear()

        async_add_entities([ESPHomeDashboardUpdateEntity(entry_data, dashboard)])

    if (
        entry_data.available
        and dashboard.last_update_success
        and dashboard.data is not None
        and dashboard.data.get(device_name)
    ):
        _async_setup_update_entity()
        return

    unsubs.extend(
        [
            entry_data.async_subscribe_device_updated(_async_setup_update_entity),
            dashboard.async_add_listener(_async_setup_update_entity),
        ]
    )


class ESPHomeDashboardUpdateEntity(
    CoordinatorEntity[ESPHomeDashboardCoordinator], UpdateEntity
):
    """Defines an ESPHome update entity."""

    _attr_has_entity_name = True
    _attr_device_class = UpdateDeviceClass.FIRMWARE
    _attr_title = "ESPHome"
    _attr_release_url = "https://esphome.io/changelog/"
    _attr_entity_registry_enabled_default = False

    def __init__(
        self, entry_data: RuntimeEntryData, coordinator: ESPHomeDashboardCoordinator
    ) -> None:
        """Initialize the update entity."""
        super().__init__(coordinator=coordinator)
        assert entry_data.device_info is not None
        self._entry_data = entry_data
        self._attr_unique_id = entry_data.device_info.mac_address
        self._attr_device_info = DeviceInfo(
            connections={
                (dr.CONNECTION_NETWORK_MAC, entry_data.device_info.mac_address)
            }
        )
        self._install_lock = asyncio.Lock()
        self._available_future: asyncio.Future[None] | None = None
        self._update_attrs()

    @callback
    def _update_attrs(self) -> None:
        """Update the supported features."""
        coordinator = self.coordinator
        device_info = self._device_info
        # Install support can change at run time
        if coordinator.last_update_success and coordinator.supports_update:
            self._attr_supported_features = UpdateEntityFeature.INSTALL
        else:
            self._attr_supported_features = NO_FEATURES
        self._attr_installed_version = device_info.esphome_version
        device = coordinator.data.get(device_info.name)
        assert device is not None
        self._attr_latest_version = device["current_version"]

    @callback
    @override
    def _handle_coordinator_update(self) -> None:
        """Handle updated data from the coordinator."""
        self._update_attrs()
        super()._handle_coordinator_update()

    @property
    def _device_info(self) -> ESPHomeDeviceInfo:
        """Return the device info."""
        assert self._entry_data.device_info is not None
        return self._entry_data.device_info

    @property
    @override
    def available(self) -> bool:
        """Return if update is available.

        During deep sleep the ESP will not be connectable (by design)
        and thus, even when unavailable, we'll show it as available.
        """
        return super().available and (
            self._entry_data.available
            or self._entry_data.expected_disconnect
            or self._device_info.has_deep_sleep
        )

    @callback
    def _handle_device_update(
        self, static_info: list[EntityInfo] | None = None
    ) -> None:
        """Handle updated data from the device."""
        if (
            self._entry_data.available
            and self._available_future
            and not self._available_future.done()
        ):
            self._available_future.set_result(None)
            self._available_future = None
        self._update_attrs()
        self.async_write_ha_state()

    @override
    async def async_added_to_hass(self) -> None:
        """Handle entity added to Home Assistant."""
        await super().async_added_to_hass()
        entry_data = self._entry_data
        self.async_on_remove(
            entry_data.async_subscribe_static_info_updated(self._handle_device_update)
        )
        self.async_on_remove(
            entry_data.async_subscribe_device_updated(self._handle_device_update)
        )

    @override
    async def async_will_remove_from_hass(self) -> None:
        """Handle entity about to be removed from Home Assistant."""
        if self._available_future and not self._available_future.done():
            self._available_future.cancel()
            self._available_future = None

    async def _async_wait_available(self) -> None:
        """Wait until the device is available."""
        # If the device has deep sleep, we need to wait for it to wake up
        # and connect to the network to be able to install the update.
        if self._entry_data.available:
            return
        self._available_future = self.hass.loop.create_future()
        try:
            await self._available_future
        finally:
            self._available_future = None

    @override
    async def async_install(
        self, version: str | None, backup: bool, **kwargs: Any
    ) -> None:
        """Install an update."""
        if self._install_lock.locked():
            raise HomeAssistantError(
                translation_domain=DOMAIN,
                translation_key="ota_in_progress",
                translation_placeholders={
                    "configuration": self._device_info.name,
                },
            )

        # Ensure only one OTA per device at a time
        async with self._install_lock:
            coordinator = self.coordinator
            api = coordinator.api
            device = coordinator.data.get(self._device_info.name)
            assert device is not None
            configuration = device["configuration"]
            if coordinator.supports_build_queue:
                # The dashboard has its own build queue
                # and can handle concurrent compile requests
                compiled = await api.compile(configuration)
            else:
                # Ensure only one compile at a time for ALL devices
                async with self.hass.data.setdefault(KEY_UPDATE_LOCK, asyncio.Lock()):
                    compiled = await api.compile(configuration)
            if not compiled:
                raise HomeAssistantError(
                    translation_domain=DOMAIN,
                    translation_key="error_compiling",
                    translation_placeholders={
                        "configuration": configuration,
                    },
                )

            # If the device uses deep sleep, there's a small chance it goes
            # to sleep right after the dashboard connects but before the OTA
            # starts. In that case, the update won't go through, so we try
            # again to catch it on its next wakeup.
            attempts = 2 if self._device_info.has_deep_sleep else 1
            try:
                for attempt in range(1, attempts + 1):
                    await self._async_wait_available()
                    if await api.upload(configuration, "OTA"):
                        break
                    if attempt == attempts:
                        raise HomeAssistantError(
                            translation_domain=DOMAIN,
                            translation_key="error_uploading",
                            translation_placeholders={
                                "configuration": configuration,
                            },
                        )
            finally:
                await self.coordinator.async_request_refresh()


class ESPHomeUpdateEntity(EsphomeEntity[UpdateInfo, UpdateState], UpdateEntity):
    """A update implementation for esphome."""

    _attr_supported_features = (
        UpdateEntityFeature.INSTALL
        | UpdateEntityFeature.PROGRESS
        | UpdateEntityFeature.RELEASE_NOTES
    )

    @callback
    @override
    def _on_static_info_update(self, static_info: EntityInfo) -> None:
        """Set attrs from static info."""
        super()._on_static_info_update(static_info)
        static_info = self._static_info
        self._attr_device_class = try_parse_enum(
            UpdateDeviceClass, static_info.device_class
        )

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

        ESPHome project versions can carry a build suffix (e.g.
        2025.11.5_c51f7548) that AwesomeVersion cannot parse. Without stripping
        it the base comparison raises and the entity is forced on for every
        build mismatch. Drop the suffix so the versions compare cleanly and we
        only report genuinely newer firmware.
        """
        return super().version_is_newer(
            latest_version.partition("_")[0], installed_version.partition("_")[0]
        )

    @property
    @esphome_state_property
    @override
    def installed_version(self) -> str:
        """Return the installed version."""
        return self._state.current_version

    @property
    @esphome_state_property
    @override
    def in_progress(self) -> bool:
        """Return if the update is in progress."""
        return self._state.in_progress

    @property
    @esphome_state_property
    @override
    def latest_version(self) -> str | None:
        """Return the latest version."""
        return self._state.latest_version

    @async_esphome_state_property
    @override
    async def async_release_notes(self) -> str | None:
        """Return the release notes."""
        if self._state.release_summary:
            return self._state.release_summary
        return None

    @property
    @esphome_state_property
    @override
    def release_url(self) -> str:
        """Return the release URL."""
        return self._state.release_url

    @property
    @esphome_state_property
    @override
    def title(self) -> str:
        """Return the title of the update."""
        return self._state.title

    @property
    @esphome_state_property
    @override
    def update_percentage(self) -> int | None:
        """Return if the update is in progress."""
        if self._state.has_progress:
            return int(self._state.progress)
        return None

    @convert_api_error_ha_error
    async def async_update(self) -> None:
        """Command device to check for update."""
        if self.available:
            self._client.update_command(
                key=self._key,
                command=UpdateCommand.CHECK,
                device_id=self._static_info.device_id,
            )

    @convert_api_error_ha_error
    @override
    async def async_install(
        self, version: str | None, backup: bool, **kwargs: Any
    ) -> None:
        """Command device to install update."""
        self._client.update_command(
            key=self._key,
            command=UpdateCommand.INSTALL,
            device_id=self._static_info.device_id,
        )
