#!/usr/bin/env python3

# Libervia plugin for Namespace Delegation (XEP-0355).
# 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 dataclasses import dataclass, field
from functools import partial
from typing import TYPE_CHECKING, NamedTuple, cast

from twisted.internet import defer, reactor
from twisted.internet.interfaces import IDelayedCall
from twisted.words.protocols.jabber import error, jid
from twisted.words.protocols.jabber.xmlstream import XMPPHandler
from twisted.words.xish import domish
from tx_xmpp import disco, itx_xmpp, pubsub
from tx_xmpp import data_form
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 SatXMPPComponent, SatXMPPEntity
from libervia.backend.core.i18n import _
from libervia.backend.core.log import getLogger
from libervia.backend.plugins.plugin_xep_0297 import NS_FORWARD, XEP_0297
from libervia.backend.tools.utils import as_deferred, ensure_deferred
from libervia.backend.tools.xml_tools import element_copy

if TYPE_CHECKING:
    from libervia.backend.core.main import LiberviaBackend


log = getLogger(__name__)

PLUGIN_INFO = {
    C.PI_NAME: "Namespace Delegation",
    C.PI_IMPORT_NAME: "XEP-0355",
    C.PI_TYPE: "XEP",
    C.PI_MODES: [C.PLUG_MODE_COMPONENT],
    C.PI_PROTOCOLS: ["XEP-0355"],
    C.PI_DEPENDENCIES: ["XEP-0297"],
    C.PI_RECOMMENDATIONS: [],
    C.PI_MAIN: "XEP_0355",
    C.PI_HANDLER: "yes",
    C.PI_DESCRIPTION: _("""Namespace delegation implementation."""),
}

NS_DELEGATION = "urn:xmpp:delegation:2"

DELEGATION_MAIN_SEP = "::"
DELEGATION_BARE_SEP = ":bare:"

SEND_HOOK_TIMEOUT = 300


@dataclass
class DelegationData:
    delegation_received: defer.Deferred[None] = field(
        default_factory=lambda: defer.Deferred()
    )


class DelegatedIQ(NamedTuple):
    """Data of IQ request being tracked for delegation."""

    #: The client that received the delegation.
    client: SatXMPPComponent
    #: The original IQ request.
    iq_request_elt: domish.Element
    #: The wrapping delegation IQ from the server.
    wrapping_iq_elt: domish.Element
    #: Timer for the request.
    timeout: IDelayedCall


