# Copyright (C) 2019 Philipp Hörist <philipp AT hoerist.com>
#
# This file is part of Gajim.
#
# SPDX-License-Identifier: GPL-3.0-only

# XEP-0384: OMEMO Encryption

from __future__ import annotations

from typing import Any
from typing import cast
from typing import Literal

import datetime as dt
from collections.abc import Generator
from dataclasses import dataclass
from pathlib import Path

import nbxmpp
from nbxmpp.errors import is_error
from nbxmpp.modules.omemo import create_omemo_message
from nbxmpp.modules.omemo import get_key_transport_message
from nbxmpp.namespaces import Namespace
from nbxmpp.protocol import JID
from nbxmpp.protocol import Message
from nbxmpp.protocol import NodeProcessed
from nbxmpp.structs import EncryptionData
from nbxmpp.structs import MessageProperties
from nbxmpp.structs import OMEMOMessage
from nbxmpp.structs import StanzaHandler
from omemo_dr.const import OMEMOTrust
from omemo_dr.exceptions import DecryptionFailed
from omemo_dr.exceptions import DuplicateMessage
from omemo_dr.exceptions import KeyExchangeMessage
from omemo_dr.exceptions import MessageNotForDevice
from omemo_dr.exceptions import SelfMessage
from omemo_dr.identitykey import IdentityKey
from omemo_dr.session_manager import OMEMOSessionManager
from omemo_dr.structs import IdentityInfo
from omemo_dr.structs import OMEMOBundle
from omemo_dr.structs import OMEMOConfig

from gajim.common import app
from gajim.common import configpaths
from gajim.common import ged
from gajim.common import types
from gajim.common.const import EncryptionInfoMsg
from gajim.common.const import Trust
from gajim.common.const import Trust as GajimTrust
from gajim.common.const import XmppUriQuery
from gajim.common.events import EncryptionInfo
from gajim.common.events import MucAdded
from gajim.common.events import MucDiscoUpdate
from gajim.common.events import SignedIn
from gajim.common.i18n import _
from gajim.common.modules.base import BaseModule
from gajim.common.modules.contacts import BareContact
from gajim.common.modules.contacts import GroupchatContact
from gajim.common.modules.contacts import GroupchatParticipant
from gajim.common.modules.util import as_task
from gajim.common.modules.util import CryptoModule
from gajim.common.modules.util import event_node
from gajim.common.modules.util import prepare_stanza
from gajim.common.modules.util import PublicKeyData
from gajim.common.storage.omemo import OMEMOStorage
from gajim.common.structs import OutgoingMessage
from gajim.common.util.decorators import cache_with_ttl
from gajim.common.util.decorators import event_filter
from gajim.common.util.text import format_fingerprint

ALLOWED_TAGS = [
    ("request", Namespace.RECEIPTS),
    ("active", Namespace.CHATSTATES),
    ("gone", Namespace.CHATSTATES),
    ("inactive", Namespace.CHATSTATES),
    ("paused", Namespace.CHATSTATES),
    ("composing", Namespace.CHATSTATES),
    ("markable", Namespace.CHATMARKERS),
    ("no-store", Namespace.HINTS),
    ("store", Namespace.HINTS),
    ("no-copy", Namespace.HINTS),
    ("no-permanent-store", Namespace.HINTS),
    ("replace", Namespace.CORRECT),
    ("thread", None),
    ("reply", Namespace.REPLY),
    ("fallback", Namespace.FALLBACK),
    ("origin-id", Namespace.SID),
    ("reactions", Namespace.REACTIONS),
]

DeviceIdT = int
IdentityT = tuple[DeviceIdT, IdentityKey]


