#!/usr/bin/env python3


# Libervia XMPP client
# 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 collections.abc import Iterable
from typing import TYPE_CHECKING, cast

from twisted.internet.interfaces import IReactorCore
from libervia.backend import G
from libervia.backend.core.i18n import _
from libervia.backend.core import exceptions
from libervia.backend.core.log import getLogger
from libervia.backend.core.core_types import SatXMPPEntity

from twisted.words.protocols.jabber import jid
from twisted.words.protocols.jabber.error import StanzaError
from twisted.internet import defer
from twisted.internet import reactor
from twisted.python import failure
from libervia.backend.core.constants import Const as C
from libervia.backend.tools import xml_tools
from libervia.backend.memory import persistent
from tx_xmpp import disco
import time

from libervia.backend.tools.utils import ensure_deferred, fork_deferred

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

log = getLogger(__name__)


TIMEOUT = 15
# Cache TTL for successful disco queries (in seconds)
DISCO_CACHE_TTL = 900
# Delay before retrying a failed disco query (in seconds)
DISCO_RETRY_DELAY = 300
# Delay before retrying a server that failed (in seconds)
DISCO_SERVER_RETRY_DELAY = 3600


class HashGenerationError(Exception):
    pass


class HashManager:
    """map object which manage hashes

    persistent storage is update when a new hash is added
    """

    def __init__(self, persistent):
        self.hashes: dict[str, disco.DiscoInfo] = {}
        self.persistent = persistent

    def __getitem__(self, key) -> disco.DiscoInfo:
        return self.hashes[key]

    def __setitem__(self, hash_: str, disco_info: disco.DiscoInfo) -> None:
        if hash_ in self.hashes:
            log.debug("ignoring hash set: it is already known")
            return
        if disco_info is None:
            msg = f"Cannot cache 'None' for hash {hash_}"
            log.exception(msg)
            raise exceptions.InternalError(msg)
        self.hashes[hash_] = disco_info
        self.persistent[hash_] = disco_info.toElement().toXml()

    def __contains__(self, hash_: str) -> bool:
        return self.hashes.__contains__(hash_)

    def load(self):
        def fill_hashes(hashes):
            for hash_, xml in hashes.items():
                element = xml_tools.ElementParser()(xml)
                disco_info = disco.DiscoInfo.fromElement(element)
                for ext_form in disco_info.extensions.values():
                    # tx_xmpp doesn't call typeCheck on reception, so we do it here
                    ext_form.typeCheck()
                if not disco_info.features and not disco_info.identities:
                    log.warning(
                        _(
                            "no feature/identity found in disco element (hash: "
                            "{cap_hash}), ignoring: {xml}"
                        ).format(cap_hash=hash_, xml=xml)
                    )
                else:
                    self.hashes[hash_] = disco_info

            log.info("Disco hashes loaded")

        d = self.persistent.load()
        d.addCallback(fill_hashes)
        return d


