#!/usr/bin/env python3

# Copyright (C) 2009-2026 Jérôme Poisson (goffi@goffi.org)

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from typing import cast
from twisted.internet import defer
from twisted.words.protocols.jabber import jid
from twisted.words.protocols.jabber.xmlstream import XMPPHandler
from twisted.words.xish import domish
from tx_xmpp import disco, itx_xmpp
from tx_xmpp import muc
from zope.interface import implementer

from libervia.backend.core import exceptions
from libervia.backend.core.constants import Const as C
from libervia.backend.core.core_types import SatXMPPEntity
from libervia.backend.core.i18n import D_, _
from libervia.backend.core.log import getLogger
from libervia.backend.plugins.plugin_xep_0045 import XEP_0045
from libervia.backend.plugins.plugin_xep_0166 import XEP_0166
from libervia.backend.plugins.plugin_xep_0167 import XEP_0167, mapping
from libervia.backend.plugins.plugin_xep_0402 import XEP_0402
from libervia.backend.plugins.plugin_xep_0482 import XEP_0482
from libervia.backend.tools import xml_tools
from libervia.backend.tools.common import data_format

log = getLogger(__name__)


PLUGIN_INFO = {
    C.PI_NAME: "Multiparty Jingle (Muji)",
    C.PI_IMPORT_NAME: "XEP-0272",
    C.PI_TYPE: "XEP",
    C.PI_PROTOCOLS: ["XEP-0272"],
    C.PI_DEPENDENCIES: ["XEP-0045", "XEP-0166", "XEP-0167", "XEP-0402", "XEP-0482"],
    C.PI_MAIN: "XEP_0272",
    C.PI_HANDLER: "yes",
    C.PI_DESCRIPTION: _(
        "Allow to run A/V conference with several participant using P2P connections. "
        "The number of participant must not be too big to make it work correctly."
    ),
}

NS_MUJI = "urn:xmpp:jingle:muji:0"
PRESENCE_MUJI = f'/presence/muji[@xmlns="{NS_MUJI}"]'


