#!/usr/bin/env python3


# Libervia plugin for managing xep-0115
# 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 base64 import b64encode
from dataclasses import dataclass, field
from hashlib import sha1
from typing import TYPE_CHECKING, Literal

from twisted.internet import defer, error
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 data_form, disco, itx_xmpp
from zope.interface import implementer

from libervia.backend import G
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

log = getLogger(__name__)


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

PRESENCE = "/presence"
NS_ENTITY_CAPABILITY = "http://jabber.org/protocol/caps"
NS_CAPS_OPTIMIZE = "http://jabber.org/protocol/caps#optimize"
CAPABILITY_UPDATE = PRESENCE + '/c[@xmlns="' + NS_ENTITY_CAPABILITY + '"]'

IMPORT_NAME = "XEP-0115"

PLUGIN_INFO = {
    C.PI_NAME: "XEP 0115 Plugin",
    C.PI_IMPORT_NAME: IMPORT_NAME,
    C.PI_TYPE: "XEP",
    C.PI_MODES: C.PLUG_MODE_BOTH,
    C.PI_PROTOCOLS: ["XEP-0115"],
    C.PI_DEPENDENCIES: [],
    C.PI_MAIN: "XEP_0115",
    C.PI_HANDLER: "yes",
    C.PI_DESCRIPTION: _("""Implementation of entity capabilities"""),
}


@dataclass
class CapsClientData:
    """Per-client data for XEP-0115 capabilities."""

    #: True if caps optimisation is available
    caps_optimize: bool = False
    #: True if a presence with caps has been sent
    caps_sent: bool = False
    #: Our own capability hash
    caps_hash: str | None = None
    #: The <c/> element built for this client's presence
    caps_elt: domish.Element | None = field(default=None)


class ByteIdentity:
    """This class manage identity as bytes (needed for i;octet sort)

    It is used for the hash generation.
    """

    def __init__(self, identity: disco.DiscoIdentity):
        assert isinstance(identity, disco.DiscoIdentity)
        self.category = identity.category.encode("utf-8")
        self.idType = identity.type.encode("utf-8")
        self.name = identity.name.encode("utf-8") if identity.name else b""
        self.lang = identity.lang.encode("utf-8") if identity.lang else b""

    def __bytes__(self):
        return b"%s/%s/%s/%s" % (self.category, self.idType, self.lang, self.name)


