#!/usr/bin/env python3

# 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/>.

import time
from typing import cast
import uuid

from sqlalchemy.orm.attributes import flag_modified
from twisted.internet import defer
from twisted.words.protocols.jabber import xmlstream
from twisted.words.protocols.jabber import jid
from twisted.words.xish import domish
from tx_xmpp import disco
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.memory.sqla import History, Message, Subject, joinedload, select
from libervia.backend.models.core import MessageData, MessageEdition
from libervia.backend.plugins.plugin_xep_0045 import XEP_0045
from libervia.backend.plugins.plugin_xep_0334 import XEP_0334
from libervia.backend.tools.common import data_format
from libervia.backend.tools.utils import aio

log = getLogger(__name__)


PLUGIN_INFO = {
    C.PI_NAME: "Last Message Correction",
    C.PI_IMPORT_NAME: "XEP-0308",
    C.PI_TYPE: "XEP",
    C.PI_PROTOCOLS: ["XEP-0308"],
    C.PI_DEPENDENCIES: ["XEP-0045", "XEP-0334"],
    C.PI_MAIN: "XEP_0308",
    C.PI_HANDLER: "yes",
    C.PI_DESCRIPTION: _("""Implementation of XEP-0308 (Last Message Correction)"""),
}

NS_MESSAGE_CORRECT = "urn:xmpp:message-correct:0"