class Discovery:
    """Manage capabilities of entities"""

    def __init__(self, host: "LiberviaBackend"):
        self.host = host
        # JID-based disco info cache, it is used for entities without XEP-0115 support or
        # when something goes wrong with hash validation.
        # Cache entries are tuples: (data, expiry_timestamp)
        #   - (disco.DiscoInfo, float): successful query, expires at timestamp
        #   - (None, float): failed query, retry after DISCO_RETRY_DELAY
        self._jid_disco_cache: dict[jid.JID, tuple[disco.DiscoInfo | None, float]] = {}
        # Track in-flight disco queries to avoid duplicate requests
        # Maps (jid, node) tuple to query's Deferred
        self._pending_queries: dict[tuple[jid.JID, str], defer.Deferred] = {}

    def load(self):
        """Load persistent hashes"""
        self.hashes = HashManager(persistent.PersistentDict("disco"))
        return self.hashes.load()

    async def has_feature(
        self,
        client: SatXMPPEntity,
        feature: str,
        jid_: jid.JID | None = None,
        node: str = "",
    ) -> bool:
        """Tell if an entity has the required feature

        @param feature: feature namespace
        @param jid_: jid of the target, or None for profile's server
        @param node: optional node to use for disco request
        @return: a Deferred which fire a boolean (True if feature is available)
        """
        try:
            disco_infos = await self.get_infos(client, jid_, node)
        except (exceptions.ExternalRequestError, exceptions.TimeOutError):
            return False
        return feature in disco_infos.features

    async def check_feature(
        self,
        client: SatXMPPEntity,
        feature: str,
        jid_: jid.JID | None = None,
        node: str = "",
    ) -> None:
        """Like has_feature, but raise an exception is feature is not Found

        @param feature: feature namespace
        @param jid_: jid of the target, or None for profile's server
        @param node: optional node to use for disco request

        @raise: exceptions.FeatureNotFound
        """
        disco_infos = await self.get_infos(client, jid_, node)
        if not feature in disco_infos.features:
            raise exceptions.FeatureNotFound()

    async def check_features(
        self,
        client: SatXMPPEntity,
        features: Iterable[str],
        jid_: jid.JID | None = None,
        identity: tuple[str, str] | None = None,
        node: str = "",
    ) -> None:
        """Like check_feature, but check several features at once, and check also identity

        @param features: features to check
        @param jid_: jid of the target, or None for profile's server
        @param node: optional node to use for disco request
        @param identity: if not None, the entity must have an identity with this
            (category, type) tuple

        @raise: exceptions.FeatureNotFound
        """
        disco_infos = await self.get_infos(client, jid_, node)
        if not set(features).issubset(disco_infos.features):
            raise exceptions.FeatureNotFound()

        if identity is not None and identity not in disco_infos.identities:
            raise exceptions.FeatureNotFound()

    async def has_identity(
        self,
        client: SatXMPPEntity,
        category: str,
        type_: str,
        jid_: jid.JID | None = None,
        node: str = "",
    ) -> bool:
        """Tell if an entity has the requested identity

        @param category: identity category
        @param type_: identity type
        @param jid_: jid of the target, or None for profile's server
        @param node: optional node to use for disco request
        @return: True if the entity has the given identity
        """
        disco_infos = await self.get_infos(client, jid_, node)
        return (category, type_) in disco_infos.identities

    def get_infos(
        self,
        client: SatXMPPEntity,
        jid_: jid.JID | None = None,
        node: str = "",
        use_cache: bool = True,
    ) -> defer.Deferred[disco.DiscoInfo]:
        """get disco infos from jid_, filling capability hash if needed

        @param client: client session
        @param jid_: jid of the target, or None for profile's server
        @param node: optional node to use for disco request
        @param use_cache: if True, use cached data if available
        @return: a Deferred which fire disco.DiscoInfo
        @raise exceptions.ExternalRequestError: disco query failed and retry delay not
            elapsed
        @raise exceptions.TimeOutError: the request was cancelled
        """
        if jid_ is None:
            jid_ = jid.JID(client.jid.host)

        try:
            if not use_cache:
                # we ignore cache, so we pretend we haven't found it
                raise KeyError
            cap_hash = self.host.memory.entity_data_get(
                client, jid_, [C.ENTITY_CAP_HASH]
            )[C.ENTITY_CAP_HASH]
        except (KeyError, exceptions.UnknownEntityError):
            # Capability hash is not available, check JID-based cache first
            if use_cache and jid_ in self._jid_disco_cache:
                data, timestamp = self._jid_disco_cache[jid_]
                elapsed = time.time() - timestamp
                if data is None:
                    # Cached failure
                    if elapsed < DISCO_RETRY_DELAY:
                        log.debug(
                            f"Cached disco failure for {jid_.full()}, retry in "
                            f"{DISCO_RETRY_DELAY - int(elapsed)} seconds"
                        )
                        return defer.fail(
                            exceptions.ExternalRequestError(
                                f"Disco query failed for {jid_.full()}, retrying in "
                                f"{DISCO_RETRY_DELAY} seconds"
                            )
                        )
                    # Retry delay expired, remove from cache and retry
                    log.debug(f"Disco retry delay expired for {jid_.full()}, retrying")
                    del self._jid_disco_cache[jid_]
                elif elapsed < DISCO_CACHE_TTL:
                    # Successful cache hit
                    log.debug(f"JID-based disco cache hit for {jid_.full()}")
                    return defer.succeed(data)
                else:
                    # Cache expired, remove and retry
                    log.debug(f"Disco cache expired for {jid_.full()}, re-querying")
                    del self._jid_disco_cache[jid_]

            # Entity-level cache miss or expired, check server-level cache
            if use_cache and (server_jid := jid.JID(jid_.host)) in self._jid_disco_cache:
                data, timestamp = self._jid_disco_cache[server_jid]
                if data is None:
                    elapsed = time.time() - timestamp
                    if elapsed < DISCO_SERVER_RETRY_DELAY:
                        msg = (
                            f"Server {server_jid} is unreachable, retry in "
                            f"{int(DISCO_SERVER_RETRY_DELAY - elapsed)} seconds"
                        )
                        log.debug(msg)
                        return defer.fail(exceptions.ExternalRequestError(msg))
                    else:
                        log.debug(
                            f"Retry delay expired for unreachable server {server_jid}"
                        )
                        del self._jid_disco_cache[server_jid]

            # No cache is used, we will do a query to return disco infos.
            # The query result will be cached in JID cache.
            query_key = (jid_, node)

            if query_key in self._pending_queries:
                log.debug(f"Reusing pending disco query for {jid_.full()}")
                return fork_deferred(self._pending_queries[query_key])

            def cleanup_query(result):
                """Remove the pending query on success or failure."""
                self._pending_queries.pop(query_key, None)
                return result

            def infos_cb(disco_infos):
                for ext_form in disco_infos.extensions.values():
                    # tx_xmpp doesn't call typeCheck on reception, so we do it here
                    # to avoid ending up with incorrect types.
                    ext_form.typeCheck()
                # Cache disco info per JID with expiry timestamp
                self._jid_disco_cache[jid_] = (disco_infos, time.time())
                return disco_infos

            def infos_eb(fail: failure.Failure):
                if fail.check(defer.CancelledError):
                    reason = "request time-out"
                    fail = failure.Failure(exceptions.TimeOutError(str(fail.value)))
                elif fail.check(StanzaError):
                    stanza_error = fail.value
                    if stanza_error.condition in (
                        "remote-server-not-found",
                        "remote-server-timeout",
                    ):
                        # Mark the server as unreachable (store as domain JID)
                        server_jid = jid.JID(jid_.host)
                        self._jid_disco_cache[server_jid] = (None, time.time())
                        log.warning(
                            f"Server {server_jid} is unreachable, will retry in "
                            f"{DISCO_SERVER_RETRY_DELAY} seconds"
                        )
                    reason = str(fail.value)
                else:
                    try:
                        reason = str(fail.value)
                    except AttributeError:
                        reason = str(fail)

                log.warning(f"can't request disco infos from {jid_.full()}: {reason}")

                # Cache the failure with timestamp to avoid repeated queries
                self._jid_disco_cache[jid_] = (None, time.time())
                raise fail

            d = client.disco.requestInfo(jid_, nodeIdentifier=node)
            d.addBoth(cleanup_query)
            d.addCallback(infos_cb)
            d.addErrback(infos_eb)
            self._pending_queries[query_key] = d
            return fork_deferred(d)
        else:
            disco_infos = self.hashes[cap_hash]
            return defer.succeed(disco_infos)

    async def get_items(
        self,
        client: SatXMPPEntity,
        jid_: jid.JID | None = None,
        node: str = "",
        use_cache: bool = True,
    ) -> disco.DiscoItems:
        """get disco items from jid_, cache them for our own server.

        @param client: client session
        @param jid_: jid of the target, or None for profile's server
        @param node: optional node to use for disco request
        @param use_cache: if True, use cached data if available
        @return: disco.DiscoItems
        """
        # FIXME: RSM is not managed here
        if jid_ is None:
            jid_ = client.server_jid

        if jid_ == client.server_jid and not node:
            # we cache items only for our own server and if node is not set
            try:
                items = G.host.memory.entity_data_get(client, jid_, ["DISCO_ITEMS"])[
                    "DISCO_ITEMS"
                ]
                log.debug("[%s] disco items are in cache" % jid_.full())
                if not use_cache:
                    # we ignore cache, so we pretend we haven't found it
                    raise KeyError
            except (KeyError, exceptions.UnknownEntityError):
                log.debug(f"Caching {jid_} disco items")
                items = await client.disco.requestItems(jid_, nodeIdentifier=node)
                G.host.memory.update_entity_data(client, jid_, "DISCO_ITEMS", items)
        else:
            try:
                items = await client.disco.requestItems(jid_, nodeIdentifier=node)
            except StanzaError as e:
                log.warning(
                    "Error while requesting items for {jid}: {reason}".format(
                        jid=jid_.full(), reason=e.condition
                    )
                )
                items = disco.DiscoItems()

        return items

    def _infos_eb(
        self, failure_: failure.Failure, entity_jid: jid.JID
    ) -> failure.Failure:
        log.warning(f"Error while requesting [{entity_jid}]: {failure_}")
        return failure_

    async def find_service_entity(
        self,
        client: SatXMPPEntity,
        category: str,
        type_: str,
        jid_: jid.JID | None = None,
    ) -> jid.JID | None:
        """Helper method to find first available entity from find_service_entities

        args are the same as for [find_service_entities]
        @param client: client session
        @param category: identity category
        @param type_: identity type
        @param jid_: jid of the target server, or None for profile's server
        @return: first found entity or None
        """
        entities = await self.host.find_service_entities(client, category, type_, jid_)
        return entities.pop() if entities else None

    async def find_service_entities(
        self,
        client: SatXMPPEntity,
        category: str,
        type_: str,
        jid_: jid.JID | None = None,
    ) -> set[jid.JID]:
        """Return all available items of an entity which correspond to (category, type_)

        @param client: client session
        @param category: identity's category
        @param type_: identitiy's type
        @param jid_: the jid of the target server (None for profile's server)
        @return: a set of found entities
        @raise defer.CancelledError: the request timed out
        """
        found_entities: set[jid.JID] = set()

        def infos_cb(infos, entity_jid):
            if (category, type_) in infos.identities:
                found_entities.add(entity_jid)

        items: disco.DiscoItems = await self.get_items(client, jid_)
        defers_list = []
        for item in items:
            info_d = self.get_infos(client, item.entity)

            info_d.addCallbacks(
                infos_cb, self._infos_eb, [item.entity], None, [item.entity]
            )
            defers_list.append(info_d)
        d = defer.DeferredList(defers_list)
        reactor.callLater(
            TIMEOUT, d.cancel
        )  # FIXME: one bad service make a general timeout
        await d
        return found_entities

    async def find_features_set(
        self,
        client: SatXMPPEntity,
        features: Iterable[str],
        identity: tuple[str, str] | None = None,
        jid_: jid.JID | None = None,
    ) -> set[jid.JID]:
        """Return entities (including jid_ and its items) offering features

        @param client: Client session.
        @param features: iterable of features which must be present
        @param identity: if not None, accept only this (category/type) identity
        @param jid_: the jid of the target server (None for profile's server)
        @return: a set of found entities
        """
        if jid_ is None:
            jid_ = jid.JID(client.jid.host)
        features = set(features)
        found_entities = set()

        def infos_cb(infos: disco.DiscoInfo, entity: jid.JID) -> None:
            if entity is None:
                log.warning(_("Received an item without JID"))
                return
            if identity is not None and identity not in infos.identities:
                return
            if features.issubset(infos.features):
                found_entities.add(entity)

        items = await self.get_items(client, jid_)
        defer_list = []
        for entity in [jid_] + [item.entity for item in items]:
            infos_d = self.get_infos(client, entity)
            infos_d.addCallbacks(infos_cb, self._infos_eb, [entity], None, [entity])
            infos_d.addTimeout(TIMEOUT, cast(IReactorCore, reactor))
            defer_list.append(infos_d)
        await defer.DeferredList(defer_list)
        return found_entities

    @ensure_deferred
    async def _disco_infos(
        self, entity_jid_s, node="", use_cache=True, profile_key=C.PROF_KEY_NONE
    ):
        """Discovery method for the bridge
        @param entity_jid_s: entity we want to discover
        @param use_cache(bool): if True, use cached data if available
        @param node(unicode): optional node to use

        @return: list of tuples
        """
        client = self.host.get_client(profile_key)
        entity = jid.JID(entity_jid_s)
        disco_infos = await self.get_infos(client, entity, node, use_cache)
        extensions = {}
        # FIXME: should extensions be serialised using tools.common.data_format?
        for form_type, form in list(disco_infos.extensions.items()):
            fields = []
            for field in form.fieldList:
                data = {"type": field.fieldType}
                for attr in ("var", "label", "desc"):
                    value = getattr(field, attr)
                    if value is not None:
                        data[attr] = value

                values = [field.value] if field.value is not None else field.values
                if field.fieldType == "boolean":
                    values = [C.bool_const(v) for v in values]
                fields.append((data, values))

            extensions[form_type or ""] = fields

        return (
            [str(f) for f in disco_infos.features],
            [
                (cat, type_, lang or "", identity.name or "")
                for (cat, type_), lang_identities in list(disco_infos.identities.items())
                for lang, identity in lang_identities.items()
            ],
            extensions,
        )

    def items2tuples(self, disco_items):
        """convert disco items to tuple of strings

        @param disco_items(iterable[disco.DiscoItem]): items
        @return G(tuple[unicode,unicode,unicode]): serialised items
        """
        for item in disco_items:
            if not item.entity:
                log.warning(_("invalid item (no jid)"))
                continue
            yield (item.entity.full(), item.nodeIdentifier or "", item.name or "")

    @ensure_deferred
    async def _disco_items(
        self,
        entity_jid_s: str,
        node: str = "",
        use_cache: bool = True,
        profile_key: str = C.PROF_KEY_NONE,
    ) -> list[tuple[str, str, str]]:
        """Discovery method for the bridge.

        @param entity_jid_s: entity we want to discover
        @param node: optional node to use
        @param use_cache: if True, use cached data if available
        @return: list of (jid, node, name) tuples
        """
        client = self.host.get_client(profile_key)
        entity = jid.JID(entity_jid_s)
        disco_items = await self.get_items(client, entity, node, use_cache)
        return list(self.items2tuples(disco_items))