#!/usr/bin/env python3

# Libervia plugin for Pubsub Node Relationships (XEP-0496).
# 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 TYPE_CHECKING, cast

from sqlalchemy import select, update
from twisted.words.protocols.jabber import jid
from twisted.words.protocols.jabber.jid import JID
from twisted.words.protocols.jabber.xmlstream import XMPPHandler
from tx_xmpp import 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
from libervia.backend.memory.sqla_mapping import (
    AccessModel,
    Affiliation,
    PublishModel,
    PubsubNode,
    SubscriptionState,
)
from libervia.backend.plugins.plugin_xep_0356 import XEP_0356
from libervia.backend.tools.utils import aio

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

log = getLogger(__name__)

PLUGIN_INFO = {
    C.PI_NAME: "Pubsub Node Relationships",
    C.PI_IMPORT_NAME: "XEP-0496",
    C.PI_TYPE: "XEP",
    C.PI_MODES: [C.PLUG_MODE_COMPONENT],
    C.PI_PROTOCOLS: ["XEP-0496"],
    C.PI_DEPENDENCIES: [
        "XEP-0060",
        "XEP-0356",
    ],
    C.PI_RECOMMENDATIONS: [],
    C.PI_MAIN: "XEP_0496",
    C.PI_HANDLER: "yes",
    C.PI_DESCRIPTION: _("""Pubsub node relationships with parent and link support."""),
}

NS_PUBSUB_RELATIONSHIPS = "urn:xmpp:pubsub-relationships:0"
PARENT_VAR = f"{{{NS_PUBSUB_RELATIONSHIPS}}}parent"
LINK_VAR = f"{{{NS_PUBSUB_RELATIONSHIPS}}}link"


