"""Backend logic for sending and subscribing to notifications.

Intended to interoperate with HTMX.

Note that this code _doesn't_ handle LLM streaming, which requires a different
approach (you can't use the builtin browser EventSource or htmx SSE support, and
you also don't need the complex Redis channel stuff on the backend).."""

import dataclasses
import inspect
import json
import logging
from dataclasses import dataclass
from datetime import datetime
from typing import Dict, Type

import redis as sync_redis
import redis.asyncio as redis
from django.conf import settings

from promptgame.gameui import redis_utils as ru

# Loose sketch of how I'm going to implement this:
#
# 1. Have an enum of event types, or a collection of different event types, or
#    something. They'll all be declared in this file. Some event types might
#    include metadata, like account numbers or balances or strings.
# 2. Have a helper function that sends events to particular accounts.
# 3. On the receiving side, deserialize the original event.
# 4. Have the frontend listener render HTML in response to the event and shove
#    it down an SSE socket as a new message.
# 5. On the frontend, have HTMX pick up the event and put it in the right place
#    (this may be with, e.g., OOB updates).
#
# I'm putting the rendering as close to the frontend as possible. My theory is
# that this is good because (1) rendering updates probably requires access to
# data related to the user receiving the event, and (2) it makes sense for the
# 'work' to be associated with the user who actually needs the notification, so
# that notifications do not get generated for users who don't need them.

log = logging.getLogger(__name__)


def notifications_channel_name(*, to_account_id: int) -> bytes:
    return f"{settings.REDIS_KEY_PREFIX}:n:{to_account_id}".encode("utf-8")


@dataclass(kw_only=True, frozen=True)
class Event:
    type_identifier: str


@dataclass(kw_only=True, frozen=True)
class NewBalance(Event):
    balance: int
    account_id: int
    type_identifier: str = "new_balance"


@dataclass(kw_only=True, frozen=True)
class LockStateChanged(Event):
    locked: bool
    account_id: int
    unlock_at_isoformat: str
    type_identifier: str = "lock_state_changed"

    @property
    def unlock_at(self) -> datetime:
        return datetime.fromisoformat(self.unlock_at_isoformat)


@dataclass(kw_only=True, frozen=True)
class AttackDetected(Event):
    attacker_account_id: int
    attack_successful: bool
    amount_lost: int
    amount_gained: int
    defender_account_id: int
    type_identifier: str = "attack_detected"


EVENT_TYPES: Dict[str, Type[Event]] = {
    e.type_identifier: e
    for e in globals().values()
    if inspect.isclass(e) and issubclass(e, Event) and e is not Event
}


def event_to_json(event: Event) -> bytes:
    return json.dumps(dataclasses.asdict(event)).encode("utf-8")


def event_from_json(json_data: bytes) -> Event:
    data = json.loads(json_data.decode("utf-8"))
    event_type = data.pop("type_identifier")
    event_cls = EVENT_TYPES[event_type]
    return event_cls(**data)


async def async_send_event(
    event: Event, to_account_id: int, redis_client: redis.Redis
) -> None:
    """Send an event to a particular account."""
    target_channel = notifications_channel_name(to_account_id=to_account_id)
    serialized_event = event_to_json(event)
    await redis_client.publish(target_channel, serialized_event)


def sync_send_event(
    event: Event, to_account_id: int, redis_client: sync_redis.Redis
) -> None:
    """Send an event to a particular account."""
    target_channel = notifications_channel_name(to_account_id=to_account_id)
    serialized_event = event_to_json(event)
    redis_client.publish(target_channel, serialized_event)


class MessageStore:
    """Stores fetches notifications within a context. Example:

    ```
    async with MessageStore(for_account_id=123) as store:
        run_code_that_sends_a_notification(to=123)
        assert len(store.get_messages()) == 1  # make sure we got the message
    ```

    This is useful for tests, where we want to make sure that code is sending
    the right number and type of notifications."""

    def __init__(self, for_account_id: int) -> None:
        self.for_account_id = for_account_id
        self.redis_client = None
        self.redis_pubsub_handle = None

    async def __aenter__(self) -> "MessageStore":
        self.redis_client = ru.get_async_client()
        self.redis_pubsub_handle = self.redis_client.pubsub()
        channel_name = notifications_channel_name(to_account_id=self.for_account_id)
        await self.redis_pubsub_handle.subscribe(channel_name)
        return self

    async def get_messages(self, wait_time: float = 0.0) -> list[Event]:
        """Get all messages received since the last time __aenter__() or
        get_messages() was called (whichever is more recent). If no messages are
        present, it waits for up to wait_time."""
        assert (
            self.redis_pubsub_handle is not None
        ), "Must use this in an async with block"
        messages = []
        timeout = wait_time
        while True:
            # ignore_subscribe_messages=True would mess this up because then we
            # wouldn't be able to distinguish between the case where there are
            # no messages left and the case where we got a subscribe message
            # (Which get_messages() would return as None)
            next_message = await self.redis_pubsub_handle.get_message(
                ignore_subscribe_messages=False, timeout=timeout
            )
            # We only sleep on the first run through. This ensures that the
            # function returns immediately if there are messages, and that the
            # call time doesn't exceed wait_time.
            timeout = 0.0
            if next_message is None:
                break
            if next_message["type"] in ("message", "pmessage"):
                decoded = event_from_json(next_message["data"])
                messages.append(decoded)
        return messages

    async def __aexit__(self, *args) -> None:
        assert self.redis_pubsub_handle is not None
        await self.redis_pubsub_handle.close()
        self.redis_client = None