class XEP_0272:

    def __init__(self, host):
        log.info(f'Plugin "{PLUGIN_INFO[C.PI_NAME]}" initialization')
        self.host = host
        host.register_namespace("muji", NS_MUJI)
        self._muc = cast(XEP_0045, host.plugins["XEP-0045"])
        self._j = cast(XEP_0166, host.plugins["XEP-0166"])
        self._rtp = cast(XEP_0167, host.plugins["XEP-0167"])
        self._bookmarks = cast(XEP_0402, host.plugins["XEP-0402"])
        self._call_invites = cast(XEP_0482, host.plugins["XEP-0482"])
        host.trigger.add("XEP-0166_initiate_elt_built", self._on_jingle_elt_built)
        host.trigger.add("XEP-0166_on_session_initiate", self._on_session_initiate)
        host.trigger.add("XEP-0482_invite_received", self._on_call_invite_received)
        host.bridge.add_method(
            "call_group_start",
            ".plugin",
            in_sign="asss",
            out_sign="s",
            method=self._call_group_start,
            async_=True,
        )
        host.bridge.add_method(
            "call_group_data_set",
            ".plugin",
            in_sign="sss",
            out_sign="",
            method=self._call_group_data_set,
            async_=True,
        )
        host.bridge.add_method(
            "call_group_end",
            ".plugin",
            in_sign="ss",
            out_sign="",
            method=self._call_group_end,
            async_=True,
        )
        host.bridge.add_signal("call_group_setup", ".plugin", signature="sss")

    def get_handler(self, client):
        return XEP_0272_handler(self)

    def _call_group_start(
        self,
        entities_s: str,
        extra_s: str,
        profile_key: str,
    ) -> defer.Deferred[str]:
        client = self.host.get_client(profile_key)
        d = defer.ensureDeferred(
            self.call_group_start(
                client, [jid.JID(e) for e in entities_s], data_format.deserialise(extra_s)
            )
        )
        d.addCallback(data_format.serialise)
        return d

    def _call_group_data_set(
        self,
        room_jid_s: str,
        call_data_s: str,
        profile_key: str,
    ) -> None:
        client = self.host.get_client(profile_key)
        defer.ensureDeferred(
            self.call_group_data_set(
                client, jid.JID(room_jid_s), data_format.deserialise(call_data_s)
            )
        )

    def _call_group_end(self, room_jid_s: str, profile_key: str) -> None:
        client = self.host.get_client(profile_key)
        defer.ensureDeferred(self.call_group_end(client, jid.JID(room_jid_s)))

    async def call_group_end(self, client: SatXMPPEntity, room_jid: jid.JID) -> None:
        """End and clean call."""
        try:
            room = self._muc.get_room(client, room_jid)
        except exceptions.NotFound:
            log.warning(f"Ignoring MUJI end for unknown room: {room_jid}")
            return
        # We remove MUJI information from presence, as required by spec.
        presence_elt = domish.Element((None, "presence"))
        presence_elt["from"] = client.jid.full()
        presence_elt["to"] = room.roomJID.full()
        await client.presence.send(presence_elt)
        muji_data = self.get_muji_data(room)
        muji_data.clear()
        muji_data["to_call"] = set()
        # remove the room bookmark (if any, e.g. created by another client) to avoid
        # auto-joining the room at next login
        try:
            await self._bookmarks.remove(client, room_jid)
        except Exception:
            log.exception(f"Can't remove the bookmark of {room_jid}")

    def on_muji_request(
        self, presence_elt: domish.Element, client: SatXMPPEntity
    ) -> None:
        """Handle the MUJI element of a presence received in a MUC room"""
        from_jid = jid.JID(presence_elt["from"])
        room_jid = from_jid.userhostJID()
        try:
            room = self._muc.get_room(client, room_jid)
        except exceptions.NotFound:
            log.warning(
                f"Ignoring MUJI element from an unknown room: {presence_elt.toXml()}"
            )
            return
        own_jid = from_jid == self._muc.get_room_user_jid(client, room_jid)
        muji_data = self.get_muji_data(room)
        muji_elt = presence_elt.muji
        assert muji_elt is not None
        preparing_state = next(muji_elt.elements(NS_MUJI, "preparing"), None) is not None

        if preparing_state:
            if own_jid:
                # we have received the broadcast of our own preparation message
                muji_data["done_collecting"] = True
                self.try_to_finish_preparation(client, room, muji_data)
            elif not muji_data.get("done_collecting", False):
                # other entities currently doing preparation
                muji_data["preparing_jids"].add(from_jid)
        elif not own_jid:
            # a peer has advertised its data
            if not muji_data.get("done_preparing", False):
                muji_data["preparing_jids"].discard(from_jid)
                log.debug(f"[{client.profile}] received call data for {from_jid}.")
                # per spec, the joining participant initiates the sessions with all
                # the other participants; the sender's real JID is in the presence
                # itself (the room roster is not up to date yet for a joiner)
                user_jid = self._get_presence_user_jid(presence_elt)
                if user_jid is None:
                    raise exceptions.ExternalRequestError(
                        f"can't get the real JID of {from_jid}: the MUJI room is "
                        "anonymous"
                    )
                muji_data["to_call"].add(user_jid)
                self.try_to_finish_preparation(client, room, muji_data)

    def _get_presence_user_jid(self, presence_elt: domish.Element) -> jid.JID | None:
        """Get the sender's real JID from the muc#user item of a room presence."""
        x_elt = next(presence_elt.elements(muc.NS_MUC_USER, "x"), None)
        if x_elt is None:
            return None
        item_elt = next(x_elt.elements(muc.NS_MUC_USER, "item"), None)
        if item_elt is None or not item_elt.hasAttribute("jid"):
            return None
        return jid.JID(item_elt["jid"])

    def try_to_finish_preparation(
        self, client: SatXMPPEntity, room: muc.Room, muji_data: dict
    ) -> None:
        """Finish preparation if possible.

        This method checks if preparations of other JIDs needs to be waited, and if not,
        finishes our own preparation.
        """
        if not muji_data.get("done_collecting", False):
            # we have not received the broadcast of our own preparation message yet;
            # a joiner receives the other occupants' data presences (join processing)
            # before that echo, so waiting for it guarantees that all their data has
            # been collected
            return
        preparing_jids = muji_data.get("preparing_jids")
        if not preparing_jids:
            muji_data = self.get_muji_data(room)
            muji_data["done_preparing"] = True
            log.debug(f"[{client.profile}] Done preparing.")

            # ask the frontend to advertise our data and to initiate the sessions
            self.host.bridge.call_group_setup(
                room.roomJID.full(),
                data_format.serialise(
                    {"to_call": [entity.full() for entity in muji_data["to_call"]]}
                ),
                client.profile,
            )

    async def call_group_data_set(
        self,
        client: SatXMPPEntity,
        room_jid: jid.JID,
        call_data: dict,
    ) -> None:
        """Advertise the group call data on the MUC.

        Per spec, the late joiner initiates the sessions with the participants
        which joined before it.

        @param client: SatXMPPEntity instance.
        @param room_jid: JID of the room used for MUJI coordination.
        @param call_data: call data similar to the one used in ``XEP-0167.call_start``.
        """
        try:
            room = self._muc.get_room(client, room_jid)
        except exceptions.NotFound:
            log.warning(f"Ignoring MUJI data for unknown room: {room_jid}")
            return
        sdp_data = mapping.parse_sdp(call_data["sdp"], self._j.ROLE_INITIATOR)
        presence_elt, muji_elt = self.generate_presence_and_muji(client, room)
        for media_type, media_data in sdp_data.items():
            if media_type in ["audio", "video"]:
                application_data = media_data["application_data"]
                content_elt = muji_elt.addElement("content")
                # XXX: the initiator will be actually the last to join, but this attribute
                #   will be ignored anyway.
                content_elt["creator"] = self._j.ROLE_INITIATOR
                content_elt["name"] = media_data["id"]
                description_elt = mapping.build_description(
                    media_type, application_data, {}
                )
                content_elt.addChild(description_elt)

                # we only want to keep payload types
                to_remove = []
                for child_elt in description_elt.children:
                    if child_elt.name != "payload-type":
                        to_remove.append(child_elt)
                for elt in to_remove:
                    description_elt.children.remove(elt)

        await client.presence.send(presence_elt)

    async def _on_jingle_elt_built(
        self, client: SatXMPPEntity, session: dict, iq_elt, jingle_elt: domish.Element
    ) -> bool:
        """Add the <muji> element to the Jingle IQ of a MUJI call"""
        metadata = session.get("metadata")
        room_jid_s = metadata.get("muji_room") if metadata else None
        if room_jid_s is not None:
            muji_elt = jingle_elt.addElement((NS_MUJI, "muji"))
            muji_elt["room"] = room_jid_s
        return True

    async def _on_session_initiate(
        self,
        client: SatXMPPEntity,
        session: dict,
        request: domish.Element,
        jingle_elt: domish.Element,
    ) -> bool:
        """Store MUJI room from a session-initiate request, if any.

        Per spec, the room is only trusted if the sender is a member of it.
        """
        muji_elt = next(jingle_elt.elements(NS_MUJI, "muji"), None)
        if muji_elt is None:
            return True
        room_jid = jid.JID(muji_elt["room"])
        if not self._muc.is_jid_in_room(client, room_jid, session["peer_jid"]):
            log.warning(
                f"Ignoring MUJI room {room_jid}: {session['peer_jid']} is not a member"
            )
            return True
        session["muji_room"] = muji_elt["room"]
        return True

    def _on_call_invite_received(self, client, message_elt, invite_elt) -> bool:
        """Handle a XEP-0482 call invite for a MUJI conference"""
        muji_elt = next(invite_elt.elements(NS_MUJI, "muji"), None)
        if muji_elt is None:
            return True
        try:
            room_jid_s = muji_elt["room"]
        except KeyError:
            log.warning(f"Ignoring MUJI invite without room: {message_elt.toXml()}")
            return True
        from_jid_s = message_elt["from"]
        log.info(
            f"MUJI call invite received from {from_jid_s} for room {room_jid_s} "
            f"[{client.profile}]"
        )
        message_id = message_elt["id"]
        try:
            sub_type = muji_elt["type"]
        except KeyError:
            sub_type = C.META_SUBTYPE_CALL_AUDIO
        action_extra = {
            "type": C.META_TYPE_CONFIRM,
            "subtype": C.META_TYPE_GROUP_CALL,
            "from_jid": from_jid_s,
            "room_jid": room_jid_s,
            "sub_type": sub_type,
        }
        confirm_msg = D_(
            "You have been invited by {user} to join the room {room}. Do you accept?"
        ).format(user=from_jid_s, room=room_jid_s)
        confirm_title = _("Group call invitation")
        d = xml_tools.defer_confirm(
            self.host,
            confirm_msg,
            confirm_title,
            profile=client.profile,
            action_extra=action_extra,
        )

        def answer_cb(accepted: bool) -> None:
            # the accept contains the connection method used to join (XEP-0482)
            method_elt = domish.Element((NS_MUJI, "muji"))
            method_elt["room"] = room_jid_s
            if accepted:
                self._call_invites.accept(
                    client, jid.JID(from_jid_s), message_id, method_elt
                )
                # join the MUJI room and start preparing the call
                defer.ensureDeferred(self._join_and_prepare(client, jid.JID(room_jid_s)))
            else:
                self._call_invites.reject(client, jid.JID(from_jid_s), message_id)

        d.addCallback(answer_cb)
        return False

    async def _join_and_prepare(self, client: SatXMPPEntity, room_jid: jid.JID) -> None:
        """Join a MUJI room and start preparing (after accepting an invitation)"""
        room = await self._muc.join(client, room_jid, {})
        if room is not None:
            await self.start_preparation(client, room)

    def get_muji_data(self, room: muc.Room) -> dict:
        """Get MUJI related data for this room

        MUJI data is stored in the room object, so it will be deleted when the room object
        itself will be deleted.
        """
        try:
            return room._xep_0272_data
        except AttributeError:
            data = room._xep_0272_data = {"preparing_jids": set(), "to_call": set()}
            return data

    def generate_presence_and_muji(
        self, client: SatXMPPEntity, room: muc.Room
    ) -> tuple[domish.Element, domish.Element]:
        """Generate a <presence> stanza with MUJI element"""
        presence_elt = domish.Element((None, "presence"))
        presence_elt["from"] = client.jid.full()
        presence_elt["to"] = room.roomJID.full()
        muji_elt = presence_elt.addElement((NS_MUJI, "muji"))
        return presence_elt, muji_elt

    async def start_preparation(self, client: SatXMPPEntity, room: muc.Room) -> None:
        """Advertise the MUJI preparation state on the MUC room"""
        presence_elt, muji_elt = self.generate_presence_and_muji(client, room)
        muji_elt.addElement("preparing")
        await client.presence.send(presence_elt)

    async def call_group_start(
        self,
        client: SatXMPPEntity,
        entities: list[jid.JID],
        extra: dict,
    ) -> dict:
        """Initiate a group call with the given peers.

        A MUC room is created, made non-anonymous, and the peers are invited in it.

        @param client: The client instance.
        @param entities: JIDs of the peers to invite.
        @param extra: Extra data.

        @return: group call data, with the following keys:

            ``room_jid``
                JID of the MUC room where the MUJI coordination is done. It may also
                be used for normal chatting.

        @raise exceptions.ExternalRequestError: The MUC room can't be created, joined,
            or made non-anonymous.
        """
        log.debug(f"{client.profile} is starting a MUJI group call with {entities}")
        room_jid = self._muc.get_unique_name(client, prefix="_muji_")
        room = await self._muc.join(client, room_jid)
        log.info(f"[{client.profile}] MUJI room created at {room_jid}")
        if not room:
            raise exceptions.ExternalRequestError("Can't create or join group chat room.")
        # Make the room non-anonymous as participants need to know the real JID of others
        try:
            await self._muc.set_room_configuration(
                client, room_jid, {"muc#roomconfig_whois": "anyone"}
            )
        except Exception:
            raise exceptions.ExternalRequestError(
                f"Can't make the MUJI room {room_jid} non-anonymous: the MUJI mesh "
                "can't be established."
            )
        await self.start_preparation(client, room)
        sub_type = (
            C.META_SUBTYPE_CALL_VIDEO if extra.get("video") else C.META_SUBTYPE_CALL_AUDIO
        )
        for entity in entities:
            self.invite(client, entity, room_jid, sub_type)
        return {"room_jid": room_jid.full()}

    def invite(
        self, client: SatXMPPEntity, entity: jid.JID, room_jid: jid.JID, sub_type: str
    ) -> None:
        """Invite an entity to the MUJI conference (XEP-0482)"""
        muji_elt = domish.Element((NS_MUJI, "muji"))
        muji_elt["room"] = room_jid.full()
        muji_elt["type"] = sub_type
        self._call_invites.invite(client, entity, [muji_elt])


@implementer(itx_xmpp.IDisco)
class XEP_0272_handler(XMPPHandler):

    def __init__(self, plugin_parent):
        self.plugin_parent = plugin_parent

    def connectionInitialized(self):
        self.xmlstream.addObserver(
            PRESENCE_MUJI, self.plugin_parent.on_muji_request, client=self.parent
        )

    def getDiscoInfo(self, requestor, target, nodeIdentifier=""):
        return [disco.DiscoFeature(NS_MUJI)]

    def getDiscoItems(self, requestor, target, nodeIdentifier=""):
        return []