from __future__ import annotations

import logging


logger = logging.getLogger(__name__)


class DmTracker:
    def __init__(self):
        self._direct_map: dict[str, str] = {}
        self._dm_rooms: set[str] = set()

    def is_direct_room(self, room_id: str) -> bool:
        return room_id in self._dm_rooms

    def get_dm_user(self, room_id: str) -> str | None:
        return self._direct_map.get(room_id)

    def add_direct_room(self, room_id: str, user_id: str) -> None:
        self._dm_rooms.add(room_id)
        self._direct_map[room_id] = user_id

    def remove_direct_room(self, room_id: str) -> None:
        self._dm_rooms.discard(room_id)
        self._direct_map.pop(room_id, None)

    def load_from_account_data(self, m_direct: dict[str, list[str]]) -> None:
        for user_id, room_ids in m_direct.items():
            for room_id in room_ids:
                self._dm_rooms.add(room_id)
                self._direct_map[room_id] = user_id
        logger.debug("Loaded %d DM rooms from m.direct account data", len(self._dm_rooms))

    def mark_candidate_dm(self, room_id: str, inviter_user_id: str) -> None:
        self._dm_rooms.add(room_id)
        self._direct_map[room_id] = inviter_user_id
        logger.debug("Marked room %s as candidate DM with %s", room_id, inviter_user_id)

    @property
    def dm_rooms(self) -> frozenset[str]:
        return frozenset(self._dm_rooms)

    def clear(self) -> None:
        self._dm_rooms.clear()
        self._direct_map.clear()