class XEP_0496:

    def __init__(self, host: "LiberviaBackend") -> None:
        log.info(f"plugin {PLUGIN_INFO[C.PI_NAME]!r} initialization")
        self._xep_0356 = cast(XEP_0356, host.plugins["XEP-0356"])
        host.register_namespace("pubsub-relationships", NS_PUBSUB_RELATIONSHIPS)

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

    async def _check_cycle_recursive(
        self,
        client: SatXMPPEntity,
        service: JID,
        visited: set[str],
        parent: str | None,
        link: str | None,
    ) -> str | None:
        """Detect recursively if setting a relationship would create a cycle.

        @param client: Client instance.
        @param service: The pubsub service JID.
        @param visited: Names of nodes already visited. Must be initialised with child
            node name.
        @param parent: The parent node identifier (if any).
        @param link: The linked node identifier (if any).
        @return: The node name that would create a cycle, or None if no cycle.
        """
        for node_name, is_parent in ((parent, True), (link, False)):
            if node_name is not None:
                if node_name in visited:
                    return node_name
                else:
                    visited.add(node_name)
                    related_node = await G.storage.get_pubsub_node(
                        client, service, node_name, with_parent=True, with_link=True
                    )
                    if related_node is None:
                        node_type = "parent" if is_parent else "link"
                        raise exceptions.NotFound(
                            f"Invalid {node_type} node: {node_name}. Visited: {visited}"
                        )
                    if is_parent:
                        new_parent = (
                            related_node.parent_node.name
                            if related_node.parent_node
                            else None
                        )
                    else:
                        # We don't check parents for linking nodes, as they have the same
                        # as linked nodes, which will result in conflicts.
                        new_parent = None

                    new_link = (
                        related_node.linked_node.name
                        if related_node.linked_node
                        else None
                    )

                    cycling_node = await self._check_cycle_recursive(
                        client,
                        service,
                        visited,
                        new_parent,
                        new_link,
                    )
                    if cycling_node is not None:
                        return cycling_node

    @aio
    async def get_linking_nodes_by_id(self, linked_node_id: int) -> list[PubsubNode]:
        """Get all nodes that have a link relationship to the given node by ID.

        @param linked_node_id: The linked node ID.
        @return: Nodes that link to this node.
        """
        async with G.storage.session() as session:
            stmt = select(PubsubNode).where(PubsubNode.linked_node_id == linked_node_id)
            result = await session.execute(stmt)
        return result.unique().scalars().all()

    async def check_relationships(
        self,
        client: SatXMPPEntity,
        service: JID,
        node: PubsubNode,
        entity_jid: jid.JID,
        new_parent: str | None,
        new_link: str | None,
    ) -> tuple[PubsubNode | None, PubsubNode | None]:
        """Validate relationship changes according to XEP-0496 rules.

        @param client: Client instance.
        @param service: The pubsub service JID.
        @param node: The node being configured.
        @param entity_jid: The entity JID to check access for.
        @param new_parent: The new parent node identifier (if any).
        @param new_link: The new linked node identifier (if any).
        @return: Tuple of (parent_node, linked_node) if they were fetched, or None.
        @raises: exceptions on validation errors.
        """
        if new_parent is not None and node.linked_node_id is not None:
            raise exceptions.ConflictError(
                "Cannot set parent on a linking node; parent is automatically managed."
            )
        if (
            new_parent is None
            and node.linked_node_id is not None
            and node.parent_node_id is not None
        ):
            raise exceptions.ConflictError(
                "Cannot clear parent on a linking node; parent is automatically managed."
            )

        parent_node: PubsubNode | None = None
        linked_node: PubsubNode | None = None

        if new_parent is not None:
            # Validate if parent node exists and entity has the right to access it.
            parent_node = await G.storage.get_pubsub_node(
                client, service, new_parent, with_affiliations=True, with_parent=True
            )
            if parent_node is None:
                raise exceptions.NotFound(f"Parent node {new_parent!r} does not exist")

            has_parent_access = await self.has_access(
                client, service, parent_node, entity_jid
            )
            if not has_parent_access:
                raise exceptions.PermissionError(
                    f"You do not have access to parent node {new_parent!r} or one of its"
                    " parents, you can't use it as parent."
                )
            parent_node_id = parent_node.id
        else:
            parent_node_id = node.parent_node_id

        if new_link is not None:
            linked_node = await G.storage.get_pubsub_node(
                client, service, new_link, with_parent=True
            )
            if linked_node is None:
                raise exceptions.NotFound(f"Linked node {new_link!r} does not exist")

            # If linked node has a parent, linking node must have the same parent
            if linked_node.parent_node_id is not None:
                # If linking node has a different parent (not None, not matching), reject
                if (
                    parent_node_id is not None
                    and parent_node_id != linked_node.parent_node_id
                ):
                    raise exceptions.ConflictError(
                        "Linking node must have the same parent as linked node"
                        f" {new_link!r}"
                    )
                # If linking node has no parent, auto-set it to match linked node's parent
                if parent_node_id is None:
                    parent_node_id = linked_node.parent_node_id
                    # Fetch the parent node so it can be used in update_values
                    parent_node = await G.storage.get_pubsub_node(
                        client, service, node_id=parent_node_id
                    )

        # Check for cycles
        cycle_node = await self._check_cycle_recursive(
            client, service, {node.name}, new_parent, new_link
        )
        if cycle_node:
            raise exceptions.ConflictError(
                "Setting this relationship would create a cycle through node"
                " {cycle_node!r}"
            )

        return parent_node, linked_node

    async def _check_node_access(
        self,
        client: SatXMPPEntity,
        service: JID,
        node: PubsubNode,
        entity_jid: jid.JID,
    ) -> bool:
        """Check if an entity has access to a single node.

        @param client: Client instance.
        @param service: The pubsub service JID.
        @param node: The node to check access for. It must have ``affiliations`` loaded.
        @param entity_jid: The entity JID to check access for.
        @return: True if the entity has access to the node.
        """
        affiliation = next(
            (
                aff_data.affiliation
                for aff_data in node.affiliations
                if aff_data.entity == entity_jid
            ),
            None,
        )

        access_model = node.access_model or AccessModel.open

        if access_model == AccessModel.open:
            return True
        elif access_model == AccessModel.presence:
            return await self._check_presence_access(client, service, node, entity_jid)
        elif access_model == AccessModel.whitelist:
            return affiliation in (
                Affiliation.owner,
                Affiliation.publisher,
                Affiliation.member,
            )
        return False

    async def _check_presence_access(
        self,
        client: SatXMPPEntity,
        service: JID,
        node: PubsubNode,
        entity_jid: jid.JID,
    ) -> bool:
        """Check presence-based access for a node.

        @param client: Client instance.
        @param service: The pubsub service JID.
        @param node: The pubsub node.
        @param entity_jid: The entity JID to check access for.
        @return: True if requester has "from" or "both" subscription with publisher.
        """
        if not node.pep:
            raise NotImplementedError("Presence access model is only implemented for PEP")

        # For PEP, owner/publisher is the service JID.
        publisher_jid = service.userhostJID()

        try:
            roster = await self._xep_0356.get_roster(client, publisher_jid)
        except Exception:
            log.warning(
                f"Failed to get roster for {publisher_jid} when checking presence access."
            )
            return False

        roster_item = roster.get(entity_jid)
        if roster_item is None:
            return False

        # According to XEP-0060, subscription "from" or "both" are accepted.
        return roster_item.subscriptionFrom

    async def has_access(
        self, client: SatXMPPEntity, service: JID, node: PubsubNode, entity_jid: jid.JID
    ) -> bool:
        """Check if an entity has access to a node and all its parents.

        @param client: Client instance.
        @param service: The pubsub service JID.
        @param node: The node to check access for.
        @param entity_jid: The entity JID to check access for.
        @return: True if the entity has access to the node and all its parents.
        """
        if not await self._check_node_access(client, service, node, entity_jid):
            return False

        current_node = node

        while current_node.parent_node:
            parent_node = await G.storage.get_pubsub_node(
                client,
                service,
                current_node.parent_node.name,
                with_affiliations=True,
                with_parent=True,
            )
            if parent_node is None:
                break

            if not await self._check_node_access(
                client, service, parent_node, entity_jid
            ):
                return False

            current_node = parent_node

        return True

    def _check_node_publish_rights(self, node: PubsubNode, entity_jid: jid.JID) -> bool:
        """Check if an entity has publish rights to a single node.

        @param node: The node to check publish rights for. It must have ``affiliations``
            and ``subscriptions`` loaded.
        @param entity_jid: The entity JID to check publish rights for.
        @return: True if the entity has publish rights to the node.
        """
        affiliation = next(
            (
                aff_data.affiliation
                for aff_data in node.affiliations
                if aff_data.entity == entity_jid
            ),
            None,
        )

        publish_model = node.publish_model or PublishModel.publishers

        if publish_model == PublishModel.open:
            return True
        elif publish_model == PublishModel.publishers:
            return affiliation in (Affiliation.owner, Affiliation.publisher)
        elif publish_model == PublishModel.subscribers:
            if affiliation in (Affiliation.owner, Affiliation.publisher):
                return True
            # Check if subscribed
            return any(
                sub.subscriber == entity_jid and sub.state == SubscriptionState.SUBSCRIBED
                for sub in node.subscriptions
            )
        return False

    async def has_publish_rights(
        self, client: SatXMPPEntity, service: JID, node: PubsubNode, entity_jid: jid.JID
    ) -> bool:
        """Check if an entity has publish rights to a node and all its parents.

        @param client: Client instance.
        @param service: The pubsub service JID.
        @param node: The node to check publish rights for. It must have ``affiliations``
            and ``subscriptions`` loaded. The later is only used in case of
            ``subscribers`` access model.
        @param entity_jid: The entity JID to check publish rights for.
        @return: True if the entity has publish rights to the node and all its parents.
        """
        if not self._check_node_publish_rights(node, entity_jid):
            return False

        # Check the entire hierarchy
        current_node = node
        while current_node.parent_node:
            parent_node = await G.storage.get_pubsub_node(
                client,
                service,
                current_node.parent_node.name,
                with_affiliations=True,
                with_subscriptions=True,
            )
            if parent_node is None:
                log.error(
                    f"Node parent doesn't exist: {client.profile=} {service=}"
                    f" node_name={current_node.parent_node.name}"
                )
                break

            if not self._check_node_publish_rights(parent_node, entity_jid):
                return False

            current_node = parent_node

        return True

    async def get_descendant_nodes(self, node: PubsubNode) -> list[PubsubNode]:
        """Get all descendant of a node.

        @param node: Parent node.
        @return: List of all descendant nodes.
        """
        descendants: list[PubsubNode] = []
        queue: list[PubsubNode] = [node]

        while queue:
            current_node = queue.pop(0)
            children = await G.storage.get_pubsub_node_children(current_node)

            for child in children:
                descendants.append(child)
                queue.append(child)

        return descendants

    @aio
    async def update_linking_nodes_parent(
        self, linked_node_id: int, new_parent_id: int | None
    ) -> None:
        """Update parent_node_id for all linking nodes when linked node parent changes.

        This method is called when a linked node's parent is updated. All nodes that
        link to this node must have their parent_node_id updated to match.

        @param linked_node_id: The linked node ID.
        @param new_parent_id: The new parent node ID (or None to clear parent).
        """
        async with G.storage.session() as session:
            stmt = (
                update(PubsubNode)
                .where(PubsubNode.linked_node_id == linked_node_id)
                .values(parent_node_id=new_parent_id)
            )
            await session.execute(stmt)
            await session.flush()


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

    def getDiscoInfo(
        self, requestor: JID, target: JID, nodeIdentifier: str = ""
    ) -> list[disco.DiscoFeature]:
        return [disco.DiscoFeature(NS_PUBSUB_RELATIONSHIPS)]

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