class XEP_0355:
    """XEP-0355: Namespace Delegation handler."""

    namespace = NS_DELEGATION

    def __init__(self, host: "LiberviaBackend") -> None:
        log.info(f"plugin {PLUGIN_INFO[C.PI_NAME]!r} initialization")
        self.host = host
        self._forward = cast(XEP_0297, host.plugins["XEP-0297"])
        self.delegated_namespaces: set[str] = set()
        self._current_iqs: dict[str, DelegatedIQ] = {}

        host.register_namespace("delegation", NS_DELEGATION)
        host.trigger.add("stream_hooks", self.add_hooks)

    def add_hooks(self, client, receive_hooks, send_hooks):
        """Register send hook to intercept outgoing IQ results/errors."""
        send_hooks.append(partial(self.on_send, client=client))
        return True

    def on_send(self, obj, client):
        """Intercept outgoing stanzas to catch delegated IQ responses."""
        if not client.is_component:
            return

        if not self._current_iqs:
            return

        if not isinstance(obj, domish.Element):
            return

        if obj.name != "iq":
            return

        if obj.getAttribute("type") not in ("result", "error"):
            return

        self.handle_outgoing_iq(obj)

    def handle_outgoing_iq(self, iq_elt: domish.Element) -> None:
        """Check if outgoing IQ is a response to a delegated stanza.

        @param iq_elt: The outgoing IQ to check.
        """
        iq_id = iq_elt.getAttribute("id")
        if not iq_id or iq_id not in self._current_iqs:
            return

        delegated_iq = self._current_iqs[iq_id]

        iq_to = iq_elt.getAttribute("to")
        iq_request_from = delegated_iq.iq_request_elt.getAttribute("from")

        # Response's "to" must match original request's "from"
        if iq_to != iq_request_from:
            log.error(
                f"IQ id conflict for {iq_id}: expected response to {iq_request_from!r}, "
                f"got {iq_to!r}"
            )
            del self._current_iqs[iq_id]
            return

        if not delegated_iq.timeout.called:
            delegated_iq.timeout.cancel()

        del self._current_iqs[iq_id]
        self.delegated_result(
            delegated_iq.client,
            delegated_iq.iq_request_elt,
            iq_elt,
            delegated_iq.wrapping_iq_elt,
        )
        raise exceptions.CancelError()

    def is_pep(self, client: SatXMPPComponent, entity_jid: jid.JID) -> bool:
        """Check if a request is a PEP request.

        @param client: Client instance.
        @param to_jid: The JID to check.
        @return: True if this is a PEP request to a local user.
        @raise exceptions.InternalError: Client is not a component (XEP-0355 can be used
            only with components).
        """
        if not client.is_component:
            raise exceptions.InternalError("XEP-0355 used with a non-component.")

        if entity_jid == client.jid:
            return False

        if entity_jid.user is None:
            return False

        if entity_jid.resource is not None:
            return False

        if not client.is_local(entity_jid):
            return False

        return True

    def get_handler(self, client: SatXMPPEntity) -> "XEP_0355_Handler":
        return XEP_0355_Handler(self)

    def on_advertise(self, message_elt: domish.Element, client: SatXMPPComponent) -> None:
        """Handle delegation advertisement from server.

        @param client: Client instance.
        @param message_elt: Delegation advertisement message.
        """
        server_jid = client.server_jid
        if jid.JID(message_elt["from"]) != server_jid:
            log.warning(
                f"Delegation advertising message received from {message_elt['from']},"
                f" while expected server JID is {server_jid}, this may be a security"
                " threat, please check your configuration and network."
            )
            raise error.StanzaError("not-allowed")

        try:
            delegation_elt = next(message_elt.elements(NS_DELEGATION, "delegation"))
        except StopIteration:
            log.warning("Invalid delegation advertisement: missing delegation element")
            raise error.StanzaError("not-acceptable")

        delegated = {}
        for child in delegation_elt.elements(NS_DELEGATION):
            if child.name != "delegated":
                log.warning(
                    f"Invalid delegation advertisement: unexpected element {child.name!r}"
                )
                raise error.StanzaError("not-acceptable")

            try:
                namespace = child["namespace"]
            except KeyError:
                log.warning(
                    "Invalid delegation advertisement: missing namespace attribute in"
                    " delegated element"
                )
                raise error.StanzaError("not-acceptable")

            delegated[namespace] = []
            for attribute_elt in child.elements(NS_DELEGATION, "attribute"):
                try:
                    delegated[namespace].append(attribute_elt["name"])
                except KeyError:
                    log.warning(
                        "Invalid delegation advertisement: missing name attribute in"
                        " attribute element"
                    )
                    raise error.StanzaError("not-acceptable")

        self.delegated_namespaces = set(delegated.keys())

        log_lines = [f"Delegations updated for {client.profile}:"]
        for ns, attrs in sorted(delegated.items()):
            if attrs:
                log_lines.append(
                    f"  - namespace {ns!r} with filtering on {len(attrs)} attribute(s):"
                    f" {', '.join(map(repr, attrs))}"
                )
            else:
                log_lines.append(f"  - namespace {ns!r}")

        log.info("\n".join(log_lines))

        delegation_data = client.get_plugin_data(C.PI_IMPORT_NAME, DelegationData)
        if not delegation_data.delegation_received.called:
            delegation_data.delegation_received.callback(None)

    def on_forward(
        self, delegation_iq_elt: domish.Element, client: SatXMPPComponent
    ) -> None:
        """Handle forwarded stanza from server.

        @param client: Client instance.
        @param delegation_iq_elt: Full delegation stanza.
        """
        server_jid = client.server_jid
        if jid.JID(delegation_iq_elt["from"]) != server_jid:
            log.warning(
                "SECURITY: forwarded stanza from unexpected server:"
                f" {delegation_iq_elt.toXml()}"
            )
            raise error.StanzaError("not-allowed")

        try:
            delegation_elt = next(delegation_iq_elt.elements(NS_DELEGATION, "delegation"))
            forwarded_elt = next(delegation_elt.elements(C.NS_FORWARD, "forwarded"))
            forwarded_iq_elt = forwarded_elt.firstChildElement()
            if forwarded_iq_elt is None:
                raise error.StanzaError("not-acceptable")
        except StopIteration:
            raise error.StanzaError("not-acceptable")

        iq_id = forwarded_iq_elt["id"]
        iq_to = forwarded_iq_elt.getAttribute("to")
        timeout = reactor.callLater(SEND_HOOK_TIMEOUT, self.on_timeout, iq_id, iq_to)
        self._current_iqs[iq_id] = DelegatedIQ(
            client, forwarded_iq_elt, delegation_iq_elt, timeout
        )

        if not forwarded_iq_elt.hasAttribute("to"):
            forwarded_iq_elt["to"] = jid.JID(forwarded_iq_elt["from"]).userhost()

        # This attribute may be used by plugins to track that the stanza is delegated.
        forwarded_iq_elt.delegated = True

        # Prevent further processing.
        delegation_iq_elt.handled = True

        client.xmlstream.dispatch(forwarded_iq_elt)

    def delegated_result(
        self,
        client: SatXMPPComponent,
        iq_req_elt: domish.Element,
        iq_resp_elt: domish.Element,
        wrapping_iq_elt: domish.Element,
    ) -> None:
        """Wrap result/error stanza and send back to server.

        @param client: Client instance.
        @param iq_req_elt: The original forwarded IQ request.
        @param iq_resp_elt: The response IQ (result or error).
        @param wrapping_iq_elt: The original delegation IQ from the server.
        """
        iq_result_elt = domish.Element((None, "iq"))
        iq_result_elt["from"] = wrapping_iq_elt["to"]
        iq_result_elt["to"] = wrapping_iq_elt["from"]
        iq_result_elt["id"] = wrapping_iq_elt["id"]
        iq_result_elt["type"] = "result"

        delegation_elt = domish.Element((NS_DELEGATION, "delegation"))

        wrapped_iq_elt = domish.Element(("jabber:client", "iq"))
        # According to the specs, we don't set a "from" here
        wrapped_iq_elt["to"] = iq_req_elt["from"]
        wrapped_iq_elt["id"] = iq_req_elt["id"]
        wrapped_iq_elt["type"] = iq_resp_elt["type"]

        for child in iq_resp_elt.children:
            wrapped_iq_elt.addChild(element_copy(child))

        forwarded_elt = XEP_0297.build_forwarded(wrapped_iq_elt)
        delegation_elt.addChild(forwarded_elt)
        iq_result_elt.addChild(delegation_elt)

        client.xmlstream.send(iq_result_elt)

    def on_timeout(self, iq_id: str, iq_to: str | None) -> None:
        """Handle timeout for a delegated IQ request.

        @param iq_id: The ID of the timed-out IQ.
        @param iq_to: The 'to' attribute of the timed-out IQ.
        """
        self._current_iqs.pop(iq_id, None)
        log.error(f"Timeout waiting for response to delegated IQ {iq_id!r} to {iq_to!r}")


