import sys
from collections.abc import Awaitable
from collections.abc import Callable
from collections.abc import Sequence
from typing import Any
from typing import Literal
from typing import TypedDict
from typing import Union

if sys.version_info >= (3, 11):
    from typing import NotRequired
else:
    from typing_extensions import NotRequired


class AMGIVersions(TypedDict):
    """
    :var version: Version of the AMGI spec
    :var spec_version: Version of the AMGI message spec this server understands-
    """

    spec_version: str
    version: Literal["2.0"]


class MessageScope(TypedDict):
    """
    :var address: The address of the batch of messages, for example, in Kafka this would be the topic
    :var headers: Includes the headers of the message
    :var payload:
        Payload of the message, which can be :py:obj:`None` or :py:obj:`bytes`. If missing, it defaults to
        :py:obj:`None`
    :var bindings:
        Protocol specific bindings, for example, when receiving a Kafka message the bindings could include the key:
        ``{"kafka": {"key": b"key"}}``
    :var state:
        A copy of the namespace passed into the lifespan corresponding to this batch. Optional; if missing the server
        does not support this feature.
    :var extensions:
        Extensions allow AMGI servers to advertise optional capabilities to applications. Extensions are provided via
        scope and are opt-in: applications MUST assume an extension is unsupported unless it is explicitly present.
    """

    type: Literal["message"]
    amgi: AMGIVersions
    address: str
    headers: Sequence[tuple[bytes, bytes]]
    payload: NotRequired[bytes | None]
    bindings: NotRequired[dict[str, dict[str, Any]]]
    state: NotRequired[dict[str, Any]]
    extensions: NotRequired[dict[str, dict[str, Any]]]


class LifespanScope(TypedDict):
    """
    :var state:
        An empty namespace where the application can persist state to be used when handling subsequent requests.
        Optional; if missing the server does not support this feature.
    """

    type: Literal["lifespan"]
    amgi: AMGIVersions
    state: NotRequired[dict[str, Any]]


class LifespanStartupEvent(TypedDict):
    type: Literal["lifespan.startup"]


class LifespanShutdownEvent(TypedDict):
    type: Literal["lifespan.shutdown"]


class LifespanStartupCompleteEvent(TypedDict):
    type: Literal["lifespan.startup.complete"]


class LifespanStartupFailedEvent(TypedDict):
    type: Literal["lifespan.startup.failed"]
    message: str


class LifespanShutdownCompleteEvent(TypedDict):
    type: Literal["lifespan.shutdown.complete"]


class LifespanShutdownFailedEvent(TypedDict):
    type: Literal["lifespan.shutdown.failed"]
    message: str


class MessageAckEvent(TypedDict):
    type: Literal["message.ack"]


class MessageNackEvent(TypedDict):
    type: Literal["message.nack"]
    message: str


class MessageSendEvent(TypedDict):
    """
    :var address: Address to send the message to
    :var headers: Headers of the message
    :var payload:
        Payload of the message, which can be :py:obj:`None`, or :py:obj:`bytes`. If missing, it defaults to
        :py:obj:`None`.
    :var bindings:
        Protocol specific bindings to send. This can be bindings for multiple protocols, allowing the server to decide
        to handle them, or ignore them.
    """

    type: Literal["message.send"]
    address: str
    headers: Sequence[tuple[bytes, bytes]]
    payload: NotRequired[bytes | None]
    bindings: NotRequired[dict[str, dict[str, Any]]]


Scope = Union[MessageScope, LifespanScope]

AMGIReceiveEvent = Union[LifespanStartupEvent, LifespanShutdownEvent]
AMGISendEvent = Union[
    LifespanStartupCompleteEvent,
    LifespanStartupFailedEvent,
    LifespanShutdownCompleteEvent,
    LifespanShutdownFailedEvent,
    MessageAckEvent,
    MessageNackEvent,
    MessageSendEvent,
]

AMGIReceiveCallable = Callable[[], Awaitable[AMGIReceiveEvent]]
AMGISendCallable = Callable[[AMGISendEvent], Awaitable[None]]

AMGIApplication = Callable[
    [
        Scope,
        AMGIReceiveCallable,
        AMGISendCallable,
    ],
    Awaitable[None],
]
