# Copyright 2023 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import asyncio
import inspect
import logging
import struct
from collections.abc import Coroutine
from inspect import getframeinfo, stack
from typing import Any, Dict, Set, Type, TypeVar, overload

import fuchsia_controller_py as fc

from ._fidl_common import (
    FIDL_EPITAPH_ORDINAL,
    Decodable,
    Encodable,
    EpitaphError,
    FidlMessage,
    FidlMeta,
    StopEventHandler,
    TXID_Type,
    parse_epitaph_value,
    parse_ordinal,
    parse_txid,
)
from ._ipc import GlobalHandleWaker, HandleWaker
from ._registry import get_registered_method

# The active TXID (mutable).
TXID: TXID_Type = 0

T = TypeVar("T", bound=Decodable)


# The TXID of a FIDL event.
EVENT_TXID: TXID_Type = 0
# Simple client ID. Monotonically increasing for each client.
_CLIENT_ID = 0
_LOGGER = logging.getLogger("fidl.client")

# FIDL transactional message header size:
# txid (4B) + at-rest flags (2B) + dynamic flags (2B) + magic (1B) + reserved (1B) + ordinal (8B)
_FIDL_MESSAGE_HEADER_SIZE = 16


class FidlClient(metaclass=FidlMeta):
    def __init__(
        self,
        channel: int | fc.Channel,
        channel_waker: HandleWaker | None = None,
    ) -> None:
        global _CLIENT_ID
        self.id = _CLIENT_ID
        _CLIENT_ID += 1
        if isinstance(channel, int):
            self._channel: fc.Channel | None = fc.Channel(channel)
        else:
            self._channel = channel
        if channel_waker is None:
            self._channel_waker: HandleWaker = GlobalHandleWaker()
        else:
            self._channel_waker = channel_waker
        self.pending_txids: Set[TXID_Type] = set({})
        self.staged_messages: Dict[TXID_Type, asyncio.Queue[FidlMessage]] = {}
        self.epitaph_received: EpitaphError | None = None
        self.epitaph_event = asyncio.Event()
        caller = getframeinfo(stack()[1][0])
        _LOGGER.debug(
            f"{self} instantiated from {caller.filename}:{caller.lineno}"
        )

    def close_cleanly(self) -> None:
        """Closes the underlying channel.

        This is so-named to avoid name conflicts with existing FIDL methods.
        """
        if self._channel is not None:
            self._channel.close()
            # self._channel = None is NOT done here, as it may be used in pending tasks.
            # The handle being closed will result in errors in those tasks.

    def __str__(self) -> str:
        return f"client:{type(self).__name__}:{self.id}"

    async def _get_staged_message(self, txid: TXID_Type) -> FidlMessage:
        res = await self.staged_messages[txid].get()
        self.staged_messages[txid].task_done()
        return res

    def _stage_message(self, txid: TXID_Type, msg: FidlMessage) -> None:
        # This should only ever happen if we're a channel reading another channel's response before
        # it has ever made a request.
        if txid not in self.staged_messages:
            self.staged_messages[txid] = asyncio.Queue(1)
        self.staged_messages[txid].put_nowait(msg)

    def _clean_staging(self, txid: TXID_Type) -> None:
        self.staged_messages.pop(txid)
        # Events are never added to this set, since they're always pending.
        if txid != EVENT_TXID:
            self.pending_txids.remove(txid)

    def _decode(
        self,
        txid: TXID_Type,
        msg: FidlMessage,
        response_cls: Type[T] | None,
        expected_ordinal: int | None = None,
    ) -> T | None:
        self._clean_staging(txid)
        if expected_ordinal is not None:
            actual_ordinal = parse_ordinal(msg)
            if actual_ordinal != expected_ordinal:
                type_name = (
                    response_cls.__name__
                    if response_cls is not None
                    else "None"
                )
                raise RuntimeError(
                    f"Ordinal mismatch when decoding {type_name}: expected {expected_ordinal}, got {actual_ordinal}"
                )
        if response_cls is None:
            return None

        handles = msg[1]
        verified_handles: list[int] = [0] * len(handles)
        for i in range(len(handles)):
            hdl = handles[i]
            if isinstance(hdl, tuple):
                verified_handles[i] = hdl[1]
            else:
                verified_handles[i] = hdl.take()
        return response_cls.decode(
            msg[0][_FIDL_MESSAGE_HEADER_SIZE:], verified_handles
        )

    async def next_event(self) -> FidlMessage | None:
        """Attempts to read the next FIDL event from this client.

        Returns:
            The next FIDL event. If ZX_ERR_PEER_CLOSED is received on the channel, will return None.
            Note: this does not check to see if the protocol supports any events, so if not this
            function could wait forever.

        Raises:
            Any exceptions other than ZX_ERR_PEER_CLOSED (fuchsia_controller_py.FcTransportStatus)
        """
        # TODO(awdavies): Raise an exception if there are no events supported for this client.
        try:
            return await self._read_and_decode(0)
        except fc.FcTransportStatus as e:
            err_code = e.code()
            if err_code != fc.FcTransportStatus.FC_ERR_FDOMAIN:
                _LOGGER.warning(
                    f"{self} received error waiting for next event: {e}"
                )
                raise e
        return None

    def _epitaph_check(self, msg: FidlMessage) -> None:
        # If the epitaph is already set, no need to continue with the remaining
        # work.
        if self.epitaph_received is not None:
            raise self.epitaph_received

        ordinal = parse_ordinal(msg)
        if ordinal == FIDL_EPITAPH_ORDINAL:
            if self.epitaph_received is None:
                self.epitaph_received = EpitaphError(parse_epitaph_value(msg))
                self.epitaph_event.set()
            if self.epitaph_received is not None:
                raise self.epitaph_received

    async def _epitaph_event_wait(self) -> None:
        await self.epitaph_event.wait()
        if self.epitaph_received is not None:
            raise self.epitaph_received

    async def _read_and_decode(
        self,
        txid: int,
        response_cls: Type[T] | None = None,
        expected_ordinal: int | None = None,
    ) -> T | FidlMessage | None:
        if txid not in self.staged_messages:
            self.staged_messages[txid] = asyncio.Queue(1)
        if self._channel is None:
            raise ValueError("Channel is already closed")
        msg: FidlMessage
        with self._channel_waker.registration(self._channel, name=str(self)):
            while True:
                if self.epitaph_received is not None:
                    raise self.epitaph_received

                # 1. Try to read from the channel.
                try:
                    msg = self._channel.read()
                    self._epitaph_check(msg)
                    recvd_txid = parse_txid(msg)
                    if recvd_txid == txid:
                        if txid != EVENT_TXID:
                            return self._decode(
                                txid, msg, response_cls, expected_ordinal
                            )
                        else:
                            self._clean_staging(txid)
                            return msg

                    if (
                        recvd_txid != EVENT_TXID
                        and recvd_txid not in self.pending_txids
                    ):
                        _LOGGER.warning(
                            f"{self} received unexpected TXID: {recvd_txid}"
                        )
                        # Unexpected TXID is often a sign of a bad server or a serious bug.
                        # We don't close the channel here, but we raise.
                        raise RuntimeError(
                            f"{self} received unexpected TXID: {recvd_txid}"
                        )

                    self._stage_message(recvd_txid, msg)
                except fc.FcTransportStatus as e:
                    if e.code() != fc.FcTransportStatus.FC_ERR_SHOULD_WAIT:
                        raise e

                # 2. Wait for either:
                #    - The channel being readable again.
                #    - A staged message for our TXID becoming available.
                #    - An epitaph being received.
                #
                #    This is handled with a task group in the event that
                #    there are task cancellations they apply to the other
                #    tasks in the task group, making cleanup predictable.
                try:
                    async with asyncio.TaskGroup() as tg:
                        read_ready_task = tg.create_task(
                            self._channel_waker.wait_ready(self._channel)
                        )
                        staged_msg_task = tg.create_task(
                            self._get_staged_message(txid)
                        )
                        epitaph_task = tg.create_task(self.epitaph_event.wait())

                        done, pending = await asyncio.wait(
                            [read_ready_task, staged_msg_task, epitaph_task],
                            return_when=asyncio.FIRST_COMPLETED,
                        )

                        for p in pending:
                            p.cancel()
                except ExceptionGroup as eg:
                    # If there is only a single exception, raise this as the
                    # main exception. We care particularly if there is an
                    # epitaph exception raised.
                    if len(eg.exceptions) == 1:
                        raise eg.exceptions[0] from None
                    raise

                if staged_msg_task in done:
                    # If we got a staged message, we're done.
                    msg = staged_msg_task.result()
                    # If read_ready_task was also done, we should "re-notify" because we didn't read.
                    if read_ready_task in done:
                        if self._channel is None:
                            raise ValueError("Channel is already closed")
                        self._channel_waker.post_ready(self._channel)

                    if txid != EVENT_TXID:
                        return self._decode(
                            txid, msg, response_cls, expected_ordinal
                        )
                    else:
                        self._clean_staging(txid)
                        return msg

                if epitaph_task in done:
                    # This will raise the epitaph error.
                    if self.epitaph_received is not None:
                        raise self.epitaph_received
                    return None

                # If only read_ready_task reached here, we just loop again and try to read().

    # Overload for two-way FIDL requests with a response payload.
    # Enables static type checking and IDE autocomplete to infer the exact response
    # object type (e.g. `Coroutine[Any, Any, EchoResponse]`) without requiring
    # `typing.cast` or leaving the return type as `Any` or `T | None`.
    # See: https://mypy.readthedocs.io/en/stable/more_types.html#type-checking-calls-to-overloads
    @overload
    def _send_two_way_fidl_request(
        self,
        ordinal: int,
        library: str,
        msg_obj: Encodable | None,
        response_cls: Type[T],
    ) -> Coroutine[Any, Any, T]:
        ...

    # Overload for two-way FIDL requests with an empty response payload (-> ()).
    # Returns `Coroutine[..., None]` statically rather than `Coroutine[..., T | None]`,
    # allowing IDEs and type checkers to know the call resolves to `None`.
    @overload
    def _send_two_way_fidl_request(
        self,
        ordinal: int,
        library: str,
        msg_obj: Encodable | None,
        response_cls: None = None,
    ) -> Coroutine[Any, Any, None]:
        ...

    def _send_two_way_fidl_request(
        self,
        ordinal: int,
        library: str,
        msg_obj: Encodable | None,
        response_cls: Type[T] | None = None,
    ) -> Coroutine[Any, Any, T | None]:
        """Sends a two-way asynchronous FIDL request.

        Args:
            ordinal: The method ordinal (for encoding).
            library: The FIDL library from which this method ordinal exists.
            msg_obj: The object being sent.
            response_cls: The class of the response object.

        Returns:
            The object from the two-way function.
        """
        global TXID
        TXID += 1
        self.pending_txids.add(TXID)
        self._send_one_way_fidl_request(TXID, ordinal, library, msg_obj)

        async def result(txid: int) -> Any:
            return await self._read_and_decode(
                txid, response_cls, expected_ordinal=ordinal
            )

        return result(TXID)

    def _send_one_way_fidl_request(
        self,
        txid: int,
        ordinal: int,
        library: str,
        msg_obj: Encodable | None,
    ) -> None:
        """Sends a synchronous one-way FIDL request.

        Args:
            ordinal: The method ordinal (for encoding).
            library: The FIDL library from which this method ordinal exists.
            msg_obj: The object being sent.
        """
        if msg_obj is not None:
            payload_bytes, handles = msg_obj.encode()
        else:
            payload_bytes, handles = b"", []

        header = struct.pack("<IHBBQ", txid, 0x02, 0x00, 0x01, ordinal)
        encoded_msg = header + payload_bytes

        if self._channel is None:
            raise ValueError("Channel is already closed")
        self._channel.write((encoded_msg, handles))


