#!/usr/bin/env python3

# Libervia plugin for Chat State Notifications Protocol (XEP-0085)
# Copyright (C) 2009-2026 Jérôme Poisson (goffi@goffi.org)
# Copyright (C) 2009-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 dataclasses import dataclass, field
from typing import Any, TYPE_CHECKING, cast

from twisted.internet import defer
from twisted.internet import reactor
from twisted.internet import error as internet_error
from twisted.internet.base import DelayedCall
from twisted.internet.interfaces import IReactorTime
from twisted.words.protocols.jabber.jid 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

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


log = getLogger(__name__)


NS_XMPP_CLIENT = "jabber:client"
NS_CHAT_STATES = "http://jabber.org/protocol/chatstates"
CHAT_STATES = ["active", "inactive", "gone", "composing", "paused"]
MESSAGE_TYPES = ["chat", "groupchat"]
PARAM_KEY = "Notifications"
PARAM_NAME = "Enable chat state notifications"
ENTITY_KEY = PARAM_KEY + "_" + PARAM_NAME
DELETE_VALUE = "DELETE"

PLUGIN_INFO = {
    C.PI_NAME: "Chat State Notifications Protocol Plugin",
    C.PI_IMPORT_NAME: "XEP-0085",
    C.PI_TYPE: "XEP",
    C.PI_PROTOCOLS: ["XEP-0085"],
    C.PI_DEPENDENCIES: [],
    C.PI_MAIN: "XEP_0085",
    C.PI_HANDLER: "yes",
    C.PI_DESCRIPTION: _("""Implementation of Chat State Notifications Protocol"""),
}

IMPORT_NAME = PLUGIN_INFO[C.PI_IMPORT_NAME]


@dataclass
class ChatStatesData:
    """Per-client data for XEP-0085 chat state notifications."""

    # peer_jid -> Allowed (True)/Denied (False)
    disco_cache: dict[JID, bool] = field(default_factory=dict)
    # Chat state machines, keyed by peer_jid
    machines: dict[JID, "ChatStateMachine"] = field(default_factory=dict)


# Describe the internal transitions that are triggered
# by a timer. Beside that, external transitions can be
# runned to target the states "active" or "composing".
# Delay is specified here in seconds.
TRANSITIONS = {
    "active": {"next_state": "inactive", "delay": 120},
    "inactive": {"next_state": "gone", "delay": 480},
    "gone": {"next_state": "", "delay": 0},
    "composing": {"next_state": "paused", "delay": 30},
    "paused": {"next_state": "inactive", "delay": 450},
}


class UnknownChatStateException(Exception):
    """
    This error is raised when an unknown chat state is used.
    """

    pass


