# Copyright (c) 2022 Tulir Asokan
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
from __future__ import annotations

import asyncio
import logging
import sys

from mautrix import __optional_imports__
from mautrix.appservice import AppService
from mautrix.client import Client, InternalEventType, SyncStore
from mautrix.crypto import CryptoStore, OlmMachine, PgCryptoStore, RejectKeyShare, StateStore
from mautrix.errors import EncryptionError, MForbidden, MNotFound, SessionNotFound
from mautrix.types import (
    JSON,
    DeviceIdentity,
    EncryptedEvent,
    EncryptedMegolmEventContent,
    EventFilter,
    EventType,
    Filter,
    LoginType,
    MessageEvent,
    RequestedKeyInfo,
    RoomEventFilter,
    RoomFilter,
    RoomID,
    RoomKeyWithheldCode,
    Serializable,
    StateEvent,
    StateFilter,
    TrustState,
)
from mautrix.util import background_task
from mautrix.util.async_db import Database
from mautrix.util.logging import TraceLogger

from .. import bridge as br
from .crypto_state_store import PgCryptoStateStore


class EncryptionManager:
    loop: asyncio.AbstractEventLoop
    log: TraceLogger = logging.getLogger("mau.bridge.e2ee")

    client: Client
    crypto: OlmMachine
    crypto_store: CryptoStore | SyncStore
    crypto_db: Database | None
    state_store: StateStore

    min_send_trust: TrustState
    key_sharing_enabled: bool
    appservice_mode: bool
    periodically_delete_expired_keys: bool
    delete_outdated_inbound: bool

    bridge: br.Bridge
    az: AppService
    _id_prefix: str
    _id_suffix: str

    _share_session_events: dict[RoomID, asyncio.Event]
    _key_delete_task: asyncio.Task | None

    def __init__(
        self,
        bridge: br.Bridge,
        homeserver_address: str,
        user_id_prefix: str,
        user_id_suffix: str,
        db_url: str,
    ) -> None:
        self.loop = bridge.loop or asyncio.get_event_loop()
        self.bridge = bridge
        self.az = bridge.az
        self.device_name = bridge.name
        self._id_prefix = user_id_prefix
        self._id_suffix = user_id_suffix
        self._share_session_events = {}
        pickle_key = "mautrix.bridge.e2ee"
        self.crypto_db = Database.create(
            url=db_url,
            upgrade_table=PgCryptoStore.upgrade_table,
            log=logging.getLogger("mau.crypto.db"),
        )
        self.crypto_store = PgCryptoStore("", pickle_key, self.crypto_db)
        self.state_store = PgCryptoStateStore(self.crypto_db, bridge.get_portal)
        default_http_retry_count = bridge.config.get("homeserver.http_retry_count", None)
        self.client = Client(
            base_url=homeserver_address,
            mxid=self.az.bot_mxid,
            loop=self.loop,
            sync_store=self.crypto_store,
            log=self.log.getChild("client"),
            default_retry_count=default_http_retry_count,
            state_store=self.bridge.state_store,
        )
        self.crypto = OlmMachine(self.client, self.crypto_store, self.state_store)
        self.client.add_event_handler(InternalEventType.SYNC_STOPPED, self._exit_on_sync_fail)
        self.crypto.allow_key_share = self.allow_key_share
        verification_levels = bridge.config["bridge.encryption.verification_levels"]
        self.min_send_trust = TrustState.parse(verification_levels["send"])
        self.crypto.share_keys_min_trust = TrustState.parse(verification_levels["share"])
        self.crypto.send_keys_min_trust = TrustState.parse(verification_levels["receive"])
        self.key_sharing_enabled = bridge.config["bridge.encryption.allow_key_sharing"]
        self.appservice_mode = bridge.config["bridge.encryption.appservice"]
        if self.appservice_mode:
            self.az.otk_handler = self.crypto.handle_as_otk_counts
            self.az.device_list_handler = self.crypto.handle_as_device_lists
            self.az.to_device_handler = self.crypto.handle_as_to_device_event

        self.periodically_delete_expired_keys = False
        self.delete_outdated_inbound = False
        self._key_delete_task = None
        del_cfg = bridge.config["bridge.encryption.delete_keys"]
        if del_cfg:
            self.crypto.delete_outbound_keys_on_ack = del_cfg["delete_outbound_on_ack"]
            self.crypto.dont_store_outbound_keys = del_cfg["dont_store_outbound"]
            self.crypto.delete_previous_keys_on_receive = del_cfg["delete_prev_on_new_session"]
            self.crypto.ratchet_keys_on_decrypt = del_cfg["ratchet_on_decrypt"]
            self.crypto.delete_fully_used_keys_on_decrypt = del_cfg["delete_fully_used_on_decrypt"]
            self.crypto.delete_keys_on_device_delete = del_cfg["delete_on_device_delete"]
            self.periodically_delete_expired_keys = del_cfg["periodically_delete_expired"]
            self.delete_outdated_inbound = del_cfg["delete_outdated_inbound"]
        self.crypto.disable_device_change_key_rotation = bridge.config[
            "bridge.encryption.rotation.disable_device_change_key_rotation"
        ]

    async def _exit_on_sync_fail(self, data) -> None:
        if data["error"]:
            self.log.critical("Exiting due to crypto sync error")
            sys.exit(32)

    async def allow_key_share(self, device: DeviceIdentity, request: RequestedKeyInfo) -> bool:
        if not self.key_sharing_enabled:
            self.log.debug(
                f"Key sharing not enabled, ignoring key request from "
                f"{device.user_id}/{device.device_id}"
            )
            return False
        elif device.trust == TrustState.BLACKLISTED:
            raise RejectKeyShare(
                f"Rejecting key request from blacklisted device "
                f"{device.user_id}/{device.device_id}",
                code=RoomKeyWithheldCode.BLACKLISTED,
                reason="Your device has been blacklisted by the bridge",
            )
        elif await self.crypto.resolve_trust(device) >= self.crypto.share_keys_min_trust:
            portal = await self.bridge.get_portal(request.room_id)
            if portal is None:
                raise RejectKeyShare(
                    f"Rejecting key request for {request.session_id} from "
                    f"{device.user_id}/{device.device_id}: room is not a portal",
                    code=RoomKeyWithheldCode.UNAVAILABLE,
                    reason="Requested room is not a portal",
                )
            user = await self.bridge.get_user(device.user_id)
            if not await user.is_in_portal(portal):
                raise RejectKeyShare(
                    f"Rejecting key request for {request.session_id} from "
                    f"{device.user_id}/{device.device_id}: user is not in portal",
                    code=RoomKeyWithheldCode.UNAUTHORIZED,
                    reason="You're not in that portal",
                )
            self.log.debug(
                f"Accepting key request for {request.session_id} from "
                f"{device.user_id}/{device.device_id}"
            )
            return True
        else:
            raise RejectKeyShare(
                f"Rejecting key request from unverified device "
                f"{device.user_id}/{device.device_id}",
                code=RoomKeyWithheldCode.UNVERIFIED,
                reason="Your device is not trusted by the bridge",
            )

    def _ignore_user(self, user_id: str) -> bool:
        return (
            user_id.startswith(self._id_prefix)
            and user_id.endswith(self._id_suffix)
            and user_id != self.az.bot_mxid
        )

    async def handle_member_event(self, evt: StateEvent) -> None:
        if self._ignore_user(evt.state_key):
            # We don't want to invalidate group sessions because a ghost left or joined
            return
        await self.crypto.handle_member_event(evt)

    async def _share_session_lock(self, room_id: RoomID) -> bool:
        try:
            event = self._share_session_events[room_id]
        except KeyError:
            self._share_session_events[room_id] = asyncio.Event()
            return True
        else:
            await event.wait()
            return False

    async def encrypt(
        self, room_id: RoomID, event_type: EventType, content: Serializable | JSON
    ) -> tuple[EventType, EncryptedMegolmEventContent]:
        try:
            encrypted = await self.crypto.encrypt_megolm_event(room_id, event_type, content)
        except EncryptionError:
            self.log.debug("Got EncryptionError, sharing group session and trying again")
            if await self._share_session_lock(room_id):
                try:
                    users = await self.az.state_store.get_members_filtered(
                        room_id, self._id_prefix, self._id_suffix, self.az.bot_mxid
                    )
                    await self.crypto.share_group_session(room_id, users)
                finally:
                    self._share_session_events.pop(room_id).set()
            encrypted = await self.crypto.encrypt_megolm_event(room_id, event_type, content)
        return EventType.ROOM_ENCRYPTED, encrypted

    async def decrypt(self, evt: EncryptedEvent, wait_session_timeout: int = 5) -> MessageEvent:
        try:
            decrypted = await self.crypto.decrypt_megolm_event(evt)
        except SessionNotFound as e:
            if not wait_session_timeout:
                raise
            self.log.debug(
                f"Couldn't find session {e.session_id} trying to decrypt {evt.event_id},"
                f" waiting {wait_session_timeout} seconds..."
            )
            got_keys = await self.crypto.wait_for_session(
                evt.room_id, e.session_id, timeout=wait_session_timeout
            )
            if got_keys:
                self.log.debug(
                    f"Got session {e.session_id} after waiting, "
                    f"trying to decrypt {evt.event_id} again"
                )
                decrypted = await self.crypto.decrypt_megolm_event(evt)
            else:
                raise
        self.log.trace("Decrypted event %s: %s", evt.event_id, decrypted)
        return decrypted

    async def start(self) -> None:
        flows = await self.client.get_login_flows()
        if not flows.supports_type(LoginType.APPSERVICE):
            self.log.critical(
                "Encryption enabled in config, but homeserver does not support appservice login"
            )
            sys.exit(30)
        self.log.debug("Logging in with bridge bot user")
        if self.crypto_db:
            try:
                await self.crypto_db.start()
            except Exception as e:
                self.bridge._log_db_error(e)
        await self.crypto_store.open()
        device_id = await self.crypto_store.get_device_id()
        if device_id:
            self.log.debug(f"Found device ID in database: {device_id}")
        # We set the API token to the AS token here to authenticate the appservice login
        # It'll get overridden after the login
        self.client.api.token = self.az.as_token
        await self.client.login(
            login_type=LoginType.APPSERVICE,
            device_name=self.device_name,
            device_id=device_id,
            store_access_token=True,
            update_hs_url=False,
        )
        await self.crypto.load()
        if not device_id:
            await self.crypto_store.put_device_id(self.client.device_id)
            self.log.debug(f"Logged in with new device ID {self.client.device_id}")
        elif self.crypto.account.shared:
            await self._verify_keys_are_on_server()
        if self.appservice_mode:
            self.log.info("End-to-bridge encryption support is enabled (appservice mode)")
        else:
            _ = self.client.start(self._filter)
            self.log.info("End-to-bridge encryption support is enabled (sync mode)")
        if self.delete_outdated_inbound:
            deleted = await self.crypto_store.redact_outdated_group_sessions()
            if len(deleted) > 0:
                self.log.debug(
                    f"Deleted {len(deleted)} inbound keys which lacked expiration metadata"
                )
        if self.periodically_delete_expired_keys:
            self._key_delete_task = background_task.create(self._periodically_delete_keys())
        background_task.create(self._resync_encryption_info())

    async def _resync_encryption_info(self) -> None:
        rows = await self.crypto_db.fetch(
            """SELECT room_id FROM mx_room_state WHERE encryption='{"resync":true}'"""
        )
        room_ids = [row["room_id"] for row in rows]
        if not room_ids:
            return
        self.log.debug(f"Resyncing encryption state event in rooms: {room_ids}")
        for room_id in room_ids:
            try:
                evt = await self.client.get_state_event(room_id, EventType.ROOM_ENCRYPTION)
            except (MNotFound, MForbidden) as e:
                self.log.debug(f"Failed to get encryption state in {room_id}: {e}")
                q = """
                    UPDATE mx_room_state SET encryption=NULL
                    WHERE room_id=$1 AND encryption='{"resync":true}'
                """
                await self.crypto_db.execute(q, room_id)
            else:
                self.log.debug(f"Resynced encryption state in {room_id}: {evt}")
                q = """
                    UPDATE crypto_megolm_inbound_session SET max_age=$1, max_messages=$2
                    WHERE room_id=$3 AND max_age IS NULL and max_messages IS NULL
                """
                await self.crypto_db.execute(
                    q, evt.rotation_period_ms, evt.rotation_period_msgs, room_id
                )

    async def _verify_keys_are_on_server(self) -> None:
        self.log.debug("Making sure keys are still on server")
        try:
            resp = await self.client.query_keys([self.client.mxid])
        except Exception:
            self.log.critical(
                "Failed to query own keys to make sure device still exists", exc_info=True
            )
            sys.exit(33)
        try:
            own_keys = resp.device_keys[self.client.mxid][self.client.device_id]
            if len(own_keys.keys) > 0:
                return
        except KeyError:
            pass
        self.log.critical("Existing device doesn't have keys on server, resetting crypto")
        await self.crypto.crypto_store.delete()
        await self.client.logout_all()
        sys.exit(34)

    async def stop(self) -> None:
        if self._key_delete_task:
            self._key_delete_task.cancel()
            self._key_delete_task = None
        self.client.stop()
        await self.crypto_store.close()
        if self.crypto_db:
            await self.crypto_db.stop()

    @property
    def _filter(self) -> Filter:
        all_events = EventType.find("*")
        return Filter(
            account_data=EventFilter(types=[all_events]),
            presence=EventFilter(not_types=[all_events]),
            room=RoomFilter(
                include_leave=False,
                state=StateFilter(not_types=[all_events]),
                timeline=RoomEventFilter(not_types=[all_events]),
                account_data=RoomEventFilter(not_types=[all_events]),
                ephemeral=RoomEventFilter(not_types=[all_events]),
            ),
        )

    async def _periodically_delete_keys(self) -> None:
        while True:
            deleted = await self.crypto_store.redact_expired_group_sessions()
            if deleted:
                self.log.info(f"Deleted expired megolm sessions: {deleted}")
            else:
                self.log.debug("No expired megolm sessions found")
            await asyncio.sleep(24 * 60 * 60)