class XEP_0115:
    def __init__(self, host: "LiberviaBackend") -> None:
        self.host = host
        log.info(_("Plugin XEP_0115 initialization"))
        host.trigger.add("Presence send", self._presence_trigger)
        host.trigger.add("disco_info", self._disco_info_trigger)

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

    async def _prepare_caps(self, client: SatXMPPEntity) -> None:
        """Generate our own capability hash for a client.

        Hash is generated per client instance as disco infos/identities may change
        (plugins loaded/unloaded, config changes, etc.).

        @param client: the client entity to generate the hash for.
        @raises exceptions.InternalError: The client's own disco info is malformed.
        """
        data = client.get_plugin_data(IMPORT_NAME, CapsClientData)

        # optimize check
        data.caps_optimize = await G.host.has_feature(client, NS_CAPS_OPTIMIZE)
        if data.caps_optimize:
            log.info(_("Caps optimisation enabled"))
            data.caps_sent = False
        else:
            log.warning(_("Caps optimisation not available"))

        old_hash = data.caps_hash
        _infos = await client.discoHandler.info(client.jid, client.jid, "")
        disco_infos = disco.DiscoInfo()
        for item in _infos:
            disco_infos.append(item)

        # Validate our own disco info before hashing.
        if not self._validate_caps(client.jid, "", "", disco_infos):
            raise exceptions.InternalError(
                "Our own disco#info response is malformed. Caps cannot be generated."
            )

        cap_hash = self.generate_hash(disco_infos)
        if cap_hash != old_hash:
            log.info(f"Capability hash changed: [{old_hash}] -> [{cap_hash}]")
        data.caps_hash = cap_hash
        log.info(f"Our capability hash has been generated: [{cap_hash}]")
        log.debug("Generating capability domish.Element")
        c_elt = domish.Element((NS_ENTITY_CAPABILITY, "c"))
        c_elt["hash"] = "sha-1"
        c_elt["node"] = C.APP_URL
        c_elt["ver"] = cap_hash
        data.caps_elt = c_elt
        if data.caps_optimize:
            data.caps_sent = False
        if cap_hash not in G.host.memory.disco.hashes:
            G.host.memory.disco.hashes[cap_hash] = disco_infos
            G.host.memory.update_entity_data(
                client, client.jid, C.ENTITY_CAP_HASH, cap_hash
            )

    def _disco_info_trigger(
        self,
        requestor: jid.JID | None,
        target: jid.JID | None,
        nodeIdentifier: str,
    ) -> tuple[Literal[True], None] | tuple[Literal[False], defer.Deferred[list]]:
        """Callback for the ``disco_info`` trigger.

        Intercepts caps node queries to return the cached disco info directly.

        @param requestor: entity that sent the disco request
        @param target: entity the disco request was sent to
        @param nodeIdentifier: optional node being queried
        @return: ``(False, deferred)`` with cached disco items for known caps queries, or
            ``(True, None)`` to let normal flow continue
        """
        # Check if this is a caps node query (format: node#ver)
        if nodeIdentifier and "#" in nodeIdentifier:
            parts = nodeIdentifier.rsplit("#", 1)
            if len(parts) == 2:
                caps_node, cap_ver = parts
                # Only respond if this matches our caps node
                if caps_node == C.APP_URL and cap_ver in G.host.memory.disco.hashes:
                    disco_infos = G.host.memory.disco.hashes[cap_ver]
                    # Return all items from the cached disco info
                    return False, defer.succeed(list(disco_infos))
        return True, None

    def _presence_add_elt(self, client: SatXMPPEntity, obj: domish.Element) -> None:
        data = client.get_plugin_data(IMPORT_NAME, CapsClientData)
        if data.caps_optimize:
            if data.caps_sent:
                return
            data.caps_sent = True
        if data.caps_elt is not None:
            obj.addChild(data.caps_elt)

    def _presence_trigger(
        self,
        client: SatXMPPEntity,
        obj: domish.Element,
        presence_d: defer.Deferred[domish.Element],
    ) -> bool:
        data = client.get_plugin_data(IMPORT_NAME, CapsClientData)
        if data.caps_elt is None:
            presence_d.addCallback(
                lambda __: defer.ensureDeferred(self._prepare_caps(client))
            )

        presence_d.addCallback(lambda __: self._presence_add_elt(client, obj))
        presence_d.addErrback(G.host.log_errback)
        return True

    def _update(self, presence_elt: domish.Element, client: SatXMPPEntity) -> None:
        defer.ensureDeferred(self.update(client, presence_elt))

    async def update(self, client: SatXMPPEntity, presence_elt: domish.Element) -> None:
        """
        Manage the capabilities of the entity

        Check if we know the version of this capabilities and get the capabilities if
        necessary

        @param presence_elt: <presence/> element to parse.
        """
        from_jid = jid.JID(presence_elt["from"])
        c_elem = next(presence_elt.elements(NS_ENTITY_CAPABILITY, "c"), None)
        if c_elem is None:
            # Presence without caps, nothing to do
            return
        try:
            c_ver = c_elem["ver"]
            c_hash = c_elem["hash"]
            c_node = c_elem["node"]
        except KeyError:
            log.warning(_("Received invalid capabilities tag: %s") % c_elem.toXml())
            return

        if c_ver in G.host.memory.disco.hashes:
            # we already know the hash, we update the jid entity
            log.debug(
                "hash [%(hash)s] already in cache, updating entity [%(jid)s]"
                % {"hash": c_ver, "jid": from_jid.full()}
            )
            G.host.memory.update_entity_data(client, from_jid, C.ENTITY_CAP_HASH, c_ver)
            return

        if c_hash != "sha-1":  # unknown hash method
            log.warning(
                _(
                    "Unknown hash method for entity capabilities: [{hash_method}] "
                    "(entity: {entity_jid}, node: {node})"
                ).format(hash_method=c_hash, entity_jid=from_jid, node=c_node)
            )
            return

        try:
            # Per XEP-0115 §6.2, include the caps node in the disco#info request
            caps_node = f"{c_node}#{c_ver}"
            disco_infos = await G.host.get_disco_infos(client, from_jid, node=caps_node)
        except error.ConnectionDone:
            return
        except Exception as e:
            log.error(
                _("Couldn't retrieve disco info for {jid}: {error}").format(
                    jid=from_jid.full(), error=str(e)
                )
            )
            return

        # Validate disco info structure per XEP-0115 §5.4
        if not self._validate_caps(from_jid, c_ver, c_node, disco_infos):
            return

        # disco is validated, we can compute hash
        cap_hash = self.generate_hash(disco_infos)

        # Compare computed hash with received verification string
        if cap_hash != c_ver:
            log.warning(
                _(
                    "Computed hash differs from received hash for {jid}:\n"
                    "received: [{received}]\n"
                    "computed: [{computed}]\n"
                    "(node: {node})"
                ).format(
                    jid=from_jid.full(),
                    received=c_ver,
                    computed=cap_hash,
                    node=c_node,
                )
            )
            # It's invalid, we exit early and won't cache.
            return

        # We're all good, we can cache the disco info
        G.host.memory.disco.hashes[cap_hash] = disco_infos
        G.host.memory.update_entity_data(client, from_jid, C.ENTITY_CAP_HASH, cap_hash)
        log.debug(f"Caps validated and cached for {from_jid}: {cap_hash}")

    def generate_hash(self, services: disco.DiscoInfo) -> str:
        """Generate a unique hash for given service

        Hash algorithm is the one described in XEP-0115 §5.1 (Generation Method).
        @param services: disco infos as returned by discoHandler.info()

        """
        # Step 1: initialize empty string S
        s = []
        # Step 2: sort identities by category, then type, then xml:lang (i;octet
        # collation)
        byte_identities = [
            ByteIdentity(service)
            for service in services
            if isinstance(service, disco.DiscoIdentity)
        ]
        # i;octet (RFC 4790 §9.3) sorting.
        # Python's bytes comparison is equivalent to i;octet
        byte_identities.sort(key=lambda i: (i.category, i.idType, i.lang))
        # Step 3: append each identity formatted as cat/type/lang/name followed by '<'
        for identity in byte_identities:
            s.append(bytes(identity))
            s.append(b"<")
        # features
        byte_features = [
            service.encode("utf-8")
            for service in services
            if isinstance(service, disco.DiscoFeature)
        ]
        # Step 4: sort the features
        byte_features.sort()
        # Step 5: append each feature followed by '<'
        for feature in byte_features:
            s.append(feature)
            s.append(b"<")

        # Step 6: sort forms by FORM_TYPE; Step 7: process each form's fields
        ext = list(services.extensions.values())
        ext.sort(key=lambda f: f.formNamespace.encode("utf-8"))
        for extension in ext:
            # Step 7.a: append FORM_TYPE field's value (formNamespace) followed by '<'
            s.append(extension.formNamespace.encode("utf-8"))
            s.append(b"<")
            # Step 7.b: sort fields by "var" attribute
            fields = extension.fieldList
            fields.sort(key=lambda f: f.var.encode("utf-8"))
            for field in fields:
                # Step 7.c: "For each field other than FORM_TYPE"
                if field.var == "FORM_TYPE":
                    continue
                # Step 7.c.a: append var attribute followed by '<'
                s.append(field.var.encode("utf-8"))
                s.append(b"<")
                # Step 7.c.b: sort values (use raw XML character data per spec)
                values = [v.encode("utf-8") for v in field.raw_values]
                values.sort()
                # Step 7.c.c: append each value followed by '<'
                for value in values:
                    s.append(value)
                    s.append(b"<")

        # Step 9: compute verification string (SHA-1 -> Base64)
        cap_hash = b64encode(sha1(b"".join(s)).digest()).decode("utf-8")
        log.debug(_("Capability hash generated: [{cap_hash}]").format(cap_hash=cap_hash))
        return cap_hash

    def _validate_caps(
        self,
        from_jid: jid.JID,
        c_ver: str,
        c_node: str,
        disco_infos: disco.DiscoInfo,
    ) -> bool:
        """Validate capabilities disco info structure.

        @param from_jid: The JID that sent the presence with caps.
        @param c_ver: The verification string from the caps element.
        @param c_node: The node attribute from the caps element.
        @param disco_infos: The fetched disco#info response.
        @return: True if disco info is well-formed, False otherwise.
        """
        # §5.4 step 3.c: duplicate identities with same category/type/lang/name
        seen_identities = set()
        for item in disco_infos:
            if isinstance(item, disco.DiscoIdentity):
                lang = item.lang or ""
                identity_key = (item.category, item.type, lang)
                if identity_key in seen_identities:
                    log.warning(
                        _(
                            "Duplicate identity found in caps for {jid}: category={cat},"
                            " type={typ}, lang={lang}"
                        ).format(
                            jid=from_jid.full(),
                            cat=item.category,
                            typ=item.type,
                            lang=lang,
                        )
                    )
                    return False
                seen_identities.add(identity_key)

        # §5.4 step 3.d: duplicate features
        seen_features = set()
        for item in disco_infos:
            if isinstance(item, disco.DiscoFeature):
                if item in seen_features:
                    log.warning(
                        _("Duplicate feature found in caps for {jid}: {feat}").format(
                            jid=from_jid.full(), feat=item
                        )
                    )
                    return False
                seen_features.add(item)

        # §5.4 step 3.e: duplicate FORM_TYPEs or FORM_TYPE with multiple values
        seen_form_types = set()
        for form_type, form in disco_infos.extensions.items():
            if form_type in seen_form_types:
                log.warning(
                    _("Duplicate FORM_TYPE found in caps for {jid}: {ft}").format(
                        jid=from_jid.full(), ft=form_type
                    )
                )
                return False

            form_type_field = None
            for field in form.fieldList:
                if field.var == "FORM_TYPE":
                    form_type_field = field
                    break

            if form_type_field is None:
                # §5.4 step 3.f: ignore forms without FORM_TYPE field
                log.debug(
                    _(
                        "Form without FORM_TYPE field in caps for {jid}, ignoring form"
                    ).format(jid=from_jid.full())
                )
                continue

            if len(form_type_field.values) != 1:
                log.warning(
                    _(
                        "FORM_TYPE field has multiple values in caps for {jid}, ignoring"
                        " form"
                    ).format(jid=from_jid.full())
                )
                continue

            seen_form_types.add(form_type)

        return True


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

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

    @property
    def client(self) -> SatXMPPEntity:
        return self.parent  # type: ignore[return-value]

    def _caps_node_cb(self, disco_infos, ver):
        """Callback for caps node query."""
        G.host.memory.disco.hashes[ver] = disco_infos
        log.debug(f"Cached disco info for caps hash {ver}")

    def _caps_node_eb(self, failure_, server_jid, ver):
        """Errback for caps node query."""
        log.warning(f"Failed to query caps node for {server_jid} (ver={ver}): {failure_}")

    def connectionInitialized(self) -> None:
        assert self.xmlstream is not None
        # Check if server advertised caps in stream features (XEP-0115 §6.3)
        # xmlstream.features is a dict mapping (ns, name) to element
        caps_key = (NS_ENTITY_CAPABILITY, "c")
        if caps_key in self.xmlstream.features:
            c_elt = self.xmlstream.features[caps_key]
            ver = c_elt.getAttribute("ver")
            node = c_elt.getAttribute("node")
            if ver:
                server_jid = self.client.server_jid
                G.host.memory.update_entity_data(
                    self.client, server_jid, C.ENTITY_CAP_HASH, ver
                )

                # If we already have the disco info, we're done
                if ver in G.host.memory.disco.hashes:
                    log.debug(
                        f"Cached server caps from stream feature: {server_jid} ver={ver}"
                    )
                else:
                    # Query the caps node to get disco info (XEP-0115 §6.2)
                    caps_node = f"{node}#{ver}" if node else ver
                    log.debug(f"Querying caps node {caps_node} for server {server_jid}")
                    # We chain to `conn_deferred` to avoid race condition and be sure that
                    # we have `client.disco` set (see
                    # https://www.goffi.org/tickets/view/484)
                    d = self.client.conn_deferred
                    d.addCallback(
                        lambda __: self.client.disco.requestInfo(
                            server_jid, nodeIdentifier=caps_node
                        )
                    )
                    d.addCallback(self._caps_node_cb, ver)
                    d.addErrback(self._caps_node_eb, server_jid, ver)

        self.xmlstream.addObserver(
            CAPABILITY_UPDATE, self.plugin_parent._update, client=self.client
        )

    def getDiscoInfo(
        self,
        requestor: SatXMPPEntity,
        target: SatXMPPEntity,
        nodeIdentifier: str = "",
    ) -> list[disco.DiscoFeature | disco.DiscoIdentity | data_form.Form]:
        """
        Handle disco#info requests, including caps node queries.

        When nodeIdentifier is in the format 'node#ver' (e.g.,
        'http://libervia.org#abc123'), this is a caps node query per XEP-0115 §6.2
        (Discovering Capabilities). We return the cached disco info for that verification
        string.

        @param requestor: The entity that sent the request.
        @param target: The entity the request was sent to.
        @param nodeIdentifier: The optional node being queried.
        @return: List of disco items (features, identities, extensions).
        """

        # For non-caps queries, return our caps features
        return [
            disco.DiscoFeature(NS_ENTITY_CAPABILITY),
            disco.DiscoFeature(NS_CAPS_OPTIMIZE),
        ]

    def getDiscoItems(
        self,
        requestor: SatXMPPEntity,
        target: SatXMPPEntity,
        nodeIdentifier: str = "",
    ) -> list[disco.DiscoItem]:
        return []