# SPDX-FileCopyrightText: Christian Amsüss and the aiocoap contributors
#
# SPDX-License-Identifier: MIT

"""Helpers for the implementation of RFC7959 blockwise transfers"""

from typing import Awaitable, Callable

from . import numbers
from .numbers.optionnumbers import OptionNumber
from .numbers import codes
from .error import ConstructionRenderableError
from .message import Message
from .optiontypes import BlockOption
from .util.asyncio.timeoutdict import TimeoutDict


def _extract_block_key(message):
    """Extract a key that hashes equally for all blocks of a blockwise
    operation from a request message.

    See discussion at <https://mailarchive.ietf.org/arch/msg/core/I-6LzAL6lIUVDA6_g9YM3Zjhg8E>.
    """

    return (
        message.remote.blockwise_key,
        message.code,
        message.get_cache_key(
            [
                OptionNumber.BLOCK1,
                OptionNumber.BLOCK2,
                OptionNumber.OBSERVE,
            ]
        ),
    )


class ContinueException(ConstructionRenderableError):
    """Not an error in the CoAP sense, but an error in the processing sense,
    indicating that no complete request message is available for processing.

    It reflects back the request's block1 option when rendered.
    """

    def __init__(self, block1):
        self.block1 = block1

    def to_message(self):
        m = super().to_message()
        m.opt.block1 = self.block1
        return m

    code = codes.CONTINUE


class IncompleteException(ConstructionRenderableError):
    code = codes.REQUEST_ENTITY_INCOMPLETE


class Block1Spool:
    def __init__(self):
        # FIXME: introduce an actual parameter here
        self._assemblies = TimeoutDict(numbers.TransportTuning().MAX_TRANSMIT_WAIT)

    def feed_and_take(self, req: Message) -> Message:
        """Assemble the request into the spool. This either produces a
        reassembled request message, or raises either a Continue or a Request
        Entity Incomplete exception.

        Requests without block1 are simply passed through."""

        if req.opt.block1 is None:
            return req

        block_key = _extract_block_key(req)

        if req.opt.block1.block_number == 0:
            # silently discarding any old incomplete operation
            self._assemblies[block_key] = req
        else:
            try:
                self._assemblies[block_key]._append_request_block(req)
            except KeyError:
                # KeyError: Received unmatched blockwise response
                # ValueError: Failed to assemble -- gaps or overlaps in data
                raise IncompleteException from None

        if req.opt.block1.more:
            raise ContinueException(req.opt.block1)
        else:
            return self._assemblies[block_key]
            # which happens to carry the last block's block1 option


class Block2Cache:
    """A cache of responses to a give block key.

    Use this when result rendering is expensive, not idempotent or has varying
    output -- otherwise it's often better to calculate the full response again
    and serve chunks.
    """

    def __init__(self):
        # FIXME: introduce an actual parameter here
        self._completes = TimeoutDict(numbers.TransportTuning().MAX_TRANSMIT_WAIT)

    async def extract_or_insert(
        self, req: Message, response_builder: Callable[[], Awaitable[Message]]
    ):
        """Given a request message,

        * if it is querying a particular block, look it up in the cache or
          raise Request Entity Incomplete.
        * otherwise,
          * await the response builder
          * return the response if it doesn't need chunking, or
          * return the first chunk and store it for later use

        """
        block_key = _extract_block_key(req)

        if req.opt.block2 is None or req.opt.block2.block_number == 0:
            assembled = await response_builder()
        else:
            try:
                assembled = self._completes[block_key]
            except KeyError:
                raise IncompleteException from None

        if (
            len(assembled.payload) > req.remote.maximum_payload_size
            or req.opt.block2 is not None
            and len(assembled.payload) > req.opt.block2.size
        ):
            self._completes[block_key] = assembled

            block2 = req.opt.block2 or BlockOption.BlockwiseTuple(
                0, 0, req.remote.maximum_block_size_exp
            )
            return assembled._extract_block(
                block2.block_number,
                block2.size_exponent,
                req.remote.maximum_payload_size,
            )
        else:
            return assembled