from __future__ import annotations

import json
import logging
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Collection, Iterator, Optional, Type

from slixmpp import JID, Iq, Message, Presence
from slixmpp.exceptions import XMPPError
from sqlalchemy import Engine, delete, select, update
from sqlalchemy.orm import Session, attributes
from sqlalchemy.sql.functions import count

from ..util.archive_msg import HistoryMessage
from ..util.types import URL, CachedPresence
from ..util.types import Hat as HatTuple
from ..util.types import MamMetadata, MucAffiliation, MucRole
from .meta import Base
from .models import (
    ArchivedMessage,
    ArchivedMessageSource,
    Attachment,
    Avatar,
    Contact,
    ContactSent,
    GatewayUser,
    Hat,
    LegacyIdsMulti,
    Participant,
    Room,
    XmppIdsMulti,
    XmppToLegacyEnum,
    XmppToLegacyIds,
)

if TYPE_CHECKING:
    from ..contact.contact import LegacyContact
    from ..group.participant import LegacyParticipant
    from ..group.room import LegacyMUC


class EngineMixin:
    def __init__(self, engine: Engine):
        self._engine = engine

    @contextmanager
    def session(self, **session_kwargs) -> Iterator[Session]:
        global _session
        if _session is not None:
            yield _session
            return
        with Session(self._engine, **session_kwargs) as session:
            _session = session
            try:
                yield session
            finally:
                _session = None


class UpdatedMixin(EngineMixin):
    model: Type[Base] = NotImplemented

    def __init__(self, *a, **kw):
        super().__init__(*a, **kw)
        with self.session() as session:
            session.execute(
                delete(self.model).where(~self.model.updated)  # type:ignore
            )
            session.execute(update(self.model).values(updated=False))  # type:ignore
            session.commit()

    def get_by_pk(self, pk: int) -> Optional[Base]:
        with self.session() as session:
            return session.execute(
                select(self.model).where(self.model.id == pk)  # type:ignore
            ).scalar()


class SlidgeStore(EngineMixin):
    def __init__(self, engine: Engine):
        super().__init__(engine)
        self.users = UserStore(engine)
        self.avatars = AvatarStore(engine)
        self.contacts = ContactStore(engine)
        self.mam = MAMStore(engine)
        self.multi = MultiStore(engine)
        self.attachments = AttachmentStore(engine)
        self.rooms = RoomStore(engine)
        self.sent = SentStore(engine)
        self.participants = ParticipantStore(engine)


class UserStore(EngineMixin):
    def new(self, jid: JID, legacy_module_data: dict) -> GatewayUser:
        if jid.resource:
            jid = JID(jid.bare)
        with self.session(expire_on_commit=False) as session:
            user = session.execute(
                select(GatewayUser).where(GatewayUser.jid == jid)
            ).scalar()
            if user is not None:
                return user
            user = GatewayUser(jid=jid, legacy_module_data=legacy_module_data)
            session.add(user)
            session.commit()
        return user

    def update(self, user: GatewayUser):
        # https://github.com/sqlalchemy/sqlalchemy/discussions/6473
        attributes.flag_modified(user, "legacy_module_data")
        attributes.flag_modified(user, "preferences")
        with self.session() as session:
            session.add(user)
            session.commit()

    def get_all(self) -> Iterator[GatewayUser]:
        with self.session() as session:
            yield from session.execute(select(GatewayUser)).scalars()

    def get(self, jid: JID) -> Optional[GatewayUser]:
        with self.session() as session:
            return session.execute(
                select(GatewayUser).where(GatewayUser.jid == jid.bare)
            ).scalar()

    def get_by_stanza(self, stanza: Iq | Message | Presence) -> Optional[GatewayUser]:
        return self.get(stanza.get_from())

    def delete(self, jid: JID) -> None:
        with self.session() as session:
            session.delete(self.get(jid))
            session.commit()


class AvatarStore(EngineMixin):
    def get_by_url(self, url: URL | str) -> Optional[Avatar]:
        with self.session() as session:
            return session.execute(select(Avatar).where(Avatar.url == url)).scalar()

    def get_by_legacy_id(self, legacy_id: str) -> Optional[Avatar]:
        with self.session() as session:
            return session.execute(
                select(Avatar).where(Avatar.legacy_id == legacy_id)
            ).scalar()

    def get_by_pk(self, pk: int) -> Optional[Avatar]:
        with self.session() as session:
            return session.execute(select(Avatar).where(Avatar.id == pk)).scalar()

    def delete_by_pk(self, pk: int):
        with self.session() as session:
            session.execute(delete(Avatar).where(Avatar.id == pk))
            session.commit()

    def get_all(self) -> Iterator[Avatar]:
        with self.session() as session:
            yield from session.execute(select(Avatar)).scalars()


