#!/usr/bin/env python3

# Libervia plugin for Privileged Entity (XEP-0356).
# 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 textwrap import dedent
from typing import TYPE_CHECKING, Final, Literal, cast
import uuid

from twisted.internet import defer, reactor
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
from tx_xmpp.xmppim import RosterItem
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 SatXMPPComponent, SatXMPPEntity
from libervia.backend.core.i18n import _
from libervia.backend.core.log import getLogger
from libervia.backend.plugins.plugin_xep_0115 import XEP_0115
from libervia.backend.plugins.plugin_xep_0297 import NS_FORWARD, XEP_0297

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

log = getLogger(__name__)

PLUGIN_INFO = {
    C.PI_NAME: "Privileged Entity",
    C.PI_IMPORT_NAME: "XEP-0356",
    C.PI_TYPE: "XEP",
    C.PI_MODES: [C.PLUG_MODE_COMPONENT],
    C.PI_PROTOCOLS: ["XEP-0356"],
    C.PI_DEPENDENCIES: ["XEP-0115", "XEP-0297"],
    C.PI_RECOMMENDATIONS: [],
    C.PI_MAIN: "XEP_0356",
    C.PI_HANDLER: "yes",
    C.PI_DESCRIPTION: dedent(_("""\
    Privileged entity implementation for components.

    This plugin provides privileged entity capabilities to components, allowing them to:
    - Access and modify user rosters (with roster push support)
    - Send messages on behalf of server
    - Send IQ stanzas on behalf of server
    - Receive presence information for managed entities and roster contacts
    """)),
}

NS_PRIVILEGE: Final = "urn:xmpp:privilege:2"
NS_ROSTER = "jabber:iq:roster"
ROSTER_TTL = 3600
IMPORT_NAME = "XEP-0356"

RosterPrivilegeType = Literal["none", "get", "set", "both"]
MessagePrivilegeType = Literal["none", "outgoing"]
PresencePrivilegeType = Literal["none", "managed_entity", "roster"]


@dataclass
class IQPermission:
    """IQ permission for a namespace."""

    get: bool
    set: bool


@dataclass
class RosterCacheEntry:
    """Cached roster data for a managed entity."""

    roster: dict[jid.JID, RosterItem]
    version: str | None
    timestamp: float


@dataclass
class PrivilegeData:
    """Data structure for privilege tracking."""

    roster_access: RosterPrivilegeType = "none"
    roster_push: bool = False
    message_privilege: MessagePrivilegeType = "none"
    presence_privilege: PresencePrivilegeType = "none"
    iq_permissions: dict[str, IQPermission] = field(default_factory=dict)


class PrivilegeError(exceptions.PermissionError):
    """A privileged operation failed"""

    pass


@dataclass
class ClientPresenceData:
    """Presence tracking data for a client."""

    # bare_jid -> set of online resource strings
    online_resources: dict[jid.JID, set[str]] = field(default_factory=dict)


@dataclass
class ClientPrivilegeData:
    privileges: PrivilegeData = field(default_factory=lambda: PrivilegeData())
    roster_cache: dict[jid.JID, RosterCacheEntry] = field(default_factory=dict)
    presence_data: ClientPresenceData = field(
        default_factory=lambda: ClientPresenceData()
    )


