"""OTA manager for Zigpy. initial implementation from: https://github.com/zigpy/zigpy/pull/1102"""

from __future__ import annotations

import asyncio
from asyncio import timeout as asyncio_timeout
from collections.abc import Callable
import contextlib
from typing import TYPE_CHECKING

import zigpy.datastructures
from zigpy.ota import OTA_FETCH_TIMEOUT
import zigpy.types as t
from zigpy.zcl import ClusterType, foundation
from zigpy.zcl.clusters.general import (
    ImageBlockCommand,
    ImagePageCommand,
    Ota,
    QueryNextImageCommand,
)

if TYPE_CHECKING:
    from typing import Self

    from zigpy.device import Device
    from zigpy.ota.providers import OtaImageWithMetadata


MAX_TIME_WITHOUT_PROGRESS = 30


def _image_block_size_for_manufacturer(
    manufacturer_code: int, requested_size: int
) -> int:
    if manufacturer_code in (4474, 4405):
        # Insta (4474) and Dresden Elektronik (4405) have issues above 40 bytes
        base_size = 40
    elif manufacturer_code == 4129:
        # Legrand (4129) is not capped
        base_size = requested_size
    else:
        # Cap everything else
        base_size = 50

    return min(base_size, requested_size)


class OTAManager:
    """Class to manage OTA updates for a device."""

    def __init__(
        self,
        device: Device,
        image: OtaImageWithMetadata,
        progress_callback=None,
        force: bool = False,
    ) -> None:
        self.device = device
        self.ota_cluster = device.find_cluster(
            cluster_id=Ota.cluster_id, cluster_type=ClusterType.Client
        )

        self.image = image
        self._image_data = image.firmware.serialize()
        self.progress_callback = progress_callback
        self.force = force

        self._upgrade_end_future = asyncio.get_running_loop().create_future()
        self._stall_timer = zigpy.datastructures.ReschedulableTimeout(
            self._stall_callback
        )

        self.stack = contextlib.ExitStack()

    def __enter__(self) -> Self:
        self.stack.enter_context(
            self.device._application.callback_for_response(
                src=self.device,
                filters=[
                    Ota.ServerCommandDefs.query_next_image.schema(),
                ],
                callback=self._image_query_req,
            )
        )

        self.stack.enter_context(
            self.device._application.callback_for_response(
                src=self.device,
                filters=[
                    Ota.ServerCommandDefs.image_block.schema(),
                ],
                callback=self._image_block_req,
            )
        )

        self.stack.enter_context(
            self.device._application.callback_for_response(
                src=self.device,
                filters=[
                    Ota.ServerCommandDefs.image_page.schema(),
                ],
                callback=self._image_page_req,
            )
        )

        self.stack.enter_context(
            self.device._application.callback_for_response(
                src=self.device,
                filters=[
                    Ota.ServerCommandDefs.upgrade_end.schema(),
                ],
                callback=self._upgrade_end,
            )
        )

        return self

    def __exit__(self, *exc_details) -> None:
        self.stack.close()

    def _stall_callback(self) -> None:
        """Handle the stall timer expiring."""
        self._finish(foundation.Status.TIMEOUT)

    def _finish(self, status: foundation.Status) -> None:
        """Finish the OTA process."""
        self._stall_timer.cancel()

        if not self._upgrade_end_future.done():
            self._upgrade_end_future.set_result(status)

    async def _image_query_req(
        self, hdr: foundation.ZCLHeader, command: QueryNextImageCommand
    ) -> None:
        """Handle image query request."""

        # If we try to send a device an old image (e.g. cache issue), don't bother
        if not self.force and (
            not self.image.check_compatibility(self.device, command)
            or not self.image.check_version(command.current_file_version)
        ):
            status = foundation.Status.NO_IMAGE_AVAILABLE
        else:
            status = foundation.Status.SUCCESS

        try:
            await self.ota_cluster.query_next_image_response(
                status=status,
                manufacturer_code=self.image.firmware.header.manufacturer_id,
                image_type=self.image.firmware.header.image_type,
                file_version=self.image.firmware.header.file_version,
                image_size=self.image.firmware.header.image_size,
                tsn=hdr.tsn,
            )
        except Exception as ex:  # noqa: BLE001
            self.device.debug("OTA query_next_image handler exception", exc_info=ex)
            status = foundation.Status.FAILURE

        if status != foundation.Status.SUCCESS:
            self._finish(status)

    async def _finish_malformed_image_block_response(self, handler: str, tsn: int):
        """Create an image block response failure."""
        try:
            await self.ota_cluster.image_block_response(
                status=foundation.Status.MALFORMED_COMMAND, tsn=tsn
            )
        except Exception as ex:  # noqa: BLE001
            self.device.debug(
                "OTA %s handler[MALFORMED_COMMAND] exception", handler, exc_info=ex
            )

        self._finish(foundation.Status.MALFORMED_COMMAND)

    async def _image_block_req(
        self, hdr: foundation.ZCLHeader, command: ImageBlockCommand
    ) -> None:
        """Handle image block request."""
        default_image_block_size = _image_block_size_for_manufacturer(
            command.manufacturer_code, command.maximum_data_size
        )
        block = self._image_data[
            command.file_offset : command.file_offset + default_image_block_size
        ]

        if not block:
            await self._finish_malformed_image_block_response(
                "image_block", tsn=hdr.tsn
            )
            return

        try:
            await self.ota_cluster.image_block_response(
                status=foundation.Status.SUCCESS,
                manufacturer_code=self.image.firmware.header.manufacturer_id,
                image_type=self.image.firmware.header.image_type,
                file_version=self.image.firmware.header.file_version,
                file_offset=command.file_offset,
                image_data=block,
                tsn=hdr.tsn,
            )

            self._stall_timer.reschedule(MAX_TIME_WITHOUT_PROGRESS)

            # Image block requests can sometimes succeed after the device aborts the
            # update. We should not allow the progress callback to be called.
            if (
                self.progress_callback is not None
                and not self._upgrade_end_future.done()
            ):
                self.progress_callback(
                    command.file_offset + len(block), len(self._image_data)
                )
        except Exception as ex:  # noqa: BLE001
            self.device.debug("OTA image_block handler exception", exc_info=ex)

    async def _image_page_req(
        self, hdr: foundation.ZCLHeader, command: ImagePageCommand
    ) -> None:
        """Handle image page request."""
        offset = command.file_offset
        max_block_size = _image_block_size_for_manufacturer(
            command.manufacturer_code, command.maximum_data_size
        )
        bytes_remaining = min(
            command.page_size, len(self._image_data) - command.file_offset
        )

        if bytes_remaining <= 0:
            await self._finish_malformed_image_block_response(
                "image_page_req",
                tsn=hdr.tsn,
            )
            return

        while bytes_remaining > 0:
            block_size = min(max_block_size, bytes_remaining)
            block = self._image_data[offset : offset + block_size]
            offset += block_size
            bytes_remaining -= block_size

            try:
                # OTA sometimes is overwhelmingly fast, we should allow higher priority
                # requests to make it through
                async with self.device.application.request_priority(
                    t.PacketPriority.LOW
                ):
                    # Once we have a way to send requests without waiting for replies,
                    # this can be converted to just `self.ota_cluster.image_block_response`.
                    await self.ota_cluster.request(
                        general=False,
                        command_id=Ota.ClientCommandDefs.image_block_response.id,
                        schema=Ota.ClientCommandDefs.image_block_response.schema,
                        expect_reply=False,
                        # kwargs
                        status=foundation.Status.SUCCESS,
                        manufacturer_code=self.image.firmware.header.manufacturer_id,
                        image_type=self.image.firmware.header.image_type,
                        file_version=self.image.firmware.header.file_version,
                        file_offset=offset - block_size,
                        image_data=block,
                    )

                self._stall_timer.reschedule(MAX_TIME_WITHOUT_PROGRESS)

                if (
                    self.progress_callback is not None
                    and not self._upgrade_end_future.done()
                ):
                    self.progress_callback(
                        offset - block_size + len(block), len(self._image_data)
                    )
            except Exception as ex:  # noqa: BLE001
                self.device.debug("OTA image_page handler exception", exc_info=ex)
                return

            # Delay according to what the device asks
            await asyncio.sleep(command.response_spacing / 1000)

    async def _upgrade_end(
        self, hdr: foundation.ZCLHeader, command: foundation.CommandSchema
    ) -> None:
        """Handle upgrade end request."""
        try:
            await self.ota_cluster.upgrade_end_response(
                manufacturer_code=self.image.firmware.header.manufacturer_id,
                image_type=self.image.firmware.header.image_type,
                file_version=self.image.firmware.header.file_version,
                current_time=0x00000000,
                upgrade_time=0x00000000,
                tsn=hdr.tsn,
            )

            self._finish(command.status)
        except Exception as ex:  # noqa: BLE001
            self.device.debug("OTA upgrade_end handler exception", exc_info=ex)
            self._finish(foundation.Status.FAILURE)

    async def notify(self) -> None:
        """Notify device of new image."""
        try:
            await self.ota_cluster.image_notify(
                payload_type=(
                    self.ota_cluster.ImageNotifyCommand.PayloadType.QueryJitter
                ),
                query_jitter=100,
            )
        except Exception as ex:  # noqa: BLE001
            self.device.debug("OTA image_notify handler exception", exc_info=ex)
            self._finish(foundation.Status.FAILURE)
        else:
            self._stall_timer.reschedule(MAX_TIME_WITHOUT_PROGRESS)

    async def wait(self) -> foundation.Status:
        """Wait for upgrade end response."""
        return await self._upgrade_end_future


async def update_firmware(
    device: Device,
    image: OtaImageWithMetadata,
    progress_callback: Callable[[int, int, float], None] | None = None,
    force: bool = False,
) -> foundation.Status:
    """Update the firmware on a Zigbee device."""
    # Fetch firmware if not already downloaded (deferred download for trusted providers)
    if image.firmware is None:
        async with asyncio_timeout(OTA_FETCH_TIMEOUT):
            image = await image.fetch()

    if force:
        # Force it to send the image even if it's the same version
        image = image.replace(
            metadata=image.metadata.replace(file_version=0xFFFFFFFF - 1),
            firmware=image.firmware.replace(
                header=image.firmware.header.replace(file_version=0xFFFFFFFF - 1)
            ),
        )

    def progress(current: int, total: int):
        progress = (100 * current) / total
        device.info(
            "OTA upgrade progress: (%d / %d): %0.4f%%",
            current,
            total,
            progress,
        )
        if progress_callback is not None:
            progress_callback(current, total, progress)

    # Ask the device to start fast polling before we send the OTA image
    async with device.fast_poll_mode():
        with OTAManager(device, image, progress_callback=progress, force=force) as ota:
            await ota.notify()
            return await ota.wait()