class SentStore(EngineMixin):
    def set_message(self, user_pk: int, legacy_id: str, xmpp_id: str) -> None:
        with self.session() as session:
            msg = XmppToLegacyIds(
                user_account_id=user_pk,
                legacy_id=legacy_id,
                xmpp_id=xmpp_id,
                type=XmppToLegacyEnum.DM,
            )
            session.add(msg)
            session.commit()

    def get_xmpp_id(self, user_pk: int, legacy_id: str) -> Optional[str]:
        with self.session() as session:
            return session.execute(
                select(XmppToLegacyIds.xmpp_id)
                .where(XmppToLegacyIds.user_account_id == user_pk)
                .where(XmppToLegacyIds.legacy_id == legacy_id)
                .where(XmppToLegacyIds.type == XmppToLegacyEnum.DM)
            ).scalar()

    def get_legacy_id(self, user_pk: int, xmpp_id: str) -> Optional[str]:
        with self.session() as session:
            return session.execute(
                select(XmppToLegacyIds.legacy_id)
                .where(XmppToLegacyIds.user_account_id == user_pk)
                .where(XmppToLegacyIds.xmpp_id == xmpp_id)
                .where(XmppToLegacyIds.type == XmppToLegacyEnum.DM)
            ).scalar()

    def set_group_message(self, user_pk: int, legacy_id: str, xmpp_id: str) -> None:
        with self.session() as session:
            msg = XmppToLegacyIds(
                user_account_id=user_pk,
                legacy_id=legacy_id,
                xmpp_id=xmpp_id,
                type=XmppToLegacyEnum.GROUP_CHAT,
            )
            session.add(msg)
            session.commit()

    def get_group_xmpp_id(self, user_pk: int, legacy_id: str) -> Optional[str]:
        with self.session() as session:
            return session.execute(
                select(XmppToLegacyIds.xmpp_id)
                .where(XmppToLegacyIds.user_account_id == user_pk)
                .where(XmppToLegacyIds.legacy_id == legacy_id)
                .where(XmppToLegacyIds.type == XmppToLegacyEnum.GROUP_CHAT)
            ).scalar()

    def get_group_legacy_id(self, user_pk: int, xmpp_id: str) -> Optional[str]:
        with self.session() as session:
            return session.execute(
                select(XmppToLegacyIds.legacy_id)
                .where(XmppToLegacyIds.user_account_id == user_pk)
                .where(XmppToLegacyIds.xmpp_id == xmpp_id)
                .where(XmppToLegacyIds.type == XmppToLegacyEnum.GROUP_CHAT)
            ).scalar()

    def set_thread(self, user_pk: int, legacy_id: str, xmpp_id: str) -> None:
        with self.session() as session:
            msg = XmppToLegacyIds(
                user_account_id=user_pk,
                legacy_id=legacy_id,
                xmpp_id=xmpp_id,
                type=XmppToLegacyEnum.THREAD,
            )
            session.add(msg)
            session.commit()

    def get_legacy_thread(self, user_pk: int, xmpp_id: str) -> Optional[str]:
        with self.session() as session:
            return session.execute(
                select(XmppToLegacyIds.legacy_id)
                .where(XmppToLegacyIds.user_account_id == user_pk)
                .where(XmppToLegacyIds.xmpp_id == xmpp_id)
                .where(XmppToLegacyIds.type == XmppToLegacyEnum.THREAD)
            ).scalar()

    def get_xmpp_thread(self, user_pk: int, legacy_id: str) -> Optional[str]:
        with self.session() as session:
            return session.execute(
                select(XmppToLegacyIds.xmpp_id)
                .where(XmppToLegacyIds.user_account_id == user_pk)
                .where(XmppToLegacyIds.legacy_id == legacy_id)
                .where(XmppToLegacyIds.type == XmppToLegacyEnum.THREAD)
            ).scalar()

    def was_sent_by_user(self, user_pk: int, legacy_id: str) -> bool:
        with self.session() as session:
            return (
                session.execute(
                    select(XmppToLegacyIds.legacy_id)
                    .where(XmppToLegacyIds.user_account_id == user_pk)
                    .where(XmppToLegacyIds.legacy_id == legacy_id)
                ).scalar()
                is not None
            )


