#!/usr/bin/env python3

# Libervia plugin for Extended Stanza Addressing (XEP-0033)
# Copyright (C) 2009-2026 Jérôme Poisson (goffi@goffi.org)
# Copyright (C) 2013-2016 Adrien Cossa (souliane@mailoo.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 Iterator, Literal, Self

from pydantic import BaseModel, model_validator
from twisted.internet import defer
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 disco, itx_xmpp
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 SatXMPPEntity
from libervia.backend.core.i18n import _
from libervia.backend.core.log import getLogger
from libervia.backend.models.core import MessageData
from libervia.backend.models.types import JIDType
from libervia.backend.tools import trigger
from libervia.backend.tools.xml_tools import element_copy

log = getLogger(__name__)


# TODO: fix Prosody "addressing" plugin to leave the concerned bcc according to the spec:
#   http://xmpp.org/extensions/xep-0033.html#addr-type-bcc "This means that the server
#   MUST remove these addresses before the stanza is delivered to anyone other than the
#   given bcc addressee or the multicast service of the bcc addressee."
#
#   http://xmpp.org/extensions/xep-0033.html#multicast "Each 'bcc' recipient MUST receive
#   only the <address type='bcc'/> associated with that addressee."

# TODO: fix Prosody "addressing" plugin to determine itself if remote servers supports
#   this XEP


PLUGIN_INFO = {
    C.PI_NAME: "Extended Stanza Addressing Protocol Plugin",
    C.PI_IMPORT_NAME: "XEP-0033",
    C.PI_TYPE: "XEP",
    C.PI_MODES: C.PLUG_MODE_BOTH,
    C.PI_PROTOCOLS: ["XEP-0033"],
    C.PI_DEPENDENCIES: [],
    C.PI_MAIN: "XEP_0033",
    C.PI_HANDLER: "yes",
    C.PI_DESCRIPTION: _(
        "Efficiently send messages to several recipients, using metadata to transmit "
        "them with main recipients (to), carbon copies (cc), and blind carbon copies "
        "(bcc) fields in a similar manner as for email."
    ),
}

NS_ADDRESS = "http://jabber.org/protocol/address"
RECIPIENT_FIELDS = ("to", "cc", "bcc")


class AddressType(BaseModel):
    jid: JIDType | None = None
    desc: str | None = None
    delivered: bool | None = None

    def set_attribute(self, address_elt: domish.Element) -> None:
        """Set <address> element attribute from this instance's data."""
        if self.jid:
            address_elt["jid"] = str(self.jid)
        if self.desc:
            address_elt["desc"] = self.desc
        if self.delivered is not None:
            address_elt["delivered"] = "true" if self.delivered else "false"

    @classmethod
    def from_element(cls, address_elt: domish.Element) -> Self:
        """Create an AddressType instance from an <address> element.

        @param address_elt: The <address> element.
        @return: AddressType instance.
        """
        if address_elt.uri != NS_ADDRESS or address_elt.name != "address":
            raise ValueError("Element is not an <address> element")

        kwargs = {}
        if address_elt.hasAttribute("jid"):
            kwargs["jid"] = jid.JID(address_elt["jid"])
        if address_elt.hasAttribute("desc"):
            kwargs["desc"] = address_elt["desc"]
        if address_elt.hasAttribute("delivered"):
            kwargs["delivered"] = address_elt["delivered"] == "true"
        return cls(**kwargs)

    def to_element(self) -> domish.Element:
        """Build the <address> element from this instance's data.

        @return: <address> element.
        """
        address_elt = domish.Element((NS_ADDRESS, "address"))
        self.set_attribute(address_elt)
        return address_elt


