#!/usr/bin/env python3

# Libervia plugin for Extended Channel Search (XEP-0461)
# 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, Any, Final, Self, cast
from pydantic import BaseModel, Field
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 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 History
from libervia.backend.models.core import MessageData
from libervia.backend.models.types import JIDType
from libervia.backend.tools.utils import ensure_deferred

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

log = getLogger(__name__)


PLUGIN_INFO = {
    C.PI_NAME: "Message Replies",
    C.PI_IMPORT_NAME: "XEP-0461",
    C.PI_TYPE: "XEP",
    C.PI_MODES: C.PLUG_MODE_BOTH,
    C.PI_DEPENDENCIES: [],
    C.PI_RECOMMENDATIONS: [],
    C.PI_MAIN: "XEP_0461",
    C.PI_HANDLER: "yes",
    C.PI_DESCRIPTION: _("Message replies support."),
}

NS_REPLY = "urn:xmpp:reply:0"
host: "LiberviaBackend | None" = None


class ReplyTo(BaseModel):
    to: JIDType | None = None
    id: str
    internal_uid: bool = Field(
        default=True,
        description="True if the id is a message UID, i.e. Libervia internal ID of "
        "message, otherwise it's an XMPP stanza ID or origin ID as specified in the XEP",
    )

    async def ensure_xmpp_id(self, client: SatXMPPEntity) -> None:
        if not self.internal_uid:
            return
        history = await G.storage.get(
            client,
            History,
            History.uid,
            self.id,
        )

        if history.uid != self.id:
            raise exceptions.InternalError(
                f"Inconsistency between given history UID {history.uid!r} and ReplyTo "
                f"id ({self.id!r})."
            )
        if history.type == C.MESS_TYPE_GROUPCHAT:
            self.id = history.stanza_id
        else:
            self.id = history.origin_id
        assert self.id
        self.internal_uid = False

    @classmethod
    def from_element(cls, reply_elt: domish.Element) -> Self:
        """Create a ReplyTo instance from a <reply> element or its parent.

        @param reply_elt: The <reply> element or a parent element.
        @return: ReplyTo instance.
        @raise exceptions.NotFound: If the <reply> element is not found.
        """
        if reply_elt.uri != NS_REPLY or reply_elt.name != "reply":
            child_file_metadata_elt = next(reply_elt.elements(NS_REPLY, "reply"), None)
            if child_file_metadata_elt is None:
                raise exceptions.NotFound("<reply> element not found")
            else:
                reply_elt = child_file_metadata_elt

        try:
            message_id = reply_elt["id"]
        except KeyError:
            raise exceptions.DataError('Missing "id" attribute.')

        return cls(to=reply_elt.getAttribute("to"), id=message_id, internal_uid=False)

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

        @return: <reply> element.
        """
        if self.internal_uid:
            raise exceptions.DataError(
                '"id" must be converted to XMPP id before calling "to_element". Please '
                '"use ensure_xmpp_id" first.'
            )
        reply_elt = domish.Element((NS_REPLY, "reply"), attribs={"id": self.id})
        if self.to is not None:
            reply_elt["to"] = self.to.full()
        return reply_elt


class XEP_0461:
    """Implementation of XEP-0461 Message Replies."""

    namespace: Final[str] = NS_REPLY

    def __init__(self, host: Any):
        log.info(f"Plugin {PLUGIN_INFO[C.PI_NAME]!r} initialization.")
        self.host = host
        host.register_namespace("reply", NS_REPLY)
        host.trigger.add("message_received", self._message_received_trigger)
        host.trigger.add("sendMessage", self._sendMessage_trigger)

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

    async def _parse_reply(
        self, client: SatXMPPEntity, message_elt: domish.Element, mess_data: MessageData
    ) -> MessageData:
        try:
            reply_to = ReplyTo.from_element(message_elt)
        except exceptions.NotFound:
            pass
        else:
            try:
                thread_id = mess_data["extra"]["thread"]
            except KeyError:
                pass
            else:
                message_type = message_elt.getAttribute("type")
                storage = G.storage
                history = await storage.get_history_from_xmpp_id(
                    client, reply_to.id, message_type
                )
                if history is None:
                    log.warning(f"Received a <reply> to an unknown message: {reply_to!r}")
                else:
                    if not history.thread:
                        await storage.add_thread(
                            history.uid, thread_id, None, is_retroactive=True
                        )
                    # We can to use internal UID to make frontends life easier.
                    reply_to.id = history.uid
                    reply_to.internal_uid = True
            # FIXME: We use mode=json to have a serialisable data in storage. When
            #  Pydantic models are fully integrated with MessageData and storage, the
            #  ReplyTo model should be used directly here instead.
            mess_data["extra"]["reply"] = reply_to.model_dump(mode="json")
        return mess_data

    def _message_received_trigger(
        self,
        client: SatXMPPEntity,
        message_elt: domish.Element,
        post_treat: defer.Deferred,
    ) -> bool:
        post_treat.addCallback(
            lambda mess_data: defer.ensureDeferred(
                self._parse_reply(client, message_elt, mess_data)
            )
        )
        return True

    @ensure_deferred
    async def _add_reply_to_elt(
        self, mess_data: MessageData, client: SatXMPPEntity
    ) -> MessageData:
        try:
            # FIXME: Once MessageData is fully moved to Pydantic, we should have directly
            #   the ReplyTo instance here.
            reply_to = ReplyTo(**mess_data["extra"]["reply"])
        except KeyError:
            log.error('"_add_reply_to_elt" should not be called when there is no reply.')
        else:
            await reply_to.ensure_xmpp_id(client)
            message_elt: domish.Element = mess_data["xml"]
            message_elt.addChild(reply_to.to_element())

        return mess_data

    def _sendMessage_trigger(
        self,
        client: SatXMPPEntity,
        mess_data: MessageData,
        pre_xml_treatments: defer.Deferred,
        post_xml_treatments: defer.Deferred,
    ) -> bool:
        try:
            extra = mess_data["extra"]
            reply_to = ReplyTo(**extra["reply"])
        except KeyError:
            pass
        else:
            if "thread" not in extra:
                # FIXME: once MessageData is fully moved to Pydantic, ReplyTo should be
                #   used directly here instead of a dump.
                extra["reply"] = reply_to.model_dump(mode="json")
                # We use parent message ID as thread ID.
                extra["thread"] = reply_to.id

            post_xml_treatments.addCallback(self._add_reply_to_elt, client)
        return True


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

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

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