#!/usr/bin/env python3

# Libervia plugin for Jingle (XEP-0176)
# 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 Dict, List, Optional
import uuid

from twisted.internet import defer
from twisted.words.protocols.jabber.xmlstream import XMPPHandler
from twisted.words.xish import domish
from tx_xmpp import disco, itx_xmpp
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 _
from libervia.backend.core.log import getLogger
from libervia.backend.tools.common import data_format

from .plugin_xep_0166 import BaseTransportHandler

log = getLogger(__name__)

NS_JINGLE_ICE_UDP = "urn:xmpp:jingle:transports:ice-udp:1"

PLUGIN_INFO = {
    C.PI_NAME: "Jingle ICE-UDP Transport Method",
    C.PI_IMPORT_NAME: "XEP-0176",
    C.PI_TYPE: "XEP",
    C.PI_MODES: C.PLUG_MODE_BOTH,
    C.PI_PROTOCOLS: ["XEP-0176"],
    C.PI_DEPENDENCIES: ["XEP-0166"],
    C.PI_RECOMMENDATIONS: [],
    C.PI_MAIN: "XEP_0176",
    C.PI_HANDLER: "yes",
    C.PI_DESCRIPTION: _("""Implementation of Jingle ICE-UDP transport"""),
}


class XEP_0176(BaseTransportHandler):
    namespace = NS_JINGLE_ICE_UDP

    def __init__(self, host):
        log.info(f"plugin {PLUGIN_INFO[C.PI_NAME]!r} initialization")
        self.host = host
        self._j = host.plugins["XEP-0166"]  # shortcut to access jingle
        self._j.register_transport(
            NS_JINGLE_ICE_UDP, self._j.TRANSPORT_DATAGRAM, self, 100
        )
        host.bridge.add_method(
            "ice_candidates_add",
            ".plugin",
            in_sign="sss",
            out_sign="",
            method=self._ice_candidates_add,
            async_=True,
        )
        host.bridge.add_signal(
            "ice_candidates_new", ".plugin", signature="sss"
        )  # args: jingle_sid, candidates_serialised, profile
        host.bridge.add_signal(
            "ice_restart", ".plugin", signature="sss"
        )  # args: jingle_sid, side ("local" or "peer"), profile

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

    def _ice_candidates_add(
        self,
        session_id: str,
        media_ice_data_s: str,
        profile_key: str,
    ):
        client = self.host.get_client(profile_key)
        return defer.ensureDeferred(
            self.ice_candidates_add(
                client,
                session_id,
                data_format.deserialise(media_ice_data_s),
            )
        )

    def build_transport(
        self, ice_data: dict, add_trickle: bool = False
    ) -> domish.Element:
        """Generate <transport> element from ICE data

        @param ice_data: a dict containing the following keys:
            - "ufrag" (str): The ICE username fragment.
            - "pwd" (str): The ICE password.
            - "candidates" (List[dict]): A list of ICE candidate dictionaries, each
              containing:
                - "component_id" (int): The component ID.
                - "foundation" (str): The candidate foundation.
                - "address" (str): The candidate IP address.
                - "port" (int): The candidate port.
                - "priority" (int): The candidate priority.
                - "transport" (str): The candidate transport protocol, e.g., "udp".
                - "type" (str): The candidate type, e.g., "host", "srflx", "prflx", or
                  "relay".
                - "generation" (str, optional): The candidate generation. Defaults to "0".
                - "network" (str, optional): The candidate network. Defaults to "0".
                - "rel_addr" (str, optional): The related address for the candidate, if
                  any.
                - "rel_port" (int, optional): The related port for the candidate, if any.

        @return: A <transport> element.
        """
        try:
            ufrag: str = ice_data["ufrag"]
            pwd: str = ice_data["pwd"]
        except KeyError as e:
            raise exceptions.DataError(f"ICE {e} must be provided")
        candidates: List[dict] = ice_data.get("candidates", [])

        candidates.sort(key=lambda c: int(c.get("priority", 0)), reverse=True)
        transport_elt = domish.Element(
            (NS_JINGLE_ICE_UDP, "transport"), attribs={"ufrag": ufrag, "pwd": pwd}
        )

        for candidate in candidates:
            try:
                candidate_elt = transport_elt.addElement("candidate")
                candidate_elt["component"] = str(candidate["component_id"])
                candidate_elt["foundation"] = candidate["foundation"]
                candidate_elt["generation"] = str(candidate.get("generation", "0"))
                candidate_elt["id"] = candidate.get("id") or str(uuid.uuid4())
                candidate_elt["ip"] = candidate["address"]
                candidate_elt["network"] = str(candidate.get("network", "0"))
                candidate_elt["port"] = str(candidate["port"])
                candidate_elt["priority"] = str(candidate["priority"])
                candidate_elt["protocol"] = candidate["transport"]
                candidate_elt["type"] = candidate["type"]
            except KeyError as e:
                raise exceptions.DataError(
                    f"Mandatory ICE candidate attribute {e} is missing"
                )

            if "rel_addr" in candidate and "rel_port" in candidate:
                candidate_elt["rel-addr"] = candidate["rel_addr"]
                candidate_elt["rel-port"] = str(candidate["rel_port"])

        if add_trickle:
            transport_elt.addElement(
                (
                    "http://gultsch.de/xmpp/drafts/jingle/transports/ice-udp/option",
                    "trickle",
                )
            )

        self.host.trigger.point("XEP-0176_build_transport", transport_elt, ice_data)

        return transport_elt

    def parse_transport(self, transport_elt: domish.Element) -> dict:
        """Parse <transport> to a dict

        @param transport_elt: <transport> element
        @return: ICE data (as in [build_transport])
        """
        try:
            ice_data = {"ufrag": transport_elt["ufrag"], "pwd": transport_elt["pwd"]}
        except KeyError as e:
            raise exceptions.DataError(
                f"<transport> is missing mandatory attribute {e}: {transport_elt.toXml()}"
            )
        ice_data["candidates"] = ice_candidates = []

        for candidate_elt in transport_elt.elements(NS_JINGLE_ICE_UDP, "candidate"):
            try:
                candidate = {
                    "component_id": int(candidate_elt["component"]),
                    "foundation": candidate_elt["foundation"],
                    "address": candidate_elt["ip"],
                    "port": int(candidate_elt["port"]),
                    "priority": int(candidate_elt["priority"]),
                    "transport": candidate_elt["protocol"],
                    "type": candidate_elt["type"],
                }
            except KeyError as e:
                raise exceptions.DataError(
                    f"Mandatory attribute {e} is missing in candidate element"
                )

            if candidate_elt.hasAttribute("generation"):
                candidate["generation"] = candidate_elt["generation"]

            if candidate_elt.hasAttribute("network"):
                candidate["network"] = candidate_elt["network"]

            if candidate_elt.hasAttribute("rel-addr"):
                candidate["rel_addr"] = candidate_elt["rel-addr"]

            if candidate_elt.hasAttribute("rel-port"):
                candidate["rel_port"] = int(candidate_elt["rel-port"])

            ice_candidates.append(candidate)

        self.host.trigger.point("XEP-0176_parse_transport", transport_elt, ice_data)

        return ice_data

    async def jingle_session_init(
        self,
        client: SatXMPPEntity,
        session: dict,
        content_name: str,
    ) -> domish.Element:
        """Create a Jingle session initiation transport element with ICE candidates.

        @param client: SatXMPPEntity object representing the client.
        @param session: Dictionary containing session data.
        @param content_name: Name of the content.
        @param ufrag: ICE username fragment.
        @param pwd: ICE password.
        @param candidates: List of ICE candidate dictionaries parsed from the
            parse_ice_candidate method.

        @return: domish.Element representing the Jingle transport element.

        @raise exceptions.DataError: If mandatory data is missing from the candidates.
        """
        content_data = session["contents"][content_name]
        transport_data = content_data["transport_data"]
        ice_data = transport_data["local_ice_data"]
        return self.build_transport(ice_data, add_trickle=True)

    async def _send_ice_buffer(
        self,
        client: SatXMPPEntity,
        session: dict,
        content_name: str,
        content_data: dict,
    ) -> None:
        """Send ``transport-info`` actions with buffered ICE candidates.

        @param client: Libervia client instance.
        @param session: Jingle session.
        @param content_name: Name (ID) of the content.
        @param content_data: Data of the content.
        @param media_type: Type of media of the content.
        """
        media_type = content_data["application_data"].get("media")
        # we check if we have any buffered ICE candidates, and send them if it's the case
        try:
            buffer = session["XEP-0176_ice_buffer"]
            buffered_ice_data = buffer.pop(media_type)
        except KeyError:
            pass
        else:
            if not buffer:
                del session["XEP-0176_ice_buffer"]
            local_jid = session["local_jid"]
            peer_jid = session["peer_jid"]
            log.debug(f"Replaying ICE buffer from {local_jid} to {peer_jid}")
            transport_elt = self.build_transport(buffered_ice_data)
            iq_elt, __ = self._j.build_action(
                client,
                self._j.A_TRANSPORT_INFO,
                session,
                content_name,
                context_elt=transport_elt,
            )
            self.host.trigger.point(
                "XEP-0176_jingle_handler_send_buffer",
                client,
                session,
                content_name,
                content_data,
                transport_elt,
                iq_elt,
            )
            await iq_elt.send()

    async def jingle_handler(
        self,
        client: SatXMPPEntity,
        action: str,
        session: dict,
        content_name: str,
        transport_elt: domish.Element,
    ) -> domish.Element:
        """Handle Jingle requests

        @param client: The SatXMPPEntity instance.
        @param action: The action to be performed with the session.
        @param session: A dictionary containing the session information.
        @param content_name: The name of the content.
        @param transport_elt: The domish.Element instance representing the transport
            element.

        @return: <transport> element
        """
        content_data = session["contents"][content_name]
        transport_data = content_data["transport_data"]
        if action in (self._j.A_PREPARE_CONFIRMATION, self._j.A_PREPARE_INITIATOR):
            peer_ice_data = self.parse_transport(transport_elt)
            transport_data["peer_ice_data"] = peer_ice_data

        elif action == self._j.A_ACCEPTED_ACK:
            buffer = session.pop("XEP-0176_handler_buffer", None)
            if buffer:
                log.debug("replaying buffered events")
                for args in buffer:
                    await self.jingle_handler(*args)
            await self._send_ice_buffer(client, session, content_name, content_data)
        elif action == self._j.A_PREPARE_RESPONDER:
            pass

        elif action == self._j.A_SESSION_ACCEPT:
            await self._send_ice_buffer(client, session, content_name, content_data)

        elif action == self._j.A_START:
            pass

        elif action == self._j.A_SESSION_INITIATE:
            # responder side, we give our candidates
            transport_elt = self.build_transport(
                transport_data["local_ice_data"], add_trickle=True
            )
        elif action == self._j.A_TRANSPORT_INFO:

            if session["state"] == self._j.STATE_PENDING:
                # Session is not yet active; we buffer the arguments to replay them
                # when the session becomes active. This makes the frontend's life easier.
                log.debug("session is not active yet, buffering transport-info element")
                buffer = session.setdefault("XEP-0176_handler_buffer", [])
                buffer.append([client, action, session, content_name, transport_elt])
                return transport_elt

            media_type = content_data["application_data"].get("media")
            new_ice_data = self.parse_transport(transport_elt)
            restart = self.update_candidates(transport_data, new_ice_data, local=False)
            if restart:
                log.debug(
                    f"Peer ICE restart detected on session {session['id']} "
                    f"[{client.profile}]"
                )
                self.host.bridge.ice_restart(session["id"], "peer", client.profile)

            ice_candidates_new_cb = session.get("ice_candidates_new_cb")
            ice_candidates_data = {media_type: new_ice_data}
            if ice_candidates_new_cb is None:
                self.host.bridge.ice_candidates_new(
                    session["id"],
                    data_format.serialise(ice_candidates_data),
                    client.profile,
                )
            else:
                ice_candidates_new_cb(client, session, ice_candidates_data)
        elif action == self._j.A_DESTROY:
            pass
        else:
            log.warning("FIXME: unmanaged action {}".format(action))

        return transport_elt

    def jingle_terminate(
        self,
        client: SatXMPPEntity,
        action: str,
        session: dict,
        content_name: str,
        reason_elt: domish.Element,
    ) -> None:
        log.debug("ICE-UDP session terminated")

    def update_candidates(
        self, transport_data: dict, new_ice_data: dict, local: bool
    ) -> bool:
        """Update ICE candidates when new one are received

        @param transport_data: transport_data of the content linked to the candidates
        @param new_ice_data: new ICE data, in the same format as returned
            by [self.parse_transport]
        @param local: True if it's our candidates, False if it's peer ones
        @return: True if there is a ICE restart
        """
        key = "local_ice_data" if local else "peer_ice_data"
        try:
            ice_data = transport_data[key]
        except KeyError:
            log.warning(f"no {key} available")
            transport_data[key] = new_ice_data
        else:
            if (
                new_ice_data["ufrag"] != ice_data["ufrag"]
                or new_ice_data["pwd"] != ice_data["pwd"]
            ):
                ice_data["ufrag"] = new_ice_data["ufrag"]
                ice_data["pwd"] = new_ice_data["pwd"]
                ice_data["candidates"] = new_ice_data["candidates"]
                return True
        return False

    async def ice_candidates_add(
        self, client: SatXMPPEntity, session_id: str, media_ice_data: dict[str, dict]
    ) -> None:
        """Called when a new ICE candidates are available for a session

        @param session_id: Session ID
        @param media_ice_data: a map from media type (audio, video) to ICE data
            ICE data must be in the same format as in [self.parse_transport]
        """
        session = self._j.get_session(client, session_id)
        iq_elt: Optional[domish.Element] = None
        content_name: str | None = None
        content_data: dict | None = None

        for media_type, new_ice_data in media_ice_data.items():
            if session["state"] == self._j.STATE_PENDING:
                log.debug(f"session not active, buffering")
                buffer = session.setdefault("XEP-0176_ice_buffer", {})
                media_buffer = buffer.setdefault(media_type, {})

                for key in ["ufrag", "pwd"]:
                    if key not in media_buffer:
                        media_buffer[key] = new_ice_data[key]
                    else:
                        if media_buffer[key] != new_ice_data[key]:
                            log.warning(
                                f"{key} conflict, new value will replace old one\n"
                                f"buffer={media_buffer[key]!r}\n"
                                f"new={new_ice_data[key]!r}"
                            )
                            media_buffer[key] = new_ice_data[key]

                media_buffer.setdefault("candidates", []).extend(
                    new_ice_data["candidates"]
                )
                continue

            for content_name, content_data in session["contents"].items():
                if content_data["application_data"].get("media") == media_type:
                    break
            else:
                log.warning(f"no media of type {media_type} has been found")
                continue
            restart = self.update_candidates(
                content_data["transport_data"], new_ice_data, True
            )
            if restart:
                log.debug(
                    f"Local ICE restart detected on session {session['id']} "
                    f"[{client.profile}]"
                )
                self.host.bridge.ice_restart(session["id"], "local", client.profile)
            transport_elt = self.build_transport(new_ice_data)
            iq_elt, __ = self._j.build_action(
                client,
                self._j.A_TRANSPORT_INFO,
                session,
                content_name,
                iq_elt=iq_elt,
                context_elt=transport_elt,
            )

        if iq_elt is not None:
            assert content_name is not None and content_data is not None
            try:
                self.host.trigger.point(
                    "XEP-0176_ice_candidate_send",
                    client,
                    session,
                    media_ice_data,
                    content_name,
                    content_data,
                    iq_elt,
                )
                await iq_elt.send()
            except Exception as e:
                log.warning(f"Could not send new ICE candidates: {e}")


@implementer(itx_xmpp.IDisco)
class XEP_0176_handler(XMPPHandler):
    def getDiscoInfo(self, requestor, target, nodeIdentifier=""):
        return [disco.DiscoFeature(NS_JINGLE_ICE_UDP)]

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