"""Async Twitch IRC client primitives.

This module provides a minimal asynchronous IRC client used for connecting to
Twitch chat. It implements low-level primitives to read raw IRC lines,
parse common Twitch messages (tags, prefix, command, params, trailing) and
dispatch typed dataclass messages to registered handlers.

The public surface is the :class:`Client` and message dataclasses at the
bottom of the module. The module is intentionally small — it is meant to be a
building block for higher-level clients.

Examples:
    Create a client, register a handler and run the serve loop::

        client = Client(auth_token="oauth:...", auth_nickname="mybot")

        @client.on(PrivmsgMessage)
        async def on_privmsg(msg):
            print(msg.channel, msg.user, msg.message)

        async def main():
            async with client:
                await client.send(JoinMessage(channel="some_channel"))
                await client.serve()

        asyncio.run(main())

"""

import asyncio as aio
import dataclasses as dc
import typing as t
import logging

from . import errors


_logger = logging.getLogger(__name__)


# TODO: move to utils module
def _to_async(fn: t.Callable[..., t.Any]) -> t.Callable[..., t.Awaitable[t.Any]]:
    if aio.iscoroutinefunction(fn):
        return fn  # type: ignore

    async def async_wrapper(*args: t.Any, **kwargs: t.Any) -> t.Any:
        return fn(*args, **kwargs)

    return async_wrapper