class XEP_0308:
    def __init__(self, host):
        log.info(f"plugin {PLUGIN_INFO[C.PI_NAME]!r} initialization")
        self.host = host
        host.register_namespace("message_correct", NS_MESSAGE_CORRECT)
        host.trigger.add("message_received", self._message_received_trigger)
        host.bridge.add_method(
            "message_edit",
            ".plugin",
            in_sign="sss",
            out_sign="",
            method=self._message_edit,
            async_=True,
        )
        self._m = cast(XEP_0045, host.plugins["XEP-0045"])
        self._h = cast(XEP_0334, host.plugins["XEP-0334"])

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

    @aio
    async def get_last_history(
        self, client: SatXMPPEntity, message_elt: domish.Element
    ) -> History | None:
        profile_id = self.host.memory.storage.profiles[client.profile]
        from_jid = jid.JID(message_elt["from"])
        message_type = message_elt.getAttribute("type", C.MESS_TYPE_NORMAL)
        async with self.host.memory.storage.session() as session:
            stmt = (
                select(History)
                .where(
                    History.profile_id == profile_id,
                    History.source == from_jid.userhost(),
                    History.dest == client.jid.userhost(),
                    History.type == message_type,
                )
                .options(joinedload(History.messages))
                .options(joinedload(History.subjects))
                .options(joinedload(History.thread))
            )
            if message_elt.type == C.MESS_TYPE_GROUPCHAT:
                stmt = stmt.where(History.source_res == from_jid.resource)

            # we want last message
            stmt = stmt.order_by(History.timestamp.desc()).limit(1)
            result = await session.execute(stmt)
            history = result.unique().scalar_one_or_none()
        return history

    async def update_history(
        self,
        client: SatXMPPEntity,
        edited_history: History,
        edit_timestamp: float,
        new_message: dict[str, str],
        new_subject: dict[str, str],
        new_extra: dict,
        previous_message: dict[str, str],
        previous_subject: dict[str, str],
        previous_extra: dict | None,
        store: bool = True,
    ) -> None:
        # FIXME: new_extra is not handled by now
        edited_history.messages = [
            Message(message=mess, language=lang) for lang, mess in new_message.items()
        ]
        edited_history.subjects = [
            Subject(subject=mess, language=lang) for lang, mess in new_subject.items()
        ]
        previous_version = {
            # this is the timestamp when this version was published
            "timestamp": edited_history.extra.get("updated", edited_history.timestamp),
            "message": previous_message,
            "subject": previous_subject,
        }
        edited_history.extra["updated"] = edit_timestamp
        if previous_extra:
            previous_extra = previous_extra.copy()
            # we must not have editions in each edition
            try:
                del previous_extra[C.MESS_EXTRA_EDITIONS]
            except KeyError:
                pass
            # extra may be important for rich content
            previous_version["extra"] = previous_extra

        if store:
            flag_modified(edited_history, "extra")
            edited_history.extra.setdefault(C.MESS_EXTRA_EDITIONS, []).append(
                previous_version
            )
            await self.host.memory.storage.add(edited_history)

            edit_data = MessageData(edited_history.serialise())
            self.host.bridge.message_update(
                edited_history.uid,
                C.MESS_UPDATE_EDIT,
                data_format.serialise(edit_data),
                client.profile,
            )

    async def _message_received_trigger(
        self,
        client: SatXMPPEntity,
        message_elt: domish.Element,
        post_treat: defer.Deferred,
    ) -> bool:
        from_jid = jid.JID(message_elt["from"])
        replace_elt = next(message_elt.elements(NS_MESSAGE_CORRECT, "replace"), None)
        if not replace_elt:
            return True
        replace_id = (replace_elt.getAttribute("id") or "").strip()
        if not replace_id:
            log.warning(f"Invalid message correction: {message_elt.toXml()}")
            return False
        # look for the corrected message by id (this also handles our own messages, e.g.
        # corrections from another device)
        message_type = message_elt.getAttribute("type", C.MESS_TYPE_NORMAL)
        edited_history = await self.host.memory.storage.get_history_from_xmpp_id(
            client, replace_id, message_type
        )
        if edited_history is None:
            # Fallback: last message from the sender (history without usable id)
            edited_history = await self.get_last_history(client, message_elt)
            if edited_history is None:
                log.warning(
                    f"No message found to correct ({replace_id}), ignoring "
                    f"correction: {message_elt.toXml()}"
                )
                return False
            check_id = edited_history.origin_id or edited_history.extra.get("message_id")
            if check_id != replace_id:
                log.warning(
                    "Can't apply correction: it doesn't reference the last one:\n"
                    f"{message_elt.toXml()}"
                )
                return False
        # For groupchat and MUC private messages the full JID must match.
        # For a plain 1:1 message, the resource may have changed (e.g. reconnection), so
        # only the bare JID is compared
        is_muc = edited_history.type == C.MESS_TYPE_GROUPCHAT or self._m.is_joined_room(
            client, jid.JID(edited_history.source)
        )

        if is_muc:
            same_author = edited_history.source_jid == from_jid
        else:
            same_author = edited_history.source_jid.userhost() == from_jid.userhost()
        if not same_author:
            log.warning(
                "Can't apply correction: correction doesn't come from same author "
                f"{edited_history.source_jid.full()}:\n{message_elt.toXml()}"
            )
            return False

        previous_message_data = edited_history.serialise()
        message_data = await client.messageProt.parse_message(message_elt)
        if not message_data["message"] and not message_data["subject"]:
            log.warning(
                "Message correction doesn't have body nor subject, we can't edit anything"
            )
            return False
        if (
            message_data["message"] == previous_message_data["message"]
            and message_data["subject"] == previous_message_data["subject"]
        ):
            # content is already up to date (e.g. our own correction redelivered by MAM
            # after reconnection)
            log.debug(f"Ignoring already applied correction: {message_elt.toXml()}")
            return False

        await self.update_history(
            client,
            edited_history,
            message_data.get("received_timestamp") or message_data["timestamp"],
            message_data["message"],
            message_data["subject"],
            message_data["extra"],
            previous_message_data["message"],
            previous_message_data["subject"],
            previous_message_data.get("extra"),
        )
        return False

    async def message_edit(
        self,
        client: SatXMPPEntity,
        message_id: str,
        edit_data: MessageEdition,
    ) -> None:
        """Edit a message

        The message can only be edited if it's the last one of the discussion.
        @param client: client instance
        @param message_id: UID of the message to edit
        @param edit_data: data to update in the message
        """
        timestamp = time.time()
        edited_history = await self.host.memory.storage.get(
            client,
            History,
            History.uid,
            message_id,
            joined_loads=[History.messages, History.subjects, History.thread],
        )
        if edited_history is None:
            raise exceptions.NotFound(
                f"message to edit not found in database ({message_id})"
            )
        if edited_history.type == C.MESS_TYPE_GROUPCHAT:
            is_group_chat = True
            # In the case of group chat, our message is sent by ourself with our room JID.
            peer_jid = edited_history.source_jid
        else:
            is_group_chat = False
            # keep the destination resource (e.g. room/nick for a MUC private message)
            peer_jid = edited_history.dest_jid
        history_data = await self.host.memory.history_get(
            client.jid, peer_jid, limit=1, profile=client.profile
        )
        if not history_data:
            raise exceptions.NotFound(
                "No message found in conversation with {peer_jid.userhost()}"
            )
        last_mess = history_data[0]
        if last_mess[0] != message_id:
            raise ValueError(
                f"{message_id} is not the last message of the discussion, we can't edit "
                "it"
            )

        ref_id = edited_history.stanza_id if is_group_chat else edited_history.origin_id
        if not ref_id:
            orig_data = edited_history.serialise()
            orig_data["from"] = jid.JID(orig_data["from"])
            orig_data["to"] = jid.JID(orig_data["to"])
            log.warning(
                f"Can't correct message {message_id}: the original message has no usable "
                f"id (stanza_id={edited_history.stanza_id}, "
                f"origin_id={edited_history.origin_id}), ignoring the edit."
            )
            return

        await self.update_history(
            client,
            edited_history,
            timestamp,
            edit_data.message,
            edit_data.subject,
            edit_data.extra,
            last_mess[4],
            last_mess[5],
            last_mess[-1],
            # message will be updated and signal sent on reception in group chat
            store=not is_group_chat,
        )

        serialised = edited_history.serialise()
        serialised["uid"] = str(uuid.uuid4())
        serialised["from"] = client.jid
        # for group chat, we want to send the correction to the room
        serialised["to"] = peer_jid.userhostJID() if is_group_chat else peer_jid

        message_elt = client.generate_message_xml(serialised)["xml"]
        replace_elt = message_elt.addElement((NS_MESSAGE_CORRECT, "replace"))
        replace_elt["id"] = ref_id
        self._h.add_hint_elements(message_elt, [self._h.HINT_STORE])
        client.send(message_elt)

    def _message_edit(self, message_id: str, edit_data_s: str, profile: str) -> None:
        client = self.host.get_client(profile)
        edit_data = MessageEdition.model_validate_json(edit_data_s)
        defer.ensureDeferred(self.message_edit(client, message_id, edit_data))


@implementer(disco.IDisco)
class XEP_0308_handler(xmlstream.XMPPHandler):
    def getDiscoInfo(self, __, target, nodeIdentifier=""):
        return [disco.DiscoFeature(NS_MESSAGE_CORRECT)]

    def getDiscoItems(self, requestor, target, nodeIdentifier=""):
        return []