class ContactStore(UpdatedMixin):
    model = Contact

    def __init__(self, *a, **k):
        super().__init__(*a, **k)
        with self.session() as session:
            session.execute(update(Contact).values(cached_presence=False))
            session.commit()

    def get_all(self, user_pk: int) -> Iterator[Contact]:
        with self.session() as session:
            yield from session.execute(
                select(Contact).where(Contact.user_account_id == user_pk)
            ).scalars()

    def get_by_jid(self, user_pk: int, jid: JID) -> Optional[Contact]:
        with self.session() as session:
            return session.execute(
                select(Contact)
                .where(Contact.jid == jid.bare)
                .where(Contact.user_account_id == user_pk)
            ).scalar()

    def get_by_legacy_id(self, user_pk: int, legacy_id: str) -> Optional[Contact]:
        with self.session() as session:
            return session.execute(
                select(Contact)
                .where(Contact.legacy_id == legacy_id)
                .where(Contact.user_account_id == user_pk)
            ).scalar()

    def update_nick(self, contact_pk: int, nick: Optional[str]) -> None:
        with self.session() as session:
            session.execute(
                update(Contact).where(Contact.id == contact_pk).values(nick=nick)
            )
            session.commit()

    def get_presence(self, contact_pk: int) -> Optional[CachedPresence]:
        with self.session() as session:
            presence = session.execute(
                select(
                    Contact.last_seen,
                    Contact.ptype,
                    Contact.pstatus,
                    Contact.pshow,
                    Contact.cached_presence,
                ).where(Contact.id == contact_pk)
            ).first()
            if presence is None or not presence[-1]:
                return None
            return CachedPresence(*presence[:-1])

    def set_presence(self, contact_pk: int, presence: CachedPresence) -> None:
        with self.session() as session:
            session.execute(
                update(Contact)
                .where(Contact.id == contact_pk)
                .values(**presence._asdict(), cached_presence=True)
            )
            session.commit()

    def reset_presence(self, contact_pk: int):
        with self.session() as session:
            session.execute(
                update(Contact)
                .where(Contact.id == contact_pk)
                .values(
                    last_seen=None,
                    ptype=None,
                    pstatus=None,
                    pshow=None,
                    cached_presence=False,
                )
            )
            session.commit()

    def set_avatar(self, contact_pk: int, avatar_pk: Optional[int]):
        with self.session() as session:
            session.execute(
                update(Contact)
                .where(Contact.id == contact_pk)
                .values(avatar_id=avatar_pk)
            )
            session.commit()

    def get_avatar_legacy_id(self, contact_pk: int) -> Optional[str]:
        with self.session() as session:
            contact = session.execute(
                select(Contact).where(Contact.id == contact_pk)
            ).scalar()
            if contact is None or contact.avatar is None:
                return None
            return contact.avatar.legacy_id

    def update(self, contact: "LegacyContact", commit=True) -> int:
        with self.session() as session:
            if contact.contact_pk is None:
                if contact.cached_presence is not None:
                    presence_kwargs = contact.cached_presence._asdict()
                    presence_kwargs["cached_presence"] = True
                else:
                    presence_kwargs = {}
                row = Contact(
                    jid=contact.jid.bare,
                    legacy_id=str(contact.legacy_id),
                    user_account_id=contact.user_pk,
                    **presence_kwargs,
                )
            else:
                row = (
                    session.query(Contact)
                    .filter(Contact.id == contact.contact_pk)
                    .one()
                )
            row.nick = contact.name
            row.is_friend = contact.is_friend
            row.added_to_roster = contact.added_to_roster
            row.updated = True
            row.extra_attributes = contact.serialize_extra_attributes()
            row.caps_ver = contact._caps_ver
            session.add(row)
            if commit:
                session.commit()
            return row.id

    def add_to_sent(self, contact_pk: int, msg_id: str) -> None:
        with self.session() as session:
            new = ContactSent(contact_id=contact_pk, msg_id=msg_id)
            session.add(new)
            session.commit()

    def pop_sent_up_to(self, contact_pk: int, msg_id: str) -> list[str]:
        result = []
        to_del = []
        with self.session() as session:
            for row in session.execute(
                select(ContactSent)
                .where(ContactSent.contact_id == contact_pk)
                .order_by(ContactSent.id)
            ).scalars():
                to_del.append(row.id)
                result.append(row.msg_id)
                if row.msg_id == msg_id:
                    break
            for row_id in to_del:
                session.execute(delete(ContactSent).where(ContactSent.id == row_id))
        return result

    def set_friend(self, contact_pk: int, is_friend: bool) -> None:
        with self.session() as session:
            session.execute(
                update(Contact)
                .where(Contact.id == contact_pk)
                .values(is_friend=is_friend)
            )
            session.commit()

    def set_added_to_roster(self, contact_pk: int, value: bool) -> None:
        with self.session() as session:
            session.execute(
                update(Contact)
                .where(Contact.id == contact_pk)
                .values(added_to_roster=value)
            )
            session.commit()

    def delete(self, contact_pk: int) -> None:
        with self.session() as session:
            session.execute(delete(Contact).where(Contact.id == contact_pk))
            session.commit()