class Client:
    """Minimal asynchronous IRC client for Twitch chat.

    The client provides convenience methods to connect, disconnect and
    read typed messages. Handlers can be registered with :meth:`on` to be
    invoked for specific message dataclasses.

    Example:

        client = Client(auth_token="oauth:...", auth_nickname="mybot")

        @client.on(PrivmsgMessage)
        async def on_privmsg(msg):
            print(msg.channel, msg.user, msg.message)

        async def main():
            async with client:
                await client.send(JoinMessage(channel="some_channel"))
                await client.serve()

        asyncio.run(main())

    """

    def __init__(
        self,
        auth_token: str,
        auth_nickname: str,
    ):
        """Create a new IRC client instance.

        Args:
            auth_token: OAuth token used for PASS during login (include "oauth:" prefix).
            auth_nickname: Nickname used for NICK command.
        """

        self._auth_token = auth_token
        self._auth_nickname = auth_nickname
        self._connection: tuple[aio.StreamReader, aio.StreamWriter] | None = None
        self._listeners: dict[t.Type, list[t.Callable[["Message"], t.Any]]] = {}

    async def __aenter__(self):
        """Enter async context and establish connection.

        Returns:
            The connected :class:`Client` instance.
        """

        await self.connect()
        return self

    async def __aexit__(self, _0, _1, _2):
        """Exit async context and disconnect the client.

        This method suppresses no exceptions; it only attempts to close the
        underlying connection.
        """

        await self.disconnect()

    def on(self, message_type: t.Type["Message"]):
        """Return a decorator that registers a handler for a message type.

        The returned decorator registers ``func`` as a handler that will be
        called when a message of type ``message_type`` is received. Handlers
        may be coroutine functions or regular callables; regular callables are
        wrapped into a coroutine before being scheduled.

        Args:
            message_type: The dataclass type of messages to handle (for
                example :class:`PrivmsgMessage`). Use the base :class:`Message`
                to register a handler that receives every parsed message.

        Returns:
            A decorator which accepts a single-argument callable (the handler)
            and registers it for the given ``message_type``. The decorator
            returns the original function for convenience so it can be used
            directly as a decorator.

        Example:

            @client.on(PrivmsgMessage)
            async def on_privmsg(msg: PrivmsgMessage):
                print(msg.channel, msg.user, msg.message)

            @client.on(Message)
            def on_any(msg: Message):
                # sync handler is supported
                print('received', msg)
        """

        def decorator(func: t.Callable[["Message"], t.Any]):
            self._listeners.setdefault(message_type, []).append(func)
            return func

        return decorator
    
    @property
    def is_connected(self):
        return self._connection is not None
    
    @property
    def is_anonymous(self) -> bool:
        """Return ``True`` if the client is authenticated as an anonymous user.

        An anonymous user is one that uses the nickname ``justinfan123123``.

        Returns:
            ``True`` if the client is authenticated as an anonymous user;
            ``False`` otherwise.
        """

        return self._auth_nickname == "justinfan123123"

    async def connect(self):
        """Open a TCP connection to Twitch IRC and authenticate.

        Raises:
            AlreadyConnectedError: If the client is already connected.
            AuthicationError: If authentication fails.
        """

        self._raise_if_connected()

        self._connection = await aio.open_connection(
            "irc.chat.twitch.tv",
            6667,
        )

        await self.send_raw(
            f"PASS {self._auth_token}\r\n"
            f"NICK {self._auth_nickname}\r\n"
            .encode()
        )

        while True:
            raw = await self.receive_raw()

            if b"NOTICE" in raw:
                raise errors.AuthicationError("Failed to authenticate")
            
            if b"376" in raw:
                break # end of MOTD; success

    async def disconnect(self):
        """Close the connection and release resources.

        Raises:
            NotConnectedError: If the client is not connected.
        """

        self._raise_if_not_connected()
        self._connection[1].close()
        await self._connection[1].wait_closed()
        self._connection = None

    async def receive_raw(self) -> bytes:
        """Read raw bytes from the connection.

        Returns:
            Bytes read from the stream.

        Raises:
            NotConnectedError: If the client is not connected.
        """

        self._raise_if_not_connected()
        receiver, _ = self._connection
        return await receiver.readline()

    async def send_raw(self, message: bytes):
        """Send raw bytes to the IRC server.

        Args:
            message: Raw bytes to write. Should include terminating CRLF when
                required by the IRC protocol.

        Raises:
            NotConnectedError: If the client is not connected.
        """

        self._raise_if_not_connected()
        _, transmitter = self._connection
        transmitter.write(message)
        await transmitter.drain()

    async def receive(self) -> "Message":
        """Read one IRC message and parse it into a dataclass.

        This method reads raw bytes using :meth:`receive_raw`, decodes the line
        and performs a small parser for the common IRC layout::

            [@tags] [:prefix] COMMAND [params] [:trailing]

        The returned object is one of the message dataclasses defined in this
        module (for example, :class:`PrivmsgMessage`). When the parser cannot
        recognize a command a generic :class:`Message` instance is returned.

        Returns:
            A message dataclass representing the parsed IRC line.

        Notes:
            - Tags (the ``@...`` prefix) are parsed into a simple ``dict``
              mapping tag name to value. Tag unescaping is intentionally not
              implemented here for simplicity.
            - Channels are normalized by stripping a leading ``#``.
        """

        raw = await self.receive_raw()
        if raw == b"":
            return Message(raw=raw)

        line = raw.decode(errors="ignore").rstrip("\r\n")

        tags: dict[str, str] = {}
        prefix = None
        trailing = None
        s = line

        # tags: @k=v;... prefix
        if s.startswith("@"):
            try:
                tags_part, s = s[1:].split(" ", 1)
            except ValueError:
                tags_part, s = s[1:], ""
            # no unescaping for simplicity
            for token in tags_part.split(";"):
                if not token:
                    continue
                if "=" in token:
                    k, v = token.split("=", 1)
                else:
                    k, v = token, ""
                tags[k] = v

        if s.startswith(":"):
            try:
                prefix, s = s[1:].split(" ", 1)
            except ValueError:
                prefix = s[1:]
                s = ""

        if " :" in s:
            before, trailing = s.split(" :", 1)
        else:
            before = s

        parts = before.split()
        command = parts[0].upper() if parts else ""
        params = parts[1:]

        user = None
        if prefix and "!" in prefix:
            user = prefix.split("!", 1)[0]

        def normalize_channel(ch: str) -> str:
            if not ch:
                return ""
            return ch.lstrip("#")

        if command == "PING":
            return PingMessage(raw=raw)

        if command == "PRIVMSG":
            channel = params[0] if params else ""
            channel = normalize_channel(channel)
            return PrivmsgMessage(raw=raw, tags=tags, channel=channel, user=user or "", message=trailing or "")

        if command == "PART":
            channel = params[0] if params else ""
            channel = normalize_channel(channel)
            return PartMessage(raw=raw, channel=channel, user=user or "")

        if command == "CLEARCHAT":
            channel = params[0] if params else ""
            channel = normalize_channel(channel)
            # trailing may contain target user
            return ClearChatMessage(raw=raw, tags=tags, channel=channel, user=trailing or "")

        if command == "CLEARMSG":
            channel = params[0] if params else ""
            channel = normalize_channel(channel)
            return ClearMessage(raw=raw, tags=tags, channel=channel, message=trailing or "")

        if command == "NOTICE":
            channel = params[0] if params else ""
            channel = normalize_channel(channel)
            return NoticeMessage(raw=raw, tags=tags, channel=channel, message=trailing or "")

        if command == "USERNOTICE":
            channel = params[0] if params else ""
            channel = normalize_channel(channel)
            return UserNoticeMessage(raw=raw, tags=tags, channel=channel, message=trailing)

        if command == "ROOMSTATE":
            channel = params[0] if params else ""
            channel = normalize_channel(channel)
            return RoomStateMessage(raw=raw, tags=tags, channel=channel)
        
        if command == "USERSTATE":
            channel = params[0] if params else ""
            channel = normalize_channel(channel)
            return UserStateMessage(raw=raw, tags=tags, channel=channel)
        
        if command == "GLOBALUSERSTATE":
            return GlobalUserStateMessage(raw=raw, tags=tags)

        if command == "RECONNECT":
            return ReconnectMessage(raw=raw)

        return Message(raw=raw)
    
    async def send(self, *messages: "Message"):
        """Send one or more message dataclasses to the IRC server.

        This convenience method accepts one or more message dataclass
        instances (for example :class:`JoinMessage`, :class:`PartMessage` or
        :class:`PrivmsgMessage`) and serializes them into the corresponding
        IRC protocol commands. Multiple messages passed to the call are sent
        sequentially.

        Args:
            *messages: One or more message dataclass instances implementing
                the :class:`Sendable` marker. When multiple messages are
                provided they are sent in the order given.

        Raises:
            AssertionError: If a provided object does not implement
                :class:`Sendable`.
            JoinChannelError: If joining a channel fails (server replies do
                not indicate a successful join).
            NotConnectedError: If the client is not connected.

        Example:

            # send a single PRIVMSG
            await client.send(PrivmsgMessage(channel='mychan', message='hi'))

            # send multiple messages in one call
            await client.send(
                JoinMessage(channel='mychan'),
                PrivmsgMessage(channel='mychan', message='hi'),
            )
        """

        if len(messages) == 1:
            message = messages[0]
        else:
            for message in messages:
                await self.send(message)
            return
        
        assert isinstance(message, Sendable), "Only Sendable messages can be sent"
        
        if isinstance(message, JoinMessage):
            await self.send_raw(
                f"JOIN #{message.channel}\r\n"
                .encode()
            )

            for i in range(4):
                raw = await self.receive_raw()
                if i == 0 and b"JOIN" not in raw:
                    raise errors.JoinChannelError()
                elif i == 2 and b"366" in raw:
                    break
        elif isinstance(message, PartMessage):
            await self.send_raw(
                f"PART #{message.channel}\r\n"
                .encode()
            )
        elif isinstance(message, PrivmsgMessage):
            assert not self.is_anonymous, "Nothing happens when sending PRIVMSG as anonymous user"

            await self.send_raw(
                f"PRIVMSG #{message.channel} :{message.message}\r\n"
                .encode()
            )

    async def serve(
        self,
        *,
        strict: bool = False,
        poll_interval: float = 0.1,
    ):
        """Run the main serve loop and dispatch parsed messages to handlers.

        This coroutine runs until cancelled or an unrecoverable exception is
        raised. It repeatedly reads messages using :meth:`receive`, dispatches
        them to registered handlers and performs minimal protocol-level
        processing (for example, replying to PING with PONG).

        The serve loop schedules handlers without awaiting them, so handler
        execution runs concurrently and will not block the main loop. Both
        coroutine and regular callables are supported as handlers; regular
        callables are wrapped into coroutines before scheduling.

        Args:
            strict: If ``True``, parsing or receive errors are raised and the
                serve loop terminates. If ``False`` (the default), parsing
                errors are logged and the loop continues processing subsequent
                messages.
            poll_interval: Time in seconds to sleep between iterations of the
                serve loop when idle. Smaller values make the loop more
                responsive but increase CPU usage; larger values reduce CPU
                usage at the cost of increased latency in dispatching
                handlers.

        Raises:
            Any exception raised from :meth:`receive` when ``strict`` is
            ``True`` is propagated to the caller. Unexpected exceptions are
            re-raised so callers can implement custom shutdown logic.

        Notes:
            - Handlers registered via :meth:`on` are scheduled and executed
              in the background; exceptions raised by handlers are not
              propagated to the serve loop.
        """

        self._raise_if_not_connected()

        try:
            while True:
                try:
                    message = await self.receive()
                except Exception as exc:
                    if strict:
                        raise exc
                    _logger.exception("Error while receiving/parsing message")
                    continue

                _logger.debug("received message: %r", message)

                if isinstance(message, PingMessage):
                    await self.send_raw(b"PONG\r\n")
                    continue

                global_listeners = self._listeners.get(Message, [])
                message_listeners = self._listeners.get(type(message), [])
                message_listeners.extend(global_listeners)

                aio.ensure_future(aio.gather(
                    *(_to_async(L)(message) for L in message_listeners)
                ))

                await aio.sleep(poll_interval)
        except aio.CancelledError:
            _logger.debug("serve loop cancelled")

    def _raise_if_connected(self):
        if self.is_connected:
            raise errors.AlreadyConnectedError()

    def _raise_if_not_connected(self):
        if not self.is_connected:
            raise errors.NotConnectedError()