class AddressesData(BaseModel):
    to: list[AddressType] | None = None
    cc: list[AddressType] | None = None
    bcc: list[AddressType] | None = None
    replyto: list[AddressType] | None = None
    replyroom: list[AddressType] | None = None
    noreply: bool | None = None
    ofrom: JIDType | None = None

    @model_validator(mode="after")
    def check_minimal_data(self) -> Self:
        has_recipient = self.to or self.cc or self.bcc
        has_metadata = self.replyto or self.replyroom or self.noreply or self.ofrom
        assert has_recipient or has_metadata, "At least one address field must be set"
        if self.noreply and (self.replyto is not None or self.replyroom is not None):
            log.warning(
                '"noreply" can\'t be used with "replyto" or "replyroom". Ignoring reply '
                f"fields ({self.replyto=}, {self.replyroom=})."
            )
            # We reset instead of raising a ValueError, because this can happen in
            # incoming messages and we should not discard them.
            self.replyto = self.replyroom = None
        return self

    @property
    def addresses(self) -> Iterator[AddressType]:
        """Iterator over all recipient addresses."""
        for field in RECIPIENT_FIELDS:
            addresses = getattr(self, field)
            if not addresses:
                continue
            yield from addresses

    @staticmethod
    def add_address_element(
        addresses_elt: domish.Element, type_: str, address: AddressType | None
    ) -> None:
        """Add <address> element to parent <addresses> element.

        @param addresses_elt: Parent <addresses> element.
        @param type_: Value of "type" attribute.
        @param address: Address data.
        """

        address_elt = addresses_elt.addElement("address")
        address_elt["type"] = type_
        if address is not None:
            address.set_attribute(address_elt)

    @classmethod
    def from_element(cls, addresses_elt: domish.Element) -> Self:
        """Create an AddressesData instance from an <addresses> element.

        @param addresses_elt: The <addresses> element or its direct parent.
        @return: AddressesData instance.
        @raise NotFound: No <addresses> element found.
        """
        if addresses_elt.uri != NS_ADDRESS or addresses_elt.name != "addresses":
            child_addresses_elt = next(
                addresses_elt.elements(NS_ADDRESS, "addresses"), None
            )
            if child_addresses_elt is None:
                raise exceptions.NotFound("<addresses> element not found")
            else:
                addresses_elt = child_addresses_elt

        kwargs = {}
        for address_elt in addresses_elt.elements(NS_ADDRESS, "address"):
            address_type = address_elt.getAttribute("type")
            if address_type in ("to", "cc", "bcc", "replyto", "replyroom"):
                try:
                    address = AddressType.from_element(address_elt)
                except Exception as e:
                    log.warning(f"Invalid <address> element: {e}\n{address_elt.toXml()}")
                else:
                    kwargs.setdefault(address_type, []).append(address)
            elif address_type == "noreply":
                kwargs["noreply"] = True
            elif address_type == "ofrom":
                kwargs["ofrom"] = jid.JID(address_elt["jid"])
            else:
                log.warning(
                    f"Invalid <address> element: unknonwn type {address_type!r}\n"
                    f"{address_elt.toXml()}"
                )
        return cls(**kwargs)

    def to_element(self) -> domish.Element:
        """Build the <addresses> element from this instance's data.

        @return: <addresses> element.
        """
        addresses_elt = domish.Element((NS_ADDRESS, "addresses"))

        if self.to:
            for address in self.to:
                self.add_address_element(addresses_elt, "to", address)
        if self.cc:
            for address in self.cc:
                self.add_address_element(addresses_elt, "cc", address)
        if self.bcc:
            for address in self.bcc:
                self.add_address_element(addresses_elt, "bcc", address)
        if self.replyto:
            for address in self.replyto:
                self.add_address_element(addresses_elt, "replyto", address)
        if self.replyroom:
            for address in self.replyroom:
                self.add_address_element(addresses_elt, "replyroom", address)
        if self.noreply:
            self.add_address_element(addresses_elt, "noreply", None)
        if self.ofrom is not None:
            address_elt = addresses_elt.addElement("address")
            address_elt["type"] = "ofrom"
            address_elt["jid"] = self.ofrom.full()

        return addresses_elt