class MAMStore(EngineMixin):
    def __init__(self, *a, **kw):
        super().__init__(*a, **kw)
        with self.session() as session:
            session.execute(
                update(ArchivedMessage).values(source=ArchivedMessageSource.BACKFILL)
            )
            session.commit()

    def nuke_older_than(self, days: int) -> None:
        with self.session() as session:
            session.execute(
                delete(ArchivedMessage).where(
                    ArchivedMessage.timestamp < datetime.now() - timedelta(days=days)
                )
            )
            session.commit()

    def add_message(
        self,
        room_pk: int,
        message: HistoryMessage,
        archive_only: bool,
        legacy_msg_id: str | None,
    ) -> None:
        with self.session() as session:
            source = (
                ArchivedMessageSource.BACKFILL
                if archive_only
                else ArchivedMessageSource.LIVE
            )
            existing = session.execute(
                select(ArchivedMessage)
                .where(ArchivedMessage.room_id == room_pk)
                .where(ArchivedMessage.stanza_id == message.id)
            ).scalar()
            if existing is not None:
                log.debug("Updating message %s in room %s", message.id, room_pk)
                existing.timestamp = message.when
                existing.stanza = str(message.stanza)
                existing.author_jid = message.stanza.get_from()
                existing.source = source
                existing.legacy_id = legacy_msg_id
                session.add(existing)
                session.commit()
                return
            mam_msg = ArchivedMessage(
                stanza_id=message.id,
                timestamp=message.when,
                stanza=str(message.stanza),
                author_jid=message.stanza.get_from(),
                room_id=room_pk,
                source=source,
                legacy_id=legacy_msg_id,
            )
            session.add(mam_msg)
            session.commit()

    def get_messages(
        self,
        room_pk: int,
        start_date: Optional[datetime] = None,
        end_date: Optional[datetime] = None,
        before_id: Optional[str] = None,
        after_id: Optional[str] = None,
        ids: Collection[str] = (),
        last_page_n: Optional[int] = None,
        sender: Optional[str] = None,
        flip=False,
    ) -> Iterator[HistoryMessage]:

        with self.session() as session:
            q = select(ArchivedMessage).where(ArchivedMessage.room_id == room_pk)
            if start_date is not None:
                q = q.where(ArchivedMessage.timestamp >= start_date)
            if end_date is not None:
                q = q.where(ArchivedMessage.timestamp <= end_date)
            if before_id is not None:
                stamp = session.execute(
                    select(ArchivedMessage.timestamp).where(
                        ArchivedMessage.stanza_id == before_id
                    )
                ).scalar()
                if stamp is None:
                    raise XMPPError(
                        "item-not-found",
                        f"Message {before_id} not found",
                    )
                q = q.where(ArchivedMessage.timestamp < stamp)
            if after_id is not None:
                stamp = session.execute(
                    select(ArchivedMessage.timestamp).where(
                        ArchivedMessage.stanza_id == after_id
                    )
                ).scalar()
                if stamp is None:
                    raise XMPPError(
                        "item-not-found",
                        f"Message {after_id} not found",
                    )
                q = q.where(ArchivedMessage.timestamp > stamp)
            if ids:
                q = q.filter(ArchivedMessage.stanza_id.in_(ids))
            if sender is not None:
                q = q.where(ArchivedMessage.author_jid == sender)
            if flip:
                q = q.order_by(ArchivedMessage.timestamp.desc())
            else:
                q = q.order_by(ArchivedMessage.timestamp.asc())
            msgs = list(session.execute(q).scalars())
            if ids and len(msgs) != len(ids):
                raise XMPPError(
                    "item-not-found",
                    "One of the requested messages IDs could not be found "
                    "with the given constraints.",
                )
            if last_page_n is not None:
                if flip:
                    msgs = msgs[:last_page_n]
                else:
                    msgs = msgs[-last_page_n:]
            for h in msgs:
                yield HistoryMessage(
                    stanza=str(h.stanza), when=h.timestamp.replace(tzinfo=timezone.utc)
                )

    def get_first(self, room_pk: int, with_legacy_id=False) -> ArchivedMessage | None:
        with self.session() as session:
            q = (
                select(ArchivedMessage)
                .where(ArchivedMessage.room_id == room_pk)
                .order_by(ArchivedMessage.timestamp.asc())
            )
            if with_legacy_id:
                q = q.filter(ArchivedMessage.legacy_id.isnot(None))
            return session.execute(q).scalar()

    def get_last(
        self, room_pk: int, source: ArchivedMessageSource | None = None
    ) -> ArchivedMessage | None:
        with self.session() as session:
            q = select(ArchivedMessage).where(ArchivedMessage.room_id == room_pk)

            if source is not None:
                q = q.where(ArchivedMessage.source == source)

            return session.execute(
                q.order_by(ArchivedMessage.timestamp.desc())
            ).scalar()

    def get_first_and_last(self, room_pk: int) -> list[MamMetadata]:
        r = []
        with self.session():
            first = self.get_first(room_pk)
            if first is not None:
                r.append(MamMetadata(first.stanza_id, first.timestamp))
            last = self.get_last(room_pk)
            if last is not None:
                r.append(MamMetadata(last.stanza_id, last.timestamp))
        return r

    def get_most_recent_with_legacy_id(
        self, room_pk: int, source: ArchivedMessageSource | None = None
    ) -> ArchivedMessage | None:
        with self.session() as session:
            q = (
                select(ArchivedMessage)
                .where(ArchivedMessage.room_id == room_pk)
                .where(ArchivedMessage.legacy_id.isnot(None))
            )
            if source is not None:
                q = q.where(ArchivedMessage.source == source)
            return session.execute(
                q.order_by(ArchivedMessage.timestamp.desc())
            ).scalar()

    def get_least_recent_with_legacy_id_after(
        self, room_pk: int, after_id: str, source=ArchivedMessageSource.LIVE
    ) -> ArchivedMessage | None:
        with self.session() as session:
            after_timestamp = (
                session.query(ArchivedMessage.timestamp)
                .filter(ArchivedMessage.legacy_id == after_id)
                .scalar()
            )
            q = (
                select(ArchivedMessage)
                .where(ArchivedMessage.room_id == room_pk)
                .where(ArchivedMessage.legacy_id.isnot(None))
                .where(ArchivedMessage.source == source)
                .where(ArchivedMessage.timestamp > after_timestamp)
            )
            return session.execute(q.order_by(ArchivedMessage.timestamp.asc())).scalar()

    def get_by_legacy_id(self, room_pk: int, legacy_id: str) -> ArchivedMessage | None:
        with self.session() as session:
            return (
                session.query(ArchivedMessage)
                .filter(ArchivedMessage.room_id == room_pk)
                .filter(ArchivedMessage.legacy_id == legacy_id)
                .first()
            )