def create_anonymous_client() -> Client:
    """Create an IRC client authenticated as an anonymous user.

    Returns:
        An :class:`Client` instance authenticated as an anonymous user.
    """

    return Client(
        auth_token="oauth:anonymous",
        auth_nickname="justinfan123123",
    )


T = t.TypeVar("T")
class ReceivableOnlyField(t.Generic[T]):
    """Marker for fields that are only populated when receiving messages."""


@dc.dataclass(kw_only=True)
class Message:
    raw: ReceivableOnlyField[bytes] = b""


class Receivable:
    """Marker for messages that can be received from the IRC server."""


class Sendable:
    """Marker for messages that can be sent to the IRC server."""


@dc.dataclass(kw_only=True)
class MessageWithTags(Message):
    tags: dict[str, str] = dc.field(default_factory=dict)


@dc.dataclass(kw_only=True)
class PingMessage(Message, Receivable):
    pass


@dc.dataclass(kw_only=True)
class ClearChatMessage(MessageWithTags, Receivable):
    channel: str
    user: str


@dc.dataclass(kw_only=True)
class ClearMessage(MessageWithTags, Receivable):
    channel: str
    message: str


@dc.dataclass(kw_only=True)
class GlobalUserStateMessage(MessageWithTags, Receivable):
    pass