class XEP_0085:
    """
    Implementation for XEP 0085
    """

    params = """
    <params>
    <individual>
    <category name="%(category_name)s" label="%(category_label)s">
        <param name="%(param_name)s" label="%(param_label)s" value="true" type="bool" security="0"/>
     </category>
    </individual>
    </params>
    """ % {
        "category_name": PARAM_KEY,
        "category_label": _(PARAM_KEY),
        "param_name": PARAM_NAME,
        "param_label": _("Enable chat state notifications"),
    }

    def __init__(self, host: "LiberviaBackend") -> None:
        log.info(_("Chat State Notifications plugin initialization"))
        self.host = host
        # parameter value is retrieved before each use
        host.memory.update_params(self.params)

        # triggers from core
        host.trigger.add("message_received", self.message_received_trigger)
        host.trigger.add("sendMessage", self.send_message_trigger)
        host.trigger.add("param_update_trigger", self.param_update_trigger)

        # args: to_s (jid as string), profile
        host.bridge.add_method(
            "chat_state_composing",
            ".plugin",
            in_sign="ss",
            out_sign="",
            method=self.chat_state_composing,
        )

        # args: from (jid as string), state in CHAT_STATES, profile
        host.bridge.add_signal("chat_state_received", ".plugin", signature="sss")

    def get_handler(self, client: SatXMPPEntity) -> "XEP_0085_handler":
        return XEP_0085_handler(self, client.profile)

    def profile_disconnected(self, client: SatXMPPEntity) -> None:
        """Eventually send a 'gone' state to all one2one contacts."""
        chat_states_data = client.get_plugin_data(IMPORT_NAME, ChatStatesData)
        for machine in list(chat_states_data.machines.values()):
            # FIXME: the "unavailable" presence stanza is received by to_jid
            # before the chat state, so it will be ignored... find a way to
            # actually defer the disconnection
            machine._onEvent("gone")
        chat_states_data.machines.clear()

    async def _check_disco_support(
        self, client: SatXMPPEntity, peer_jid: JID
    ) -> bool | None:
        """Check if a contact supports XEP-0085 via service discovery.

        For 1:1 (bare JID): checks all known resources.
        For MUC DM (full JID): checks the occupant directly.

        @param client: the client instance
        @param peer_jid: bare JID for 1:1, full JID for MUC DM
        @return: True if confirmed, False if denied, None if no resources known
        """
        data = client.get_plugin_data(IMPORT_NAME, ChatStatesData)
        if peer_jid in data.disco_cache:
            return data.disco_cache[peer_jid]

        if not peer_jid.resource:
            # 1:1 (we check all known resources)
            try:
                resources = self.host.memory.get_all_resources(client, peer_jid)
            except exceptions.UnknownEntityError:
                # entity not in cache
                return None
            if not resources:
                # no resources known yet
                return None
            for resource in resources:
                full_jid = JID(None, (peer_jid.user, peer_jid.host, resource))
                try:
                    has_feature = await self.host.memory.disco.has_feature(
                        client, NS_CHAT_STATES, full_jid
                    )
                    if has_feature:
                        data.disco_cache[peer_jid] = True
                        return True
                except (
                    exceptions.ExternalRequestError,
                    exceptions.TimeOutError,
                ) as e:
                    log.warning(
                        f"Can't check chat state availability for {full_jid}: {e}"
                    )
                    continue
            data.disco_cache[peer_jid] = False
            return False

        # We have a full JID -> This is a MUC DM
        try:
            has_feature = await self.host.memory.disco.has_feature(
                client, NS_CHAT_STATES, peer_jid
            )
        except (
            exceptions.ExternalRequestError,
            exceptions.TimeOutError,
        ):
            log.debug(
                f"Disco query for chat states support failed for {peer_jid.full()}, "
                "falling back to implicit discovery"
            )
            has_feature = False
        data.disco_cache[peer_jid] = has_feature
        return has_feature

    async def _ensure_disco_check(self, client: SatXMPPEntity, to_jid: JID) -> None:
        """Check disco support for chat states if not already done.

        Skipped for MUC groupchat.
        For 1:1: passes bare JID, checks all known resources.
        For MUC DM: passes full JID, checks occupant.

        @param client: the client instance
        @param to_jid: JID of the contact
        """
        if self._is_muc(to_jid, client.profile) and not to_jid.resource:
            # This is a MUC groupchat, we skip
            return
        data = client.get_plugin_data(IMPORT_NAME, ChatStatesData)
        if to_jid not in data.disco_cache:
            await self._check_disco_support(client, to_jid)

    def update_cache(self, entity_jid: JID, value: bool | str, profile: str) -> None:
        """Update the entity data of the given profile for one or all contacts.
        Reset the chat state(s) display if the notification has been disabled.

        @param entity_jid: contact's JID
        @param value: True, False or DELETE_VALUE to delete the entity data
        @param profile: current profile
        """
        client = self.host.get_client(profile)
        if value == DELETE_VALUE:
            self.host.memory.del_entity_datum(client, entity_jid, ENTITY_KEY)
        else:
            self.host.memory.update_entity_data(client, entity_jid, ENTITY_KEY, value)
        if not value or value == DELETE_VALUE:
            # reinit chat state UI for this contact
            self.host.bridge.chat_state_received(entity_jid.full(), "", profile)

    def param_update_trigger(
        self, name: str, value: str, category: str, type_: str, profile: str
    ) -> bool:
        """Update Chat State data according to parameter.

        @param name: Parameter name.
        @param value: "true" to activate the notifications, or any other value to delete
            it.
        @param category: Parameter category.
        @param type_: Parameter type.
        """
        if (category, name) == (PARAM_KEY, PARAM_NAME):
            client = self.host.get_client(profile)
            if C.bool(value):
                self.host.memory.update_entity_data(
                    client, C.ENTITY_ALL, ENTITY_KEY, True
                )
            else:
                self.host.memory.del_entity_datum(client, C.ENTITY_ALL, ENTITY_KEY)
            return False
        return True

    async def message_received_trigger(
        self, client: SatXMPPEntity, message: domish.Element, post_treat: bool
    ) -> bool:
        """Handle an incoming message.

        Update entity cache, init state machine, and signal frontends about
        received chat states.
        """
        profile = client.profile
        if not self.host.memory.param_get_a(PARAM_NAME, PARAM_KEY, profile_key=profile):
            return True

        from_jid = JID(message.getAttribute("from"))
        msg_type = message.getAttribute("type")
        # MUC: for groupchat we remove resource to get room bare JID.
        #      for DM (type="chat") we keep full JID to track occupant.
        is_groupchat = self._is_muc(from_jid, profile) and msg_type == "groupchat"
        if is_groupchat:
            from_jid = from_jid.userhostJID()

        # Extract any chat state element from the message (for all message types)
        chat_state = next(
            (
                child.name
                for child in message.elements()
                if child.name in CHAT_STATES and child.defaultUri == NS_CHAT_STATES
            ),
            None,
        )

        if not is_groupchat:
            # one2one or MUC DM: implicit discovery (XEP-0085 §5.1)
            # check if this is a content message (has <body/>)
            has_body = (
                next(
                    domish.generateElementsNamed(message.elements(), name="body"),
                    None,
                )
                is not None
            )

            if chat_state is not None:
                # XEP-0085 §5.1 Rule 3: any chat state proves support
                self.update_cache(from_jid, True, profile=profile)
            elif has_body and message.getAttribute("type") == "chat":
                # XEP-0085 §5.1 Rule 2: content message without chat state
                # only deactivate if not already activated
                try:
                    current = self.host.memory.get_entity_datum(
                        client, from_jid, ENTITY_KEY
                    )
                except (exceptions.UnknownEntityError, KeyError):
                    current = None
                if not current:
                    self.update_cache(from_jid, False, profile=profile)
                    return True

            # fire explicit disco check for this contact
            await self._ensure_disco_check(client, from_jid)

        # init state machine so composing states can be sent to this contact
        # XEP-0085 §5.4: only use with chat/groupchat message types
        if msg_type in MESSAGE_TYPES:
            self._chat_state_init(client, from_jid, msg_type)

        # Signal frontends about received chat states
        if chat_state is not None:
            # XEP-0085 §5.5: ignore <gone/> in groupchat
            if chat_state != "gone" or message.getAttribute("type") != "groupchat":
                self.host.bridge.chat_state_received(
                    message.getAttribute("from"), chat_state, profile
                )
        return True

    def send_message_trigger(
        self,
        client: SatXMPPEntity,
        mess_data: dict,
        pre_xml_treatments: "defer.Deferred",
        post_xml_treatments: "defer.Deferred",
    ) -> bool:
        """Add chat state to outgoing messages and initiate the state machine."""
        profile = client.profile

        async def treatment(mess_data):
            message = mess_data["xml"]
            to_jid = JID(message.getAttribute("to"))

            # XEP-0085 §4: fire disco check for this contact
            await self._ensure_disco_check(client, to_jid)

            if not await self._check_activation(client, to_jid, forceEntityData=True):
                return mess_data

            # XEP-0085 §2: check if this is a content message (has <body/>)
            has_body = (
                next(
                    domish.generateElementsNamed(message.elements(), name="body"),
                    None,
                )
                is not None
            )

            if has_body and mess_data["type"] in MESSAGE_TYPES:
                # XEP-0085 §5.4: only use with chat/groupchat message types
                # XEP-0085 §5.6: content message -> attach <active/>
                message.addElement((NS_CHAT_STATES, "active"))
                # Launch the chat state machine (init the timer)
                # MUC groupchat: strip resource to get room bare JID.
                # MUC DM (type="chat"): keep full JID to track occupant.
                target_jid = (
                    to_jid.userhostJID()
                    if self._is_muc(to_jid, profile) and mess_data["type"] == "groupchat"
                    else to_jid
                )
                self._chat_state_active(client, target_jid, mess_data["type"])
            elif "chat_state" in mess_data["extra"]:
                # Standalone notification
                state = mess_data["extra"].pop("chat_state")
                assert state in CHAT_STATES
                message.addElement((NS_CHAT_STATES, state))
            return mess_data

        post_xml_treatments.addCallback(
            lambda _: defer.ensureDeferred(treatment(mess_data))
        )
        return True

    def _is_muc(self, to_jid: JID, profile: str) -> bool:
        """Tell if that JID is a MUC or not

        @param to_jid: full or bare JID to check
        @param profile: %(doc_profile)s
        @return: bool
        """
        client = self.host.get_client(profile)
        try:
            type_ = self.host.memory.get_entity_datum(
                client, to_jid.userhostJID(), C.ENTITY_TYPE
            )
            if type_ == C.ENTITY_TYPE_MUC:
                return True
        except (exceptions.UnknownEntityError, KeyError):
            pass
        return False

    async def _check_activation(
        self, client: SatXMPPEntity, to_jid: JID, forceEntityData: bool
    ) -> bool:
        """Check whether chat state notifications should be sent to a contact.

        Activation follows a priority order:
        1. User config
        2. MUC
        3. Disco cache (preferred)
        4. Implicit discovery cache (fallback)

        @param client: the client instance
        @param to_jid: the contact's full JID (or bare if you know it's a MUC)
        @param forceEntityData: if True, a missing implicit cache entry is
            treated as True and initialized (used for initial discovery request)
        @return: True if notifications should be sent to this JID.
        """
        profile = client.profile

        # XEP-0085 §5.2: user must be able to disable globally
        if not self.host.memory.param_get_a(PARAM_NAME, PARAM_KEY, profile_key=profile):
            return False

        # XEP-0085 §5.5: MUC groupchat, no discovery required.
        # MUC DMs (to_jid has resource) go through normal checks.
        if self._is_muc(to_jid, profile) and not to_jid.resource:
            return True

        # XEP-0085 §4: explicit disco discovery (preferred over implicit)
        disco_result = await self._check_disco_support(client, to_jid)
        if disco_result is True:
            # disco confirmed support, activate unconditionally
            self.update_cache(to_jid, True, profile=profile)
            return True
        if disco_result is False:
            # disco confirmed no support, deactivate
            return False
        # disco_result is None: no resources known, fall through to implicit discovery

        # XEP-0085 §5.1: implicit discovery (fallback)
        try:
            return self.host.memory.get_entity_datum(client, to_jid, ENTITY_KEY)
        except (exceptions.UnknownEntityError, KeyError):
            if forceEntityData:
                # XEP-0085 §5.1 Rule 1: enable for initial discovery request
                self.update_cache(to_jid, True, profile=profile)
                return True
            # no discovery completed yet, wait for first message
            return False

    def _chat_state_init(
        self, client: SatXMPPEntity, to_jid: JID, mess_type: str
    ) -> None:
        """Data initialization for the chat state machine.

        @param client: the client instance
        @param to_jid: full JID for one2one, bare JID for MUC
        @param mess_type: "chat" or "groupchat"
        """
        chat_states_data = client.get_plugin_data(IMPORT_NAME, ChatStatesData)
        machines = chat_states_data.machines
        if to_jid not in machines:
            machines[to_jid] = ChatStateMachine(
                self.host, to_jid, mess_type, client.profile
            )

    def _chat_state_active(
        self, client: SatXMPPEntity, to_jid: JID, mess_type: str
    ) -> None:
        """Launch the chat state machine on "active" state.

        @param client: the client instance
        @param to_jid: full JID for one2one, bare JID for MUC
        @param mess_type: "chat" or "groupchat"
        """
        self._chat_state_init(client, to_jid, mess_type)
        chat_states_data = client.get_plugin_data(IMPORT_NAME, ChatStatesData)
        chat_states_data.machines[to_jid]._onEvent("active")

    async def chat_state_composing(self, to_jid_s: str, profile_key: str) -> None:
        """Move to the "composing" state when required.

        Called from the frontend. Checks user config, disco support,
        and implicit discovery state before sending.

        @param to_jid_s: contact full JID as a string
        @param profile_key: %(doc_profile_key)s
        """
        client = self.host.get_client(profile_key)
        to_jid = JID(to_jid_s)
        profile = client.profile

        # XEP-0085 §4: fire disco check if not already done
        await self._ensure_disco_check(client, to_jid)

        if not await self._check_activation(client, to_jid, forceEntityData=False):
            return

        # Ensure state machine exists (may not be initialised if no message
        # has been exchanged yet but disco confirmed support)
        # groupchat (bare JID) vs MUC DM (full JID with resource)
        is_groupchat = self._is_muc(to_jid, profile) and not to_jid.resource
        self._chat_state_init(client, to_jid, "groupchat" if is_groupchat else "chat")
        chat_states_data = client.get_plugin_data(IMPORT_NAME, ChatStatesData)
        chat_states_data.machines[to_jid]._onEvent("composing")


