import asyncio
import codecs

from collections.abc import AsyncIterator
from typing import TYPE_CHECKING, Optional

from websockets.frames import Frame, Opcode
from websockets.typing import Data

from sanic.exceptions import ServerError


if TYPE_CHECKING:
    from .impl import WebsocketImplProtocol

UTF8Decoder = codecs.getincrementaldecoder("utf-8")


class WebsocketFrameAssembler:
    """
    Assemble a message from frames.
    Code borrowed from aaugustin/websockets project:
    https://github.com/aaugustin/websockets/blob/6eb98dd8fa5b2c896b9f6be7e8d117708da82a39/src/websockets/sync/messages.py
    """

    __slots__ = (
        "protocol",
        "read_mutex",
        "write_mutex",
        "message_complete",
        "message_fetched",
        "get_in_progress",
        "decoder",
        "completed_queue",
        "chunks",
        "chunks_queue",
        "paused",
        "get_id",
        "put_id",
    )
    if TYPE_CHECKING:
        protocol: "WebsocketImplProtocol"
        read_mutex: asyncio.Lock
        write_mutex: asyncio.Lock
        message_complete: asyncio.Event
        message_fetched: asyncio.Event
        completed_queue: asyncio.Queue
        get_in_progress: bool
        decoder: Optional[codecs.IncrementalDecoder]
        # For streaming chunks rather than messages:
        chunks: list[Data]
        chunks_queue: Optional[asyncio.Queue[Optional[Data]]]
        paused: bool

    def __init__(self, protocol) -> None:
        self.protocol = protocol

        self.read_mutex = asyncio.Lock()
        self.write_mutex = asyncio.Lock()

        self.completed_queue = asyncio.Queue(maxsize=1)  # type: asyncio.Queue[Data]

        # put() sets this event to tell get() that a message can be fetched.
        self.message_complete = asyncio.Event()
        # get() sets this event to let put()
        self.message_fetched = asyncio.Event()

        # This flag prevents concurrent calls to get() by user code.
        self.get_in_progress = False

        # Decoder for text frames, None for binary frames.
        self.decoder = None

        # Buffer data from frames belonging to the same message.
        self.chunks = []

        # When switching from "buffering" to "streaming", we use a thread-safe
        # queue for transferring frames from the writing thread (library code)
        # to the reading thread (user code). We're buffering when chunks_queue
        # is None and streaming when it's a Queue. None is a sentinel
        # value marking the end of the stream, superseding message_complete.

        # Stream data from frames belonging to the same message.
        self.chunks_queue = None

        # Flag to indicate we've paused the protocol
        self.paused = False

    async def get(self, timeout: Optional[float] = None) -> Optional[Data]:
        """
        Read the next message.
        :meth:`get` returns a single :class:`str` or :class:`bytes`.
        If the :message was fragmented, :meth:`get` waits until the last frame
        is received, then it reassembles the message.
        If ``timeout`` is set and elapses before a complete message is
        received, :meth:`get` returns ``None``.
        """
        completed: bool
        async with self.read_mutex:
            if timeout is not None and timeout <= 0:
                if not self.message_complete.is_set():
                    return None
            if self.get_in_progress:
                # This should be guarded against with the read_mutex,
                # exception is only here as a failsafe
                raise ServerError(
                    "Called get() on Websocket frame assembler "
                    "while asynchronous get is already in progress."
                )
            self.get_in_progress = True

            # If the message_complete event isn't set yet, release the lock to
            # allow put() to run and eventually set it.
            # Locking with get_in_progress ensures only one task can get here.
            if timeout is None:
                completed = await self.message_complete.wait()
            elif timeout <= 0:
                completed = self.message_complete.is_set()
            else:
                try:
                    await asyncio.wait_for(
                        self.message_complete.wait(), timeout=timeout
                    )
                except asyncio.TimeoutError:
                    ...
                finally:
                    completed = self.message_complete.is_set()

            # Unpause the transport, if its paused
            if self.paused:
                self.protocol.resume_frames()
                self.paused = False
            if not self.get_in_progress:  # no cov
                # This should be guarded against with the read_mutex,
                # exception is here as a failsafe
                raise ServerError(
                    "State of Websocket frame assembler was modified while an "
                    "asynchronous get was in progress."
                )
            self.get_in_progress = False

            # Waiting for a complete message timed out.
            if not completed:
                return None
            if not self.message_complete.is_set():
                return None

            self.message_complete.clear()

            joiner: Data = b"" if self.decoder is None else ""
            # mypy cannot figure out that chunks have the proper type.
            message: Data = joiner.join(self.chunks)  # type: ignore
            if self.message_fetched.is_set():
                # This should be guarded against with the read_mutex,
                # and get_in_progress check, this exception is here
                # as a failsafe
                raise ServerError(
                    "Websocket get() found a message when "
                    "state was already fetched."
                )
            self.message_fetched.set()
            self.chunks = []
            # this should already be None, but set it here for safety
            self.chunks_queue = None
            return message

    async def get_iter(self) -> AsyncIterator[Data]:
        """
        Stream the next message.
        Iterating the return value of :meth:`get_iter` yields a :class:`str`
        or :class:`bytes` for each frame in the message.
        """
        async with self.read_mutex:
            if self.get_in_progress:
                # This should be guarded against with the read_mutex,
                # exception is only here as a failsafe
                raise ServerError(
                    "Called get_iter on Websocket frame assembler "
                    "while asynchronous get is already in progress."
                )
            self.get_in_progress = True

            chunks = self.chunks
            self.chunks = []
            self.chunks_queue = asyncio.Queue()

            # Sending None in chunk_queue supersedes setting message_complete
            # when switching to "streaming". If message is already complete
            # when the switch happens, put() didn't send None, so we have to.
            if self.message_complete.is_set():
                await self.chunks_queue.put(None)

            # Locking with get_in_progress ensures only one task can get here
            for c in chunks:
                yield c
            while True:
                chunk = await self.chunks_queue.get()
                if chunk is None:
                    break
                yield chunk

            # Unpause the transport, if its paused
            if self.paused:
                self.protocol.resume_frames()
                self.paused = False
            if not self.get_in_progress:  # no cov
                # This should be guarded against with the read_mutex,
                # exception is here as a failsafe
                raise ServerError(
                    "State of Websocket frame assembler was modified while an "
                    "asynchronous get was in progress."
                )
            self.get_in_progress = False
            if not self.message_complete.is_set():  # no cov
                # This should be guarded against with the read_mutex,
                # exception is here as a failsafe
                raise ServerError(
                    "Websocket frame assembler chunks queue ended before "
                    "message was complete."
                )
            self.message_complete.clear()
            if self.message_fetched.is_set():  # no cov
                # This should be guarded against with the read_mutex,
                # and get_in_progress check, this exception is
                # here as a failsafe
                raise ServerError(
                    "Websocket get_iter() found a message when state was "
                    "already fetched."
                )

            self.message_fetched.set()
            # this should already be empty, but set it here for safety
            self.chunks = []
            self.chunks_queue = None

    async def put(self, frame: Frame) -> None:
        """
        Add ``frame`` to the next message.
        When ``frame`` is the final frame in a message, :meth:`put` waits
        until the message is fetched, either by calling :meth:`get` or by
        iterating the return value of :meth:`get_iter`.
        :meth:`put` assumes that the stream of frames respects the protocol.
        If it doesn't, the behavior is undefined.
        """

        async with self.write_mutex:
            if frame.opcode is Opcode.TEXT:
                self.decoder = UTF8Decoder(errors="strict")
            elif frame.opcode is Opcode.BINARY:
                self.decoder = None
            elif frame.opcode is Opcode.CONT:
                pass
            else:
                # Ignore control frames.
                return
            data: Data
            if self.decoder is not None:
                data = self.decoder.decode(frame.data, frame.fin)
            else:
                data = frame.data
            if self.chunks_queue is None:
                self.chunks.append(data)
            else:
                await self.chunks_queue.put(data)

            if not frame.fin:
                return
            if not self.get_in_progress:
                # nobody is waiting for this frame, so try to pause subsequent
                # frames at the protocol level
                self.paused = self.protocol.pause_frames()
            # Message is complete. Wait until it's fetched to return.

            if self.chunks_queue is not None:
                await self.chunks_queue.put(None)
            if self.message_complete.is_set():
                # This should be guarded against with the write_mutex
                raise ServerError(
                    "Websocket put() got a new message when a message was "
                    "already in its chamber."
                )
            self.message_complete.set()  # Signal to get() it can serve the
            if self.message_fetched.is_set():
                # This should be guarded against with the write_mutex
                raise ServerError(
                    "Websocket put() got a new message when the previous "
                    "message was not yet fetched."
                )

            # Allow get() to run and eventually set the event.
            await self.message_fetched.wait()
            self.message_fetched.clear()
            self.decoder = None