class MultiStore(EngineMixin):
    def get_xmpp_ids(self, user_pk: int, xmpp_id: str) -> list[str]:
        with self.session() as session:
            multi = session.execute(
                select(XmppIdsMulti)
                .where(XmppIdsMulti.xmpp_id == xmpp_id)
                .where(XmppIdsMulti.user_account_id == user_pk)
            ).scalar()
            if multi is None:
                return []
            return [m.xmpp_id for m in multi.legacy_ids_multi.xmpp_ids]

    def set_xmpp_ids(
        self, user_pk: int, legacy_msg_id: str, xmpp_ids: list[str], fail=False
    ) -> None:
        with self.session() as session:
            existing = session.execute(
                select(LegacyIdsMulti)
                .where(LegacyIdsMulti.user_account_id == user_pk)
                .where(LegacyIdsMulti.legacy_id == legacy_msg_id)
            ).scalar()
            if existing is not None:
                if fail:
                    raise
                log.warning("Resetting multi for %s", legacy_msg_id)
                session.execute(
                    delete(LegacyIdsMulti)
                    .where(LegacyIdsMulti.user_account_id == user_pk)
                    .where(LegacyIdsMulti.legacy_id == legacy_msg_id)
                )
                for i in xmpp_ids:
                    session.execute(
                        delete(XmppIdsMulti)
                        .where(XmppIdsMulti.user_account_id == user_pk)
                        .where(XmppIdsMulti.xmpp_id == i)
                    )
                session.commit()
                self.set_xmpp_ids(user_pk, legacy_msg_id, xmpp_ids, True)
                return

            row = LegacyIdsMulti(
                user_account_id=user_pk,
                legacy_id=legacy_msg_id,
                xmpp_ids=[
                    XmppIdsMulti(user_account_id=user_pk, xmpp_id=i)
                    for i in xmpp_ids
                    if i
                ],
            )
            session.add(row)
            session.commit()

    def get_legacy_id(self, user_pk: int, xmpp_id: str) -> Optional[str]:
        with self.session() as session:
            multi = session.execute(
                select(XmppIdsMulti)
                .where(XmppIdsMulti.xmpp_id == xmpp_id)
                .where(XmppIdsMulti.user_account_id == user_pk)
            ).scalar()
            if multi is None:
                return None
            return multi.legacy_ids_multi.legacy_id