class XEP_0033:
    """
    Implementation for XEP-0033
    """

    def __init__(self, host):
        log.info(_("Extended Stanza Addressing plugin initialization"))
        self.host = host
        host.register_namespace("address", NS_ADDRESS)
        self.internal_data = {}
        host.trigger.add(
            "sendMessage",
            self.send_message_trigger,
            # We want this trigger to be the last one, as it may send messages.
            trigger.TriggerManager.MIN_PRIORITY,
        )
        host.trigger.add(
            "sendMessageComponent",
            self.send_message_trigger,
            # We want this trigger to be the last one, as it may send messages.
            trigger.TriggerManager.MIN_PRIORITY,
        )
        host.trigger.add("message_received", self.message_received_trigger)

    async def _stop_if_all_delivered(
        self, client: SatXMPPEntity, mess_data: MessageData, addr_data: AddressesData
    ) -> None:
        """Check if all message have been delivered, and stop workflow in this case.

        If workflow is stopped, message will be added to history and a signal will be sent
        to bridge.
        @param client: Client session.
        @param mess_data: Message data.
        @param addr_data: Addresses data.

        @raise exceptions.CancelError: All message have been delivered and workflow is
            terminated.
        """
        if all(a.delivered for a in addr_data.addresses):
            await client.message_add_to_history(mess_data)
            await client.message_send_to_bridge(mess_data)
            raise exceptions.CancelError(
                f"Message has been delivered by {PLUGIN_INFO['C.PI_NAME']}."
            )

    async def _handle_addresses(self, client, mess_data: MessageData) -> MessageData:
        """Handle Extended Stanza Addressing metadata for outgoing messages."""
        if not "addresses" in mess_data["extra"]:
            return mess_data

        if mess_data["extra"].get(C.MESS_KEY_ENCRYPTED, False):
            # TODO: Message must be encrypted for all recipients, and "to" correspond to
            #   multicast service in this case.
            raise NotImplementedError(
                "End-to-end encryption is not supported yet with multicast addressing."
            )

        data = AddressesData(**mess_data["extra"]["addresses"])
        recipients = set()
        domains: dict[str, list[AddressType]] = {}
        for address in data.addresses:
            if address.jid is None:
                raise NotImplementedError("Non JID addresses are not supported yet.")
            recipients.add(address.jid)
            try:
                domains[address.jid.host].append(address)
            except KeyError:
                domains[address.jid.host] = [address]

        to_recipient_jid = mess_data["to"]

        if to_recipient_jid.user and to_recipient_jid not in recipients:
            # If the main recipient is not a service (i.e. it has a "user" part), we want
            # to move it to the XEP-0033's "to" addresses, so we can use the multicast
            # service for <message> "to" attribute.
            to_recipient_addr = AddressType(jid=to_recipient_jid)
            if data.to is None:
                data.to = [to_recipient_addr]
            else:
                data.to.insert(0, to_recipient_addr)
            recipients.add(to_recipient_jid)
            domains.setdefault(to_recipient_jid.host, []).append(to_recipient_addr)

        # XXX: If our server doesn't handle multicast, we don't check sub-services as
        #   requested in §2.2, because except if there is a special arrangement with the
        #   server, a service at a sub-domain can't send message in the name of the main
        #   domain (e.g. "multicast.example.org" can't send message from
        #   "juliet@example.org"). So the specification is a bit dubious here, and we only
        #   use the main server multicast feature if it's present.
        if not await self.host.memory.disco.has_feature(
            client, NS_ADDRESS, client.server_jid
        ):
            # No multicast service
            log.warning(
                _(
                    f"Server of {client.profile} does not support XEP-0033 "
                    f"({PLUGIN_INFO[C.PI_IMPORT_NAME]}). We will send all messages ourselves."
                )
            )
            await self.deliver_messages(client, mess_data, data, domains)
            await self._stop_if_all_delivered(client, mess_data, data)
        else:
            # XXX: We delived ourself to multicast services because it's not correctly
            #     handled by some multicast services, notably by Prosody mod_addresses.
            # FIXME: Only do this workaround for known incomplete implementations.
            # TODO: remove this workaround when known implementations have been completed.
            if mess_data["to"] != client.server_jid:
                # We send the message to our server which will distribute it to the right
                # locations. The initial ``to`` jid has been moved to ``data.to`` above.
                # FIXME: When sub-services issue is properly handler, a sub-service JID
                #     supporting multicast should be allowed here.
                mess_data["to"] = client.server_jid
            await self.deliver_messages(
                client, mess_data, data, domains, multicast_only=True
            )
            await self._stop_if_all_delivered(client, mess_data, data)

        message_elt = mess_data["xml"]
        message_elt["to"] = str(mess_data["to"])
        message_elt.addChild(data.to_element())
        return mess_data

    async def deliver_messages(
        self,
        client,
        mess_data: MessageData,
        addr_data: AddressesData,
        domains: dict[str, list[AddressType]],
        multicast_only: bool = False,
    ) -> None:
        """Send messages to requested recipients.

        If a domain handles multicast, a single message will be send there.
        @param client: Client session.
        @param mess_data: Messsa data.
        @param addr_data: XEP-0033 addresses data.
        @param domains: Domain to addresses map.
            Note that that the addresses instances in this argument must be the same as in
            ``addr_data`` (There ``delivered`` status will be manipulated).
        @param multicast_only: if True, only multicast domain will be delivered.
        """
        # We'll modify delivered status, so we keep track here of addresses which have
        # already be delivered.
        already_delivered = [a for a in addr_data.addresses if a.delivered]
        multicast_domains = set()
        for domain, domain_addresses in domains.items():
            if domain == client.server_jid.host:

                # ``client.server_jid`` is discarded to avoid sending twice the same
                # message. ``multicast_only`` flag is set when the server supports
                # multicast, so the message will be sent to it at the end of the workflow.
                continue
            if len(domain_addresses) > 1:
                # For domains with multiple recipients, we check if we they support
                # multicast and so if we can deliver to them directly.
                if await self.host.memory.disco.has_feature(
                    client, NS_ADDRESS, jid.JID(domain)
                ):
                    multicast_domains.add(domain)

        # We remove bcc, they have a special handling.
        bcc = addr_data.bcc or []
        addr_data.bcc = None

        # Mark all addresses as "delivered" upfront, even if some won't actually be sent
        # by us (when multicast_only is set). This flag signals to multicast services that
        # they shouldn't handle these addresses. We'll remove the "delivered" status from
        # undelivered addresses post-delivery.
        for address in addr_data.addresses:
            address.delivered = True

        # First, we send multicast messages.
        for domain in multicast_domains:
            something_to_deliver = False
            for address in domains[domain]:
                if address in already_delivered:
                    continue
                # We need to mark as non delivered, so the multicast service will deliver
                # itself.
                address.delivered = False
                something_to_deliver = True

            if not something_to_deliver:
                continue

            domain_bcc = [a for a in bcc if a.jid and a.jid.host == domain]
            message_elt = element_copy(mess_data["xml"])
            # The service must only see BCC from its own domain.
            addr_data.bcc = domain_bcc
            message_elt.addChild(addr_data.to_element())
            message_elt["to"] = domain
            await client.a_send(message_elt)
            for address in domains[domain] + domain_bcc:
                # Those addresses have now been delivered.
                address.delivered = True

        if multicast_only:
            # Only addresses from multicast domains must be marked as delivered.
            for address in addr_data.addresses:
                if (
                    address.jid is not None
                    and address.jid.host not in multicast_domains
                    and address not in already_delivered
                ):
                    address.delivered = None

            # We have delivered to all multicast services, we stop here.
            # But first we need to restore BCC, without the delivered ones.
            addr_data.bcc = [a for a in bcc if not a.delivered]
            return

        # Then BCC
        for address in bcc:
            if address in already_delivered:
                continue
            if address.jid is None:
                raise NotImplementedError(
                    f"Sending to non JID address is not supported yet"
                )
            if address.jid.host in multicast_domains:
                # Address has already be handled by a multicast domain
                continue
            message_elt = element_copy(mess_data["xml"])
            # The recipient must only get its own BCC
            addr_data.bcc = [address]
            message_elt.addChild(addr_data.to_element())
            message_elt["to"] = address.jid.full()
            await client.a_send(message_elt)

        # BCC address must be removed.
        addr_data.bcc = None

        # and finally, other ones.
        message_elt = mess_data["xml"]
        message_elt.addChild(addr_data.to_element())
        non_bcc_addresses = (addr_data.to or []) + (addr_data.cc or [])
        for address in non_bcc_addresses:
            if address in already_delivered:
                continue
            if address.jid is None:
                raise NotImplementedError(
                    f"Sending to non JID address is not supported yet"
                )
            if address.jid.host in multicast_domains:
                # Multicast domains have already been delivered.
                continue
            message_elt["to"] = address.jid.full()
            await client.a_send(message_elt)

    def send_message_trigger(
        self, client, mess_data, pre_xml_treatments, post_xml_treatments
    ) -> Literal[True]:
        """Process the XEP-0033 related data to be sent"""
        post_xml_treatments.addCallback(
            lambda mess_data: defer.ensureDeferred(
                self._handle_addresses(client, mess_data)
            )
        )
        return True

    def message_received_trigger(
        self,
        client: SatXMPPEntity,
        message_elt: domish.Element,
        post_treat: defer.Deferred,
    ) -> Literal[True]:
        """Parse addresses information and add them to message data."""

        try:
            addresses = AddressesData.from_element(message_elt)
        except exceptions.NotFound:
            pass
        else:

            def post_treat_addr(mess_data: MessageData):
                mess_data["extra"]["addresses"] = addresses.model_dump(
                    mode="json", exclude_none=True
                )
                return mess_data

            post_treat.addCallback(post_treat_addr)
        return True

    def get_handler(self, client):
        return XEP_0033_handler()


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

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

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