class XEP_0356:
    """XEP-0356: Privileged Entity handler."""

    namespace = NS_PRIVILEGE

    def __init__(self, host: "LiberviaBackend") -> None:
        log.info(f"plugin {PLUGIN_INFO[C.PI_NAME]!r} initialization")
        self._xep_0115 = cast(XEP_0115, host.plugins["XEP-0115"])
        self._forward = cast(XEP_0297, host.plugins["XEP-0297"])
        host.register_namespace("privilege", NS_PRIVILEGE)

    def get_handler(self, client: SatXMPPEntity) -> "XEP0356Handler":
        """Return the handler for this plugin.

        @param client: The client entity.
        @return: The XEP0356Handler instance.
        """
        return XEP0356Handler(self)

    def profile_connected(self, client: SatXMPPEntity) -> None:
        """Called when a client is connected.

        Observers are now registered in XEP0356Handler.connectionInitialized()
        to avoid race conditions with privilege advertisements.

        @param client: The client connected.
        """
        if not client.is_component:
            return
        log.debug(f"XEP-0356 profile_connected called for {client.profile}")

    def can_access_roster(self, client: SatXMPPComponent) -> bool:
        """Check if this entity can access user rosters."""
        data = client.get_plugin_data(IMPORT_NAME, ClientPrivilegeData)
        return data.privileges.roster_access in ("get", "both")

    def can_modify_roster(self, client: SatXMPPComponent) -> bool:
        """Check if this entity can modify user rosters."""
        data = client.get_plugin_data(IMPORT_NAME, ClientPrivilegeData)
        return data.privileges.roster_access in ("set", "both")

    def can_receive_roster_pushes(self, client: SatXMPPComponent) -> bool:
        """Check if this entity receives roster pushes."""
        data = client.get_plugin_data(IMPORT_NAME, ClientPrivilegeData)
        return data.privileges.roster_push

    def can_send_messages(self, client: SatXMPPComponent) -> bool:
        """Check if this entity can send privileged messages."""
        data = client.get_plugin_data(IMPORT_NAME, ClientPrivilegeData)
        return data.privileges.message_privilege == "outgoing"

    def can_receive_managed_presence(self, client: SatXMPPComponent) -> bool:
        """Check if this entity can receive presence for managed entities."""
        data = client.get_plugin_data(IMPORT_NAME, ClientPrivilegeData)
        return data.privileges.presence_privilege in ("managed_entity", "roster")

    def can_receive_roster_presence(self, client: SatXMPPComponent) -> bool:
        """Check if this entity can receive presence for roster contacts."""
        data = client.get_plugin_data(IMPORT_NAME, ClientPrivilegeData)
        return data.privileges.presence_privilege == "roster"

    def can_send_iq(self, client: SatXMPPComponent, namespace: str, iq_type: str) -> bool:
        """Check if this entity can send privileged IQ for a namespace/type.

        @param client: The client to check.
        @param namespace: The namespace of the IQ payload.
        @param iq_type: The type of IQ ("get" or "set").
        @return: True if the entity can send this IQ.
        """
        data = client.get_plugin_data(IMPORT_NAME, ClientPrivilegeData)
        perms = data.privileges.iq_permissions.get(namespace)
        if perms is None:
            return False
        return perms.get if iq_type == "get" else perms.set

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

        @param msg: The message stanza containing privilege information.
        @param client: The client receiving the advertisement.
        """
        try:
            privilege_elt = next(msg.elements(NS_PRIVILEGE, "privilege"))
        except StopIteration:
            log.warning("Privilege advertisement without privilege element")
            return

        server_jid = client.server_jid
        if jid.JID(msg["from"]) != server_jid:
            log.warning(
                f"SECURITY: privilege advertisement from unexpected server: {msg['from']}"
            )
            raise error.StanzaError("not-allowed")

        data = client.get_plugin_data(IMPORT_NAME, ClientPrivilegeData)
        data.privileges = PrivilegeData()

        for perm in privilege_elt.elements(NS_PRIVILEGE, "perm"):
            try:
                access = perm["access"]
            except KeyError:
                log.warning(f"perm element without access attribute: {perm.toXml()}")
                continue

            if access == "roster":
                self._handle_roster_perm(data, perm)
            elif access == "message":
                self._handle_message_perm(data, perm)
            elif access == "presence":
                self._handle_presence_perm(data, perm)
            elif access == "iq":
                self._handle_iq_perm(data, perm)
            else:
                log.warning(f"Unknown privilege access type: {access}")

        log.info(
            f"Privileges updated for {client.profile}:\n"
            f"  roster: {data.privileges.roster_access} "
            f"(push={data.privileges.roster_push})\n"
            f"  message: {data.privileges.message_privilege}\n"
            f"  presence: {data.privileges.presence_privilege}\n"
            f"  iq: {list(data.privileges.iq_permissions.keys())}"
        )

    def _handle_roster_perm(
        self, data: ClientPrivilegeData, perm: domish.Element
    ) -> None:
        """Handle roster permission element."""
        try:
            perm_type = perm["type"]
        except KeyError:
            log.warning(f"roster perm without type: {perm.toXml()}")
            return

        if perm_type not in ("none", "get", "set", "both"):
            log.warning(f"Invalid roster permission type: {perm_type}")
            return

        data.privileges.roster_access = cast(RosterPrivilegeType, perm_type)

        push_attr = perm.getAttribute("push")
        if push_attr is not None:
            data.privileges.roster_push = push_attr.lower() == "true"
        else:
            data.privileges.roster_push = perm_type in ("get", "both")

    def _handle_message_perm(
        self, data: ClientPrivilegeData, perm: domish.Element
    ) -> None:
        """Handle message permission element."""
        try:
            perm_type = perm["type"]
        except KeyError:
            log.warning(f"message perm without type: {perm.toXml()}")
            return

        if perm_type not in ("none", "outgoing"):
            log.warning(f"Invalid message permission type: {perm_type}")
            return

        data.privileges.message_privilege = cast(MessagePrivilegeType, perm_type)

    def _handle_presence_perm(
        self, data: ClientPrivilegeData, perm: domish.Element
    ) -> None:
        """Handle presence permission element."""
        try:
            perm_type = perm["type"]
        except KeyError:
            log.warning(f"presence perm without type: {perm.toXml()}")
            return

        if perm_type not in ("none", "managed_entity", "roster"):
            log.warning(f"Invalid presence permission type: {perm_type}")
            return

        data.privileges.presence_privilege = cast(PresencePrivilegeType, perm_type)

    def _handle_iq_perm(self, data: ClientPrivilegeData, perm: domish.Element) -> None:
        """Handle IQ permission element."""
        for ns_elt in perm.elements(NS_PRIVILEGE, "namespace"):
            try:
                ns = ns_elt["ns"]
                perm_type = ns_elt["type"]
            except KeyError:
                log.warning(
                    f"namespace element without required attributes: {ns_elt.toXml()}"
                )
                continue

            if perm_type not in ("get", "set", "both"):
                log.warning(f"Invalid IQ permission type: {perm_type}")
                continue

            data.privileges.iq_permissions[ns] = IQPermission(
                get=perm_type in ("get", "both"),
                set=perm_type in ("set", "both"),
            )

    def on_roster_push(self, client: SatXMPPComponent, iq_elt: domish.Element) -> None:
        """Handle roster push from server.

        @param client: The client receiving the roster push.
        @param iq_elt: The roster push IQ stanza.
        """
        if not self.can_receive_roster_pushes(client):
            log.debug("Roster push received but not configured")
            return

        from_jid_str = iq_elt.getAttribute("from")
        if not from_jid_str:
            log.warning(f"Roster push without 'from' attribute: {iq_elt.toXml()}")
            return

        from_jid = jid.JID(from_jid_str)
        if from_jid.resource:
            log.warning(
                f"Roster push from full JID instead of bare JID: {iq_elt.toXml()}"
            )
            return

        data = client.get_plugin_data(IMPORT_NAME, ClientPrivilegeData)
        cache = data.roster_cache

        # Parse roster item using tx-xmpp
        query = next(iq_elt.elements(NS_ROSTER, "query"))
        item_elt = next(query.elements(NS_ROSTER, "item"))
        roster_item = RosterItem.fromElement(item_elt)
        item_jid = roster_item.entity

        cached_entry = cache.get(from_jid)
        roster = cached_entry.roster if cached_entry else {}
        if roster_item.remove:
            roster.pop(item_jid, None)
        else:
            roster[item_jid] = roster_item

        version = query.getAttribute("ver")
        cache[from_jid] = RosterCacheEntry(
            roster=roster,
            version=version,
            timestamp=reactor.seconds(),
        )

        log.debug(f"Roster push for {from_jid}: {item_jid}")

    async def get_roster(
        self, client: SatXMPPComponent, user_jid: jid.JID
    ) -> dict[jid.JID, RosterItem]:
        """Get the roster of a managed entity.

        @param client: The client making the request.
        @param user_jid: The bare JID of the user whose roster to retrieve.
        @return: Dict mapping JIDs to RosterItem objects.
        @raise PrivilegeError: If roster access is not permitted.
        """
        if not self.can_access_roster(client):
            raise PrivilegeError(f"Roster access not permitted for {user_jid}")

        data = client.get_plugin_data(IMPORT_NAME, ClientPrivilegeData)
        cache = data.roster_cache
        cached = cache.get(user_jid)
        if cached and (reactor.seconds() - cached.timestamp) < ROSTER_TTL:
            return cached.roster

        iq = client.IQ("get")
        iq["to"] = user_jid.full()
        query = domish.Element((NS_ROSTER, "query"))
        iq.addChild(query)

        result = await iq.send()

        roster = {}
        for item_elt in result.elements(NS_ROSTER, "item"):
            roster_item = RosterItem.fromElement(item_elt)
            roster[roster_item.entity] = roster_item

        version = query.getAttribute("ver")
        cache[user_jid] = RosterCacheEntry(
            roster=roster,
            version=version,
            timestamp=reactor.seconds(),
        )

        return roster

    async def send_privileged_message(
        self, client: SatXMPPComponent, message_elt: domish.Element
    ) -> None:
        """Send a message on behalf of a user.

        @param client: The client sending the message.
        @param message_elt: The message stanza to send (will be wrapped in forwarded).
        @raise PrivilegeError: If message privilege is not permitted.
        """
        if not self.can_send_messages(client):
            raise PrivilegeError("Message privilege not permitted")

        try:
            from_attr = message_elt["from"]
        except KeyError:
            raise PrivilegeError('Message must have a "from" attribute')

        from_jid = jid.JID(from_attr)
        if from_jid.resource or jid.JID(from_jid.host) != client.server_jid:
            raise PrivilegeError(
                f"from_jid must be a bare JID from our server: {from_jid}"
            )

        wrapper = domish.Element((None, "message"))
        wrapper["from"] = client.jid.full()
        wrapper["to"] = client.server_jid.full()
        wrapper["id"] = str(uuid.uuid4())

        privilege = domish.Element((NS_PRIVILEGE, "privilege"))
        forwarded = self._forward.build_forwarded(message_elt)
        privilege.addChild(forwarded)
        wrapper.addChild(privilege)

        client.xmlstream.send(wrapper)

    async def send_privileged_iq(
        self, client: SatXMPPComponent, iq_elt: domish.Element
    ) -> domish.Element:
        """Send an IQ stanza on behalf of a user.

        @param client: The client sending the IQ.
        @param iq_elt: The IQ stanza to send (will be wrapped).
        @return: The IQ response.
        @raise PrivilegeError: If IQ privilege is not permitted.
        """
        data = client.get_plugin_data(IMPORT_NAME, ClientPrivilegeData)
        if not data.privileges.iq_permissions:
            raise PrivilegeError("IQ privilege not permitted")

        to_attr = iq_elt.getAttribute("to")
        if not to_attr:
            raise PrivilegeError("IQ must have a 'to' attribute")

        user_jid = jid.JID(to_attr)
        if user_jid.resource or jid.JID(user_jid.host) != client.server_jid:
            raise PrivilegeError(
                f"user_jid must be a bare JID from our server: {user_jid}"
            )

        iq_type = cast(str, iq_elt.getAttribute("type", ""))
        first_child = iq_elt.firstChildElement()
        if first_child is None:
            raise PrivilegeError("IQ has no payload")

        namespace = first_child.uri
        if not self.can_send_iq(client, namespace, iq_type):
            raise PrivilegeError(
                f"IQ privilege not permitted for namespace {namespace} type {iq_type}"
            )

        wrapper = client.IQ(iq_type)
        wrapper["to"] = user_jid.full()

        priv_iq = domish.Element((NS_PRIVILEGE, "privileged_iq"))
        iq_copy = domish.Element(("jabber:client", "iq"))
        iq_copy["type"] = iq_type
        iq_copy["id"] = iq_elt.getAttribute("id", str(uuid.uuid4()))
        if iq_elt.getAttribute("to"):
            iq_copy["to"] = iq_elt["to"]
        iq_copy.addChild(first_child)
        priv_iq.addChild(iq_copy)
        wrapper.addChild(priv_iq)

        result = await wrapper.send()
        return self._unwrap_privileged_iq_result(result)

    def _unwrap_privileged_iq_result(
        self, iq_result_elt: domish.Element
    ) -> domish.Element:
        """Unwrap a privileged IQ response.

        @param iq_result_elt: The IQ result stanza.
        @return: The unwrapped inner IQ.
        @raise PrivilegeError: If the response format is invalid.
        """
        try:
            privilege = next(iq_result_elt.elements(NS_PRIVILEGE, "privilege"))
            forwarded = next(privilege.elements(NS_FORWARD, "forwarded"))
            return next(forwarded.elements("jabber:client", "iq"))
        except StopIteration:
            log.error(f"Invalid privileged IQ response: {iq_result_elt.toXml()}")
            raise PrivilegeError("Invalid response format")

    def _on_presence(
        self, presence_elt: domish.Element, client: SatXMPPComponent
    ) -> None:
        defer.ensureDeferred(self.on_presence(presence_elt, client))

    async def on_presence(
        self, presence_elt: domish.Element, client: SatXMPPComponent
    ) -> None:
        """Handle incoming presence stanza.

        Capabilities are parsed with XEP-0115.

        @param presence_elt: The presence stanza.
        @param client: The client instance.
        """
        try:
            from_jid = jid.JID(presence_elt["from"])
        except (KeyError, RuntimeError):
            log.exception("Can't parse 'from' attribute.")
            raise

        self._track_presence(from_jid, presence_elt, client)

        try:
            await self._xep_0115.update(client, presence_elt)
        except Exception:
            log.exception("Failed to parse capabilities in XEP-0115")

        await G.host.trigger.async_point(
            "XEP-0356_presence_received", client, from_jid, presence_elt
        )

    def _track_presence(
        self, from_jid: jid.JID, presence_elt: domish.Element, client: SatXMPPComponent
    ) -> None:
        """Track online resources per bare JID.

        @param from_jid: The JID sending presence.
        @param presence_elt: The presence stanza.
        @param client: The client instance.
        """
        resource = from_jid.resource
        if resource is None:
            # Bare JID presence, ignore
            return

        bare_jid = from_jid.userhostJID()
        data = client.get_plugin_data(IMPORT_NAME, ClientPrivilegeData)
        online_resources = data.presence_data.online_resources

        presence_type = presence_elt.getAttribute("type")
        if presence_type == "unavailable":
            online_resources.setdefault(bare_jid, set()).discard(resource)
        else:
            online_resources.setdefault(bare_jid, set()).add(resource)


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

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

    @property
    def client(self) -> SatXMPPComponent:
        """Get the client (component) this handler is attached to."""
        assert self.parent is not None
        return cast(SatXMPPComponent, self.parent)

    def connectionInitialized(self) -> None:
        """Register observers after stream initialization.

        This is called immediately after authentication, ensuring we don't miss
        privilege advertisements that may be sent by the server right after
        the component authenticates.
        """
        if not self.client.is_component:
            return

        self.client.xmlstream.addObserver(
            f"/message/privilege[@xmlns='{NS_PRIVILEGE}']",
            self.plugin_parent.on_advertise,
            client=self.client,
        )
        self.client.xmlstream.addObserver(
            "/presence", self.plugin_parent._on_presence, client=self.client
        )
        log.debug(f"XEP-0356 observers registered for {self.client.profile}")

    def getDiscoInfo(
        self,
        requestor: jid.JID,
        target: jid.JID,
        nodeIdentifier: str = "",
    ) -> list[disco.DiscoFeature | disco.DiscoIdentity]:
        """Get disco info.

        @param requestor: JID of the requesting entity
        @param target: JID of the target entity
        @param nodeIdentifier: Node identifier
        @return: Disco features
        """
        return [disco.DiscoFeature(NS_PRIVILEGE)]

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

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