class AttachmentStore(EngineMixin):
    def get_url(self, legacy_file_id: str) -> Optional[str]:
        with self.session() as session:
            return session.execute(
                select(Attachment.url).where(
                    Attachment.legacy_file_id == legacy_file_id
                )
            ).scalar()

    def set_url(self, user_pk: int, legacy_file_id: str, url: str) -> None:
        with self.session() as session:
            att = session.execute(
                select(Attachment)
                .where(Attachment.legacy_file_id == legacy_file_id)
                .where(Attachment.user_account_id == user_pk)
            ).scalar()
            if att is None:
                att = Attachment(
                    legacy_file_id=legacy_file_id, url=url, user_account_id=user_pk
                )
                session.add(att)
            else:
                att.url = url
            session.commit()

    def get_sims(self, url: str) -> Optional[str]:
        with self.session() as session:
            return session.execute(
                select(Attachment.sims).where(Attachment.url == url)
            ).scalar()

    def set_sims(self, url: str, sims: str) -> None:
        with self.session() as session:
            session.execute(
                update(Attachment).where(Attachment.url == url).values(sims=sims)
            )
            session.commit()

    def get_sfs(self, url: str) -> Optional[str]:
        with self.session() as session:
            return session.execute(
                select(Attachment.sfs).where(Attachment.url == url)
            ).scalar()

    def set_sfs(self, url: str, sfs: str) -> None:
        with self.session() as session:
            session.execute(
                update(Attachment).where(Attachment.url == url).values(sfs=sfs)
            )
            session.commit()

    def remove(self, legacy_file_id: str) -> None:
        with self.session() as session:
            session.execute(
                delete(Attachment).where(Attachment.legacy_file_id == legacy_file_id)
            )
            session.commit()