@implementer(itx_xmpp.IDisco)
class XEP_0355_Handler(XMPPHandler):
    """Handler for delegation stanzas on a per-client basis."""

    def __init__(self, plugin_parent: XEP_0355) -> None:
        self.plugin_parent = plugin_parent

    @property
    def client(self) -> SatXMPPComponent:
        assert self.parent is not None
        return self.parent

    def connectionInitialized(self) -> None:
        """Handle connection initialization for delegation stanzas."""
        self.client.xmlstream.addObserver(
            f"/message/delegation[@xmlns='{NS_DELEGATION}']",
            self.plugin_parent.on_advertise,
            client=self.client,
        )
        self.client.xmlstream.addObserver(
            f"/iq[@type='set']/delegation[@xmlns='{NS_DELEGATION}']"
            f"/forwarded[@xmlns='{NS_FORWARD}']",
            self.plugin_parent.on_forward,
            client=self.client,
        )

    @ensure_deferred
    async def getDiscoInfo(
        self,
        requestor: jid.JID,
        target: jid.JID,
        nodeIdentifier: str = "",
    ) -> list[disco.DiscoFeature | disco.DiscoIdentity | data_form.Form]:
        """Get disco info for delegation nesting.

        Server queries this handler for features/identities to nest in its own disco info.
        Node format: urn:xmpp:delegation:2::<namespace> for server
                     urn:xmpp:delegation:2:bare:<namespace> for bare JID

        @param requestor: JID of the requesting entity
        @param target: JID of the target entity
        @param nodeIdentifier: Node identifier
        @return: Disco features/identities for the delegated namespace
        """
        if not nodeIdentifier.startswith(NS_DELEGATION):
            return []

        # Extract namespace from delegation node identifier
        namespace: str | None = None
        if DELEGATION_MAIN_SEP in nodeIdentifier:
            _, namespace = nodeIdentifier.split(DELEGATION_MAIN_SEP, 1)
        elif DELEGATION_BARE_SEP in nodeIdentifier:
            _, namespace = nodeIdentifier.split(DELEGATION_BARE_SEP, 1)
        else:
            log.error(f"Unexpected disco node: {nodeIdentifier}")

        if not namespace:
            return []

        return await self._get_disco_info_for_namespace(namespace, requestor, target)

    def getDiscoItems(
        self,
        requestor: jid.JID,
        target: jid.JID,
        nodeIdentifier: str = "",
    ) -> list[disco.DiscoItem]:
        """Get disco items.

        Actual disco#items delegation is handled via stanza forwarding.

        @param requestor: JID of the requesting entity
        @param target: JID of the target entity
        @param nodeIdentifier: Node identifier
        @return: list of disco items
        """
        return []

    async def _get_disco_info_for_namespace(
        self,
        namespace: str,
        requestor: jid.JID,
        target: jid.JID,
    ) -> list[disco.DiscoFeature | disco.DiscoIdentity | data_form.Form]:
        """Collect disco info for a delegated namespace from all handlers.

        @param namespace: The delegated namespace
        @return: List of features/identities matching the namespace
        """
        # Normalize pubsub namespace: pubsub uses several namespaces starting with
        # NS_PUBSUB; e.g. http://jabber.org/protocol/pubsub#owner. We return the same
        # disco info for all pubsub variants (matches SàT Pubsub reference implementation)
        normalized_namespace = namespace
        if namespace.startswith(pubsub.NS_PUBSUB):
            normalized_namespace = pubsub.NS_PUBSUB

        disco_info: list[disco.DiscoFeature | disco.DiscoIdentity | data_form.Form] = []

        for handler in self.client.handlers:
            if not hasattr(handler, "getDiscoInfo"):
                continue

            try:
                handler_disco_info = await as_deferred(
                    handler.getDiscoInfo, requestor, target, ""
                )
                if not handler_disco_info:
                    continue

                for info in handler_disco_info:
                    if isinstance(info, disco.DiscoFeature):
                        if info.startswith(normalized_namespace):
                            disco_info.append(info)
                    elif isinstance(info, disco.DiscoIdentity):
                        # Check if identity matches namespace (pubsub category for pubsub
                        # namespace)
                        # FIXME: we should handle other identities
                        if (
                            normalized_namespace == pubsub.NS_PUBSUB
                            and info.category == "pubsub"
                        ):
                            disco_info.append(info)
                    elif isinstance(info, data_form.Form):
                        if info.formNamespace and info.formNamespace.startswith(
                            normalized_namespace
                        ):
                            disco_info.append(info)
            except Exception:
                log.exception(f"Error collecting disco info from handler")

        # Add PEP identity for pubsub namespace
        # FIXME: to be moved to Pubsub component
        if normalized_namespace == pubsub.NS_PUBSUB:
            disco_info.append(disco.DiscoIdentity("pubsub", "pep", "PEP"))

        return disco_info