@dc.dataclass(kw_only=True)
class NoticeMessage(MessageWithTags, Receivable):
    channel: str
    message: str


@dc.dataclass(kw_only=True)
class JoinMessage(Message, Receivable, Sendable):
    """
    AS RECEIVER:
        Sent when a user has join a chatroom.

    AS SENDER:
        Send to join a chatroom.
    """

    channel: str
    user: ReceivableOnlyField[str] = ""


@dc.dataclass(kw_only=True)
class PartMessage(Message, Receivable, Sendable):
    """
    AS RECEIVER:
        Sent when a user has left a chatroom.

    AS SENDER:
        Send to leave a chatroom.
    """

    channel: str
    user: ReceivableOnlyField[str] = ""


@dc.dataclass(kw_only=True)
class PrivmsgMessage(MessageWithTags, Receivable, Sendable):
    """
    AS RECEIVER:
        Sent when a user posts a chat message in the chat room.
        
    AS SENDER:
        Send to post a message to a chatroom.
    """

    channel: str
    user: ReceivableOnlyField[str] = ""
    message: str


@dc.dataclass(kw_only=True)
class ReconnectMessage(Message, Receivable):
    pass


@dc.dataclass(kw_only=True)
class RoomStateMessage(MessageWithTags, Receivable):
    channel: str


@dc.dataclass(kw_only=True)
class UserNoticeMessage(MessageWithTags, Receivable):
    channel: str
    message: t.Optional[str]


@dc.dataclass(kw_only=True)
class UserStateMessage(MessageWithTags, Receivable):
    channel: str