class RoomStore(UpdatedMixin):
    model = Room

    def __init__(self, *a, **kw):
        super().__init__(*a, **kw)
        with self.session() as session:
            session.execute(
                update(Room).values(
                    subject_setter=None, user_resources=None, history_filled=False
                )
            )
            session.commit()

    def set_avatar(self, room_pk: int, avatar_pk: int | None) -> None:
        with self.session() as session:
            session.execute(
                update(Room).where(Room.id == room_pk).values(avatar_id=avatar_pk)
            )
            session.commit()

    def get_avatar_legacy_id(self, room_pk: int) -> Optional[str]:
        with self.session() as session:
            room = session.execute(select(Room).where(Room.id == room_pk)).scalar()
            if room is None or room.avatar is None:
                return None
            return room.avatar.legacy_id

    def get_by_jid(self, user_pk: int, jid: JID) -> Optional[Room]:
        if jid.resource:
            raise TypeError
        with self.session() as session:
            return session.execute(
                select(Room)
                .where(Room.user_account_id == user_pk)
                .where(Room.jid == jid)
            ).scalar()

    def get_by_legacy_id(self, user_pk: int, legacy_id: str) -> Optional[Room]:
        with self.session() as session:
            return session.execute(
                select(Room)
                .where(Room.user_account_id == user_pk)
                .where(Room.legacy_id == legacy_id)
            ).scalar()

    def update_subject_setter(self, room_pk: int, subject_setter: str | None):
        with self.session() as session:
            session.execute(
                update(Room)
                .where(Room.id == room_pk)
                .values(subject_setter=subject_setter)
            )
            session.commit()

    def update(self, muc: "LegacyMUC") -> int:
        with self.session() as session:
            if muc.pk is None:
                row = Room(
                    jid=muc.jid,
                    legacy_id=str(muc.legacy_id),
                    user_account_id=muc.user_pk,
                )
            else:
                row = session.query(Room).filter(Room.id == muc.pk).one()

            row.updated = True
            row.extra_attributes = muc.serialize_extra_attributes()
            row.name = muc.name
            row.description = muc.description
            row.user_resources = (
                None
                if not muc._user_resources
                else json.dumps(list(muc._user_resources))
            )
            row.muc_type = muc.type
            row.subject = muc.subject
            row.subject_date = muc.subject_date
            row.subject_setter = muc.subject_setter
            row.participants_filled = muc._participants_filled
            row.n_participants = muc._n_participants
            row.user_nick = muc.user_nick
            session.add(row)
            session.commit()
            return row.id

    def update_subject_date(
        self, room_pk: int, subject_date: Optional[datetime]
    ) -> None:
        with self.session() as session:
            session.execute(
                update(Room).where(Room.id == room_pk).values(subject_date=subject_date)
            )
            session.commit()

    def update_subject(self, room_pk: int, subject: Optional[str]) -> None:
        with self.session() as session:
            session.execute(
                update(Room).where(Room.id == room_pk).values(subject=subject)
            )
            session.commit()

    def update_description(self, room_pk: int, desc: Optional[str]) -> None:
        with self.session() as session:
            session.execute(
                update(Room).where(Room.id == room_pk).values(description=desc)
            )
            session.commit()

    def update_name(self, room_pk: int, name: Optional[str]) -> None:
        with self.session() as session:
            session.execute(update(Room).where(Room.id == room_pk).values(name=name))
            session.commit()

    def update_n_participants(self, room_pk: int, n: Optional[int]) -> None:
        with self.session() as session:
            session.execute(
                update(Room).where(Room.id == room_pk).values(n_participants=n)
            )
            session.commit()

    def delete(self, room_pk: int) -> None:
        with self.session() as session:
            session.execute(delete(Room).where(Room.id == room_pk))
            session.execute(delete(Participant).where(Participant.room_id == room_pk))
            session.commit()

    def set_resource(self, room_pk: int, resources: set[str]) -> None:
        with self.session() as session:
            session.execute(
                update(Room)
                .where(Room.id == room_pk)
                .values(
                    user_resources=(
                        None if not resources else json.dumps(list(resources))
                    )
                )
            )
            session.commit()

    def nickname_is_available(self, room_pk: int, nickname: str) -> bool:
        with self.session() as session:
            return (
                session.execute(
                    select(Participant)
                    .where(Participant.room_id == room_pk)
                    .where(Participant.nickname == nickname)
                ).scalar()
                is None
            )

    def set_participants_filled(self, room_pk: int, val=True) -> None:
        with self.session() as session:
            session.execute(
                update(Room).where(Room.id == room_pk).values(participants_filled=val)
            )
            session.commit()

    def set_history_filled(self, room_pk: int, val=True) -> None:
        with self.session() as session:
            session.execute(
                update(Room).where(Room.id == room_pk).values(history_filled=True)
            )
            session.commit()

    def get_all(self, user_pk: int) -> Iterator[Room]:
        with self.session() as session:
            yield from session.execute(
                select(Room).where(Room.user_account_id == user_pk)
            ).scalars()