class ChatStateMachine:
    """
    This class represents a chat state, between one profile and
    one target contact. A timer is used to move from one state
    to the other. The initialization is done through the "active"
    state which is internally set when a message is sent. The state
    "composing" can be set externally (through the bridge by a
    frontend). Other states are automatically set with the timer.
    """

    def __init__(
        self, host: "LiberviaBackend", to_jid: JID, mess_type: str, profile: str
    ) -> None:
        """
        Initialization need to store the target, message type
        and a profile for sending later messages.
        """
        self.host = host
        self.to_jid = to_jid
        self.mess_type = mess_type
        self.profile = profile
        self.state = None
        self.timer: DelayedCall | None = None

    def _onEvent(self, state: str) -> None:
        """
        Move to the specified state, eventually send the
        notification to the contact (the "active" state is
        automatically sent with each message) and set the timer.
        """
        assert state in TRANSITIONS
        transition = TRANSITIONS[state]
        assert "next_state" in transition and "delay" in transition

        if state != self.state and state != "active":
            if state != "gone" or self.mess_type != "groupchat":
                # send a new message without body
                log.debug(
                    "sending state '{state}' to {jid}".format(
                        state=state, jid=self.to_jid.full()
                    )
                )
                client = self.host.get_client(self.profile)
                mess_data = {
                    "from": client.jid,
                    "to": self.to_jid,
                    "uid": "",
                    "message": {},
                    "type": self.mess_type,
                    "subject": {},
                    "extra": {},
                }
                client.generate_message_xml(cast(Any, mess_data))
                message_elt = cast(domish.Element, mess_data["xml"])
                message_elt.addElement((NS_CHAT_STATES, state))
                client.send(message_elt)

        self.state = state
        if self.timer is not None:
            try:
                self.timer.cancel()
            except internet_error.AlreadyCalled:
                pass

        if transition["next_state"] and transition["delay"] > 0:
            reactor_time = cast(IReactorTime, reactor)
            self.timer = reactor_time.callLater(
                transition["delay"], self._onEvent, transition["next_state"]
            )


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

    def __init__(self, plugin_parent: "XEP_0085", profile: str) -> None:
        self.plugin_parent = plugin_parent
        self.host = plugin_parent.host
        self.profile = profile

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

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