class EventHandlerBase(
    metaclass=FidlMeta,
    required_class_variables=[
        ("library", str),
        ("method_map", dict),
    ],
):
    """Base object for doing FIDL client event handling."""

    client: FidlClient
    method_map: Dict[int, Any]

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

    def __str__(self) -> str:
        return f"event:{type(self.client).__name__}:{self.client.id}"

    async def serve(self) -> None:
        while True:
            msg = await self.client.next_event()
            # msg is None if the channel has been closed.
            if msg is None:
                break
            if not await self._handle_request(msg):
                break

    async def _handle_request(self, msg: FidlMessage) -> bool:
        try:
            await self._handle_request_helper(msg)
            return True
        except StopEventHandler:
            return False

    async def _handle_request_helper(self, msg: FidlMessage) -> None:
        ordinal = parse_ordinal(msg)
        method = get_registered_method(ordinal)
        if method is None:
            raise RuntimeError(f"Unknown ordinal {ordinal}")
        _, response_cls = method  # Events are registered as response_cls

        handles = msg[1]
        verified_handles: list[int] = [0] * len(handles)
        for i in range(len(handles)):
            hdl = handles[i]
            if isinstance(hdl, tuple):
                verified_handles[i] = hdl[1]
            else:
                verified_handles[i] = hdl.take()

        if response_cls is not None:
            request_obj = response_cls.decode(msg[0][16:], verified_handles)
        else:
            request_obj = None

        method_meta = self.method_map[ordinal]
        method_lambda = getattr(self, method_meta.name)
        if request_obj is not None:
            res = method_lambda(request_obj)
        else:
            res = method_lambda()
        if inspect.isawaitable(res):
            await res