@dataclass
class OMEMOPublicKeyData(PublicKeyData):
    device_id: int
    public_key: IdentityKey

    @classmethod
    def from_identity_info(cls, info: IdentityInfo) -> OMEMOPublicKeyData:
        last_seen = None
        if info.last_seen is not None:
            last_seen = dt.datetime.fromtimestamp(info.last_seen, dt.UTC)

        return cls(
            active=info.active,
            address=JID.from_string(info.address),
            label=info.label,
            last_seen=last_seen,
            fingerprint=info.public_key.get_fingerprint(),
            trust=GajimTrust[info.trust.name],
            device_id=info.device_id,
            public_key=info.public_key,
        )

    def pretty_fingerprint(self, wrap: bool = True) -> str:
        fingerprint = self.public_key.get_fingerprint()
        return format_fingerprint(fingerprint, "OMEMO", wrap=wrap)


class OMEMO(BaseModule, CryptoModule):
    _nbxmpp_extends = "OMEMO"
    _nbxmpp_methods = [
        "set_devicelist",
        "request_devicelist",
        "set_bundle",
        "request_bundle",
    ]

    def __init__(self, client: types.Client) -> None:
        BaseModule.__init__(self, client)

        self.handlers = [
            StanzaHandler(
                name="message",
                callback=self._message_received,
                ns=Namespace.OMEMO_TEMP,
                priority=9,
            ),
        ]

        self._register_pubsub_handler(self._devicelist_notification_received)
        self.register_events(
            [
                ("signed-in", ged.CORE, self._on_signed_in),
                ("muc-disco-update", ged.PREGUI, self._on_muc_disco_update),
                ("muc-added", ged.PREGUI, self._on_muc_added),
            ]
        )

        self.allow_groupchat = True

        self._own_jid = self._client.get_own_jid().bare

        data_dir = Path(configpaths.get("MY_DATA"))
        db_path = data_dir / f"omemo_{self._own_jid}.db"
        storage = OMEMOStorage(self._account, db_path, self._log)

        omemo_config = OMEMOConfig(
            default_prekey_amount=100,
            min_prekey_amount=80,
            spk_archive_seconds=86400 * 15,
            spk_cycle_seconds=86400 * 30,
            unacknowledged_count=2000,
        )

        self._backend = OMEMOSessionManager(
            self._own_jid, storage, omemo_config, self._account
        )
        self._backend.register_signal("republish-bundle", self._on_republish_bundle)

        self._omemo_groupchats: set[str] = set()
        self._muc_temp_store: dict[bytes, str] = {}

    @event_filter(["account"])
    def _on_signed_in(self, _event: SignedIn) -> None:
        self._log.info("Publish our bundle after sign in")
        self.set_bundle()
        self.request_devicelist()

    @event_filter(["account"])
    def _on_muc_disco_update(self, event: MucDiscoUpdate) -> None:
        before = self._is_omemo_groupchat(event.jid)
        self._check_if_omemo_capable(event.jid)
        after = self._is_omemo_groupchat(event.jid)

        if not before and after:
            self._set_group_members(event.jid)

    @event_filter(["account"])
    def _on_muc_added(self, event: MucAdded) -> None:
        client = app.get_client(event.account)
        contact = client.get_module("Contacts").get_contact(event.jid)
        if not isinstance(contact, GroupchatContact):
            self._log.warning("%s is not a groupchat contact", contact)
            return

        # Event is triggert on every join, avoid multiple connects
        contact.disconnect_all_from_obj(self)
        contact.connect("room-joined", self._on_room_joined)
        contact.connect("user-affiliation-changed", self._on_affiliation_change)
        contact.connect("room-affiliation-changed", self._on_affiliation_change)
        contact.connect("room-affiliations-complete", self._on_affiliation_change)

    def _on_room_joined(self, contact: GroupchatContact, _signal_name: str) -> None:
        self._check_if_omemo_capable(contact.jid)

    def _on_affiliation_change(
        self,
        contact: GroupchatContact,
        _signal_name: str,
        *args: Any,
    ) -> None:

        if not self._is_omemo_groupchat(contact.jid):
            return

        self._set_group_members(contact.jid)

    def _set_group_members(self, muc_jid: JID) -> None:
        self._log.info("Set group members: %s", muc_jid)
        affiliations = self._client.get_module("MUC").get_affiliations(muc_jid)
        members: set[JID] = set()
        members.update(*affiliations.values())
        str_jids = set(map(str, members))

        self._backend.set_group_members(str(muc_jid), str_jids)

        for jid in str_jids:
            if not self._is_contact_in_roster(jid):
                self._log.info("%s not in roster, query devicelist...", jid)
                self._request_device_list_ttl(jid)

    def _on_republish_bundle(
        self, _session: OMEMOSessionManager, _signal_name: str, bundle: OMEMOBundle
    ) -> None:

        self.set_bundle(bundle=bundle)

    @property
    def backend(self) -> OMEMOSessionManager:
        return self._backend

    def get_our_public_key(self) -> OMEMOPublicKeyData | None:
        device_id, identity_key = self._backend.get_our_identity()

        return OMEMOPublicKeyData(
            active=True,
            address=self._get_own_bare_jid(),
            label=None,
            last_seen=None,
            fingerprint=identity_key.get_fingerprint(),
            trust=Trust.VERIFIED,
            device_id=device_id,
            public_key=identity_key,
        )

    def get_public_keys(self, jid: JID, *, is_groupchat: bool) -> list[PublicKeyData]:
        identity_infos = self._backend.get_identity_infos(
            str(jid),
            only_active=is_groupchat,
            trust=OMEMOTrust.UNDECIDED if is_groupchat else None,
        )

        return [OMEMOPublicKeyData.from_identity_info(info) for info in identity_infos]

    def set_public_key_trust(
        self, public_key_data: PublicKeyData, trust: Trust
    ) -> None:
        public_key_data = cast(OMEMOPublicKeyData, public_key_data)
        self._backend.set_trust(
            str(public_key_data.address),
            public_key_data.public_key,
            OMEMOTrust[trust.name],
        )

    def remove_public_key(self, public_key_data: PublicKeyData) -> None:
        public_key_data = cast(OMEMOPublicKeyData, public_key_data)
        self._backend.delete_session(
            str(public_key_data.address),
            public_key_data.device_id,
            delete_identity=True,
        )

    def check_send_preconditions(self, contact: types.ChatContactT) -> bool:
        jid = str(contact.jid)
        if isinstance(contact, GroupchatContact):
            if not self._groupchat_pre_conditions_satisfied(contact):
                return False

        else:
            if not self._chat_pre_conditions_satisfied(contact):
                return False

        if self.backend.get_identity_infos(
            jid, only_active=True, trust=OMEMOTrust.UNDECIDED
        ):
            self._log.info("Undecided keys for %s", jid)
            app.ged.raise_event(
                EncryptionInfo(
                    account=contact.account,
                    jid=contact.jid,
                    type=EncryptionInfoMsg.UNDECIDED_FINGERPRINTS,
                    message=EncryptionInfoMsg.UNDECIDED_FINGERPRINTS.value,
                )
            )
            return False

        self._log.debug("Sending message to %s", jid)

        return True

    def _groupchat_pre_conditions_satisfied(
        self, contact: types.GroupchatContactT
    ) -> bool:

        if not self._is_omemo_groupchat(contact.jid):
            app.ged.raise_event(
                EncryptionInfo(
                    account=contact.account,
                    jid=contact.jid,
                    type=EncryptionInfoMsg.BAD_MUC_CONFIG,
                    message=EncryptionInfoMsg.BAD_MUC_CONFIG.value.format(
                        encryption="OMEMO"
                    ),
                )
            )
            return False

        jid = str(contact.jid)
        has_trusted_keys = False
        for member_jid in self.backend.get_group_members(jid, without_self=False):
            self._request_bundles_for_new_devices(member_jid)
            if self._has_trusted_keys(member_jid):
                has_trusted_keys = True

        if not has_trusted_keys:
            self._log.info("No trusted keys for %s", jid)
            app.ged.raise_event(
                EncryptionInfo(
                    account=contact.account,
                    jid=contact.jid,
                    type=EncryptionInfoMsg.NO_FINGERPRINTS,
                    message=EncryptionInfoMsg.NO_FINGERPRINTS.value,
                )
            )
            return False
        return True

    def _chat_pre_conditions_satisfied(self, contact: types.ChatContactT) -> bool:

        jid = str(contact.jid)
        if not self.backend.get_devices(jid, without_self=True):
            self.request_devicelist(jid)
            app.ged.raise_event(
                EncryptionInfo(
                    account=contact.account,
                    jid=contact.jid,
                    type=EncryptionInfoMsg.QUERY_DEVICES,
                    message=EncryptionInfoMsg.QUERY_DEVICES.value,
                )
            )
            return False

        if not self._has_trusted_keys(jid):
            self._log.info("No trusted keys for %s", jid)
            app.ged.raise_event(
                EncryptionInfo(
                    account=contact.account,
                    jid=contact.jid,
                    type=EncryptionInfoMsg.NO_FINGERPRINTS,
                    message=EncryptionInfoMsg.NO_FINGERPRINTS.value,
                )
            )
            return False
        return True

    def _is_omemo_groupchat(self, room_jid: JID) -> bool:
        return str(room_jid) in self._omemo_groupchats

    def encrypt_message(self, message: OutgoingMessage) -> bool:
        if not message.has_text():
            return False

        remote_jid = message.contact.jid

        client = app.get_client(self._account)
        contact = client.get_module("Contacts").get_contact(remote_jid)

        text = message.get_text()
        assert text is not None
        omemo_message = self.backend.encrypt(
            str(remote_jid), text, groupchat=contact.is_groupchat
        )
        if omemo_message is None:
            raise Exception("Encryption error")

        assert omemo_message.payload is not None

        create_omemo_message(
            message.get_stanza(),
            omemo_message,  # type: ignore
            node_whitelist=ALLOWED_TAGS,  # type: ignore
        )

        if message.is_groupchat:
            self._muc_temp_store[omemo_message.payload] = text

        message.set_encryption(
            EncryptionData(
                protocol="OMEMO",
                key="Unknown",
                trust=GajimTrust[OMEMOTrust.VERIFIED.name],
            )
        )

        self._debug_print_stanza(message.get_stanza())
        return True

    def _send_key_transport_message(
        self, typ: Literal["chat", "groupchat"], jid: str, devices: list[int]
    ) -> None:

        omemo_message = self.backend.encrypt_key_transport(jid, devices)
        if omemo_message is None:
            self._log.warning("Key transport message to %s (%s) failed", jid, devices)
            return

        transport_message = get_key_transport_message(typ, jid, omemo_message)  # type: ignore
        self._log.info("Send key transport message to %s (%s)", jid, devices)
        self._client.send_stanza(transport_message)

    def _message_received(
        self,
        _client: types.NBXMPPClient,
        stanza: Message,
        properties: MessageProperties,
    ) -> None:

        if properties.omemo is None:
            return

        if properties.carbon is not None and properties.carbon.is_sent:
            from_jid = self._own_jid

        elif properties.is_mam_message:
            from_jid = self._process_mam_message(properties)

        elif properties.from_muc:
            from_jid = self._process_muc_message(properties)

        else:
            assert properties.jid is not None
            from_jid = properties.jid.bare

        if from_jid is None:
            return

        self._log.info("Message received from: %s", from_jid)

        assert isinstance(properties.omemo, OMEMOMessage)
        try:
            plaintext, fingerprint, trust = self.backend.decrypt_message(
                properties.omemo,  # type: ignore
                from_jid,
            )
        except (KeyExchangeMessage, DuplicateMessage):
            raise NodeProcessed

        except SelfMessage:
            if not properties.from_muc:
                raise NodeProcessed

            if properties.omemo.payload not in self._muc_temp_store:
                self._log.warning("Can't decrypt own group chat message")
                return

            del self._muc_temp_store[properties.omemo.payload]
            return

        except DecryptionFailed:
            return

        except MessageNotForDevice:
            if properties.omemo.payload is None:
                # Key Transport message for another device
                return

            plaintext = _(
                "This message was encrypted with OMEMO, but not for your device."
            )
            # Neither trust nor fingerprint can be verified if we didn't
            # successfully decrypt the message
            trust = OMEMOTrust.UNTRUSTED
            fingerprint = None

        prepare_stanza(stanza, plaintext)
        self._debug_print_stanza(stanza)
        properties.encrypted = EncryptionData(
            protocol="OMEMO", key=fingerprint or "Unknown", trust=GajimTrust[trust.name]
        )

    def _process_muc_message(self, properties: MessageProperties) -> str | None:
        assert properties.jid is not None
        resource = properties.jid.resource
        contact = self._client.get_module("Contacts").get_contact(properties.jid)
        assert isinstance(contact, GroupchatParticipant)
        if contact.real_jid is not None:
            return contact.real_jid.bare

        self._log.error("Unable to find jid for group chat message from %s", resource)
        return None

    def _process_mam_message(self, properties: MessageProperties) -> str | None:
        assert properties.mam is not None
        assert properties.from_ is not None
        self._log.info("Message received, archive: %s", properties.mam.archive)
        if properties.from_muc:
            self._log.info("MUC MAM Message received")
            if properties.muc_user is None or properties.muc_user.jid is None:
                self._log.warning(
                    "Received MAM message which can not be mapped to a real jid"
                )
                return None
            return properties.muc_user.jid.bare
        return properties.from_.bare

    def _is_contact_in_roster(self, jid: str) -> bool:
        if jid == self._own_jid:
            return True

        roster_item = self._client.get_module("Roster").get_item(JID.from_string(jid))
        if roster_item is None:
            return False

        contact = self._client.get_module("Contacts").get_contact(jid)
        assert isinstance(contact, BareContact)
        return contact.subscription == "both"

    def _check_if_omemo_capable(self, jid: JID) -> None:
        disco_info = app.storage.cache.get_last_disco_info(jid)
        assert disco_info is not None
        if disco_info.muc_is_members_only and disco_info.muc_is_nonanonymous:
            self._log.info("OMEMO room discovered: %s", jid)
            self._omemo_groupchats.add(str(jid))
        else:
            self._log.info("OMEMO room removed due to config change: %s", jid)
            self._omemo_groupchats.discard(str(jid))

    def _request_bundles_for_new_devices(self, jid_: str) -> None:
        for jid in [jid_, self._own_jid]:
            device_ids = self.backend.get_devices_without_sessions(jid)
            for device_id in device_ids:
                self._request_bundle_ttl(jid, device_id)

    def _has_trusted_keys(self, jid: str) -> bool:
        infos = self.backend.get_identity_infos(
            jid, only_active=True, trust=[OMEMOTrust.VERIFIED, OMEMOTrust.BLIND]
        )
        return bool(infos)

    def set_bundle(self, bundle: OMEMOBundle | None = None) -> None:
        if bundle is None:
            bundle = self.backend.get_bundle(Namespace.OMEMO_TEMP)
        self._nbxmpp("OMEMO").set_bundle(bundle, self.backend.get_our_device())

    @cache_with_ttl(ttl=7200)
    def _request_bundle_ttl(self, jid: str, device_id: int) -> None:
        self.request_bundle(jid, device_id)

    @as_task
    def request_bundle(self, jid: str, device_id: int) -> Generator[Any, Any]:
        _task = yield  # noqa: F841

        self._log.info("Request device bundle %s %s", device_id, jid)

        bundle = yield self._nbxmpp("OMEMO").request_bundle(jid, device_id)

        if is_error(bundle) or bundle is None:
            self._log.info("Bundle request failed: %s %s: %s", jid, device_id, bundle)
            return

        try:
            self.backend.build_session(jid, bundle)
        except Exception as error:
            self._log.error("Building session failed: %s", error)
            return

        self._log.info("Session created for: %s", jid)
        # TODO: In MUC we should send a groupchat message
        self._send_key_transport_message("chat", jid, [device_id])

        # Trigger dialog to trust new Fingerprints if
        # the Chat Window is Open
        app.ged.raise_event(
            EncryptionInfo(
                account=self._account,
                jid=JID.from_string(jid),
                type=EncryptionInfoMsg.SESSION_BUILD,
                message=EncryptionInfoMsg.SESSION_BUILD.value.format(
                    device_id=device_id
                ),
            )
        )

    def set_devicelist(self, devicelist: list[int] | None = None) -> None:
        devicelist_: set[int] = {self.backend.get_our_device()}
        if devicelist is not None:
            devicelist_.update(devicelist)
        self._log.info("Publishing own devicelist: %s", devicelist_)
        self._nbxmpp("OMEMO").set_devicelist(devicelist_)

    def clear_keylist(self) -> None:
        self.backend.update_devicelist(self._own_jid, [self.backend.get_our_device()])
        self.set_devicelist()

    @cache_with_ttl(ttl=7200)
    def _request_device_list_ttl(self, jid: str) -> None:
        self.request_devicelist(jid)

    @as_task
    def request_devicelist(self, jid: str | None = None) -> Generator[Any, Any]:
        _task = yield  # noqa: F841

        if jid is None:
            jid = self._own_jid

        self._log.info("Request devicelist for %s", jid)

        devicelist = yield self._nbxmpp("OMEMO").request_devicelist(jid=jid)
        if is_error(devicelist) or devicelist is None:
            self._log.info("Devicelist request failed: %s %s", jid, devicelist)
            devicelist = []

        self._process_devicelist_update(jid, cast(list[int], devicelist))

    @event_node(Namespace.OMEMO_TEMP_DL)
    def _devicelist_notification_received(
        self,
        _client: types.NBXMPPClient,
        _stanza: Message,
        properties: MessageProperties,
    ) -> None:

        assert properties.pubsub_event is not None
        if properties.pubsub_event.retracted:
            return

        devicelist = cast(list[int], properties.pubsub_event.data or [])

        self._process_devicelist_update(str(properties.jid), devicelist)

    def _process_devicelist_update(self, jid: str, devicelist: list[int]) -> None:
        own_devices = self._client.get_own_jid().bare_match(jid)
        if own_devices:
            jid = self._own_jid

        self._log.info("Received device list for %s: %s", jid, devicelist)
        # Pass a copy, we need the full list for potential set_devicelist()
        self.backend.update_devicelist(jid, list(devicelist))

        if own_devices:
            if not self.backend.is_our_device_published():
                # Our own device_id is not in the list, it could be
                # overwritten by some other client
                self.set_devicelist(devicelist)

        self._request_bundles_for_new_devices(jid)

    def _debug_print_stanza(self, stanza: nbxmpp.Node) -> None:
        stanzastr = "\n" + stanza.__str__(fancy=True)
        stanzastr = stanzastr[0:-1]
        self._log.debug(stanzastr)

    def compose_trust_uri(self, jid: JID) -> str:
        verified_identities = [
            (info.device_id, info.public_key)
            for info in self._backend.get_identity_infos(
                jid.bare, only_active=True, trust=OMEMOTrust.VERIFIED
            )
        ]
        if self._client.is_own_jid(jid):
            verified_identities.insert(0, self._backend.get_our_identity())

        query = (
            (
                XmppUriQuery.MESSAGE.value,
                [
                    (f"omemo-sid-{sid}", ik.get_fingerprint())
                    for sid, ik in verified_identities
                ],
            )
            if verified_identities
            else None
        )
        uri = jid.new_as_bare().to_iri(query)
        return uri

    def cleanup(self) -> None:
        BaseModule.cleanup(self)
        self._backend.destroy()
        del self._backend