class ParticipantStore(EngineMixin):
    def __init__(self, *a, **kw):
        super().__init__(*a, **kw)
        with self.session() as session:
            session.execute(delete(Participant))
            session.execute(delete(Hat))
            session.commit()

    def add(self, room_pk: int, nickname: str) -> int:
        with self.session() as session:
            existing = session.execute(
                select(Participant.id)
                .where(Participant.room_id == room_pk)
                .where(Participant.nickname == nickname)
            ).scalar()
            if existing is not None:
                return existing
            participant = Participant(room_id=room_pk, nickname=nickname)
            session.add(participant)
            session.commit()
            return participant.id

    def get_by_nickname(self, room_pk: int, nickname: str) -> Optional[Participant]:
        with self.session() as session:
            return session.execute(
                select(Participant)
                .where(Participant.room_id == room_pk)
                .where(Participant.nickname == nickname)
            ).scalar()

    def get_by_resource(self, room_pk: int, resource: str) -> Optional[Participant]:
        with self.session() as session:
            return session.execute(
                select(Participant)
                .where(Participant.room_id == room_pk)
                .where(Participant.resource == resource)
            ).scalar()

    def get_by_contact(self, room_pk: int, contact_pk: int) -> Optional[Participant]:
        with self.session() as session:
            return session.execute(
                select(Participant)
                .where(Participant.room_id == room_pk)
                .where(Participant.contact_id == contact_pk)
            ).scalar()

    def get_all(self, room_pk: int, user_included=True) -> Iterator[Participant]:
        with self.session() as session:
            q = select(Participant).where(Participant.room_id == room_pk)
            if not user_included:
                q = q.where(~Participant.is_user)
            yield from session.execute(q).scalars()

    def get_for_contact(self, contact_pk: int) -> Iterator[Participant]:
        with self.session() as session:
            yield from session.execute(
                select(Participant).where(Participant.contact_id == contact_pk)
            ).scalars()

    def update(self, participant: "LegacyParticipant") -> None:
        with self.session() as session:
            session.execute(
                update(Participant)
                .where(Participant.id == participant.pk)
                .values(
                    resource=participant.jid.resource,
                    affiliation=participant.affiliation,
                    role=participant.role,
                    presence_sent=participant._presence_sent,  # type:ignore
                    # hats=[self.add_hat(h.uri, h.title) for h in participant._hats],
                    is_user=participant.is_user,
                    contact_id=(
                        None
                        if participant.contact is None
                        else participant.contact.contact_pk
                    ),
                )
            )
            session.commit()

    def add_hat(self, uri: str, title: str) -> Hat:
        with self.session() as session:
            existing = session.execute(
                select(Hat).where(Hat.uri == uri).where(Hat.title == title)
            ).scalar()
            if existing is not None:
                return existing
            hat = Hat(uri=uri, title=title)
            session.add(hat)
            session.commit()
            return hat

    def set_presence_sent(self, participant_pk: int) -> None:
        with self.session() as session:
            session.execute(
                update(Participant)
                .where(Participant.id == participant_pk)
                .values(presence_sent=True)
            )
            session.commit()

    def set_affiliation(self, participant_pk: int, affiliation: MucAffiliation) -> None:
        with self.session() as session:
            session.execute(
                update(Participant)
                .where(Participant.id == participant_pk)
                .values(affiliation=affiliation)
            )
            session.commit()

    def set_role(self, participant_pk: int, role: MucRole) -> None:
        with self.session() as session:
            session.execute(
                update(Participant)
                .where(Participant.id == participant_pk)
                .values(role=role)
            )
            session.commit()

    def set_hats(self, participant_pk: int, hats: list[HatTuple]) -> None:
        with self.session() as session:
            part = session.execute(
                select(Participant).where(Participant.id == participant_pk)
            ).scalar()
            if part is None:
                raise ValueError
            part.hats.clear()
            for h in hats:
                hat = self.add_hat(*h)
                if hat in part.hats:
                    continue
                part.hats.append(hat)
            session.commit()

    def delete(self, participant_pk: int) -> None:
        with self.session() as session:
            session.execute(delete(Participant).where(Participant.id == participant_pk))

    def get_count(self, room_pk: int) -> int:
        with self.session() as session:
            return session.query(
                count(Participant.id).filter(Participant.room_id == room_pk)
            ).scalar()


log = logging.getLogger(__name__)
_session: Optional[Session] = None
