// Copyright 2020 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{
    collections::{BTreeMap, BTreeSet, HashMap, HashSet},
    sync::Arc,
};

use dashmap::DashMap;
use matrix_sdk_common::{
    api::r0::to_device::DeviceIdOrAllDevices,
    events::{
        room::{encrypted::EncryptedEventContent, history_visibility::HistoryVisibility},
        AnyMessageEventContent, EventType,
    },
    identifiers::{DeviceId, DeviceIdBox, RoomId, UserId},
    uuid::Uuid,
};
use serde_json::Value;
use tracing::{debug, info};

use crate::{
    error::{EventError, MegolmResult, OlmResult},
    olm::{Account, InboundGroupSession, OutboundGroupSession, Session, ShareState},
    store::{Changes, Store},
    Device, EncryptionSettings, OlmError, ToDeviceRequest,
};

#[derive(Debug, Clone)]
pub struct GroupSessionManager {
    account: Account,
    /// Store for the encryption keys.
    /// Persists all the encryption keys so a client can resume the session
    /// without the need to create new keys.
    store: Store,
    /// The currently active outbound group sessions.
    outbound_group_sessions: Arc<DashMap<RoomId, OutboundGroupSession>>,
    /// A map from the request id to the group session that the request belongs
    /// to. Used to mark requests belonging to the session as shared.
    outbound_sessions_being_shared: Arc<DashMap<Uuid, OutboundGroupSession>>,
}

impl GroupSessionManager {
    const MAX_TO_DEVICE_MESSAGES: usize = 20;

    pub(crate) fn new(account: Account, store: Store) -> Self {
        Self {
            account,
            store,
            outbound_group_sessions: Arc::new(DashMap::new()),
            outbound_sessions_being_shared: Arc::new(DashMap::new()),
        }
    }

    pub fn invalidate_group_session(&self, room_id: &RoomId) -> bool {
        if let Some(s) = self.outbound_group_sessions.get(room_id) {
            s.invalidate_session();
            true
        } else {
            false
        }
    }

    pub fn mark_request_as_sent(&self, request_id: &Uuid) {
        if let Some((_, s)) = self.outbound_sessions_being_shared.remove(request_id) {
            s.mark_request_as_sent(request_id);
        }
    }

    /// Get an outbound group session for a room, if one exists.
    ///
    /// # Arguments
    ///
    /// * `room_id` - The id of the room for which we should get the outbound
    /// group session.
    pub fn get_outbound_group_session(&self, room_id: &RoomId) -> Option<OutboundGroupSession> {
        #[allow(clippy::map_clone)]
        self.outbound_group_sessions.get(room_id).map(|s| s.clone())
    }

    pub async fn encrypt(
        &self,
        room_id: &RoomId,
        content: AnyMessageEventContent,
    ) -> MegolmResult<EncryptedEventContent> {
        let session = if let Some(s) = self.get_outbound_group_session(room_id) {
            s
        } else {
            panic!("Session wasn't created nor shared");
        };

        if session.expired() {
            panic!("Session expired");
        }

        let content = session.encrypt(content).await;

        let mut changes = Changes::default();
        changes.outbound_group_sessions.push(session);
        self.store.save_changes(changes).await?;

        Ok(content)
    }

    /// Create a new outbound group session.
    ///
    /// This also creates a matching inbound group session and saves that one in
    /// the store.
    pub async fn create_outbound_group_session(
        &self,
        room_id: &RoomId,
        settings: EncryptionSettings,
    ) -> OlmResult<(OutboundGroupSession, InboundGroupSession)> {
        let (outbound, inbound) = self
            .account
            .create_group_session_pair(room_id, settings)
            .await
            .map_err(|_| EventError::UnsupportedAlgorithm)?;

        let _ = self
            .outbound_group_sessions
            .insert(room_id.to_owned(), outbound.clone());
        Ok((outbound, inbound))
    }

    pub async fn get_or_create_outbound_session(
        &self,
        room_id: &RoomId,
        settings: EncryptionSettings,
    ) -> OlmResult<(OutboundGroupSession, Option<InboundGroupSession>)> {
        // Get the cached session, if there isn't one load one from the store
        // and put it in the cache.
        let outbound_session = if let Some(s) = self.outbound_group_sessions.get(room_id) {
            Some(s.clone())
        } else if let Some(s) = self.store.get_outbound_group_sessions(room_id).await? {
            self.outbound_group_sessions
                .insert(room_id.clone(), s.clone());

            Some(s)
        } else {
            None
        };

        // If there is no session or the session has expired or is invalid,
        // create a new one.
        if let Some(s) = outbound_session {
            if s.expired() || s.invalidated() {
                self.create_outbound_group_session(room_id, settings)
                    .await
                    .map(|(o, i)| (o, i.into()))
            } else {
                Ok((s, None))
            }
        } else {
            self.create_outbound_group_session(room_id, settings)
                .await
                .map(|(o, i)| (o, i.into()))
        }
    }

    /// Encrypt the given content for the given devices and create a to-device
    /// requests that sends the encrypted content to them.
    async fn encrypt_session_for(
        &self,
        content: Value,
        devices: &[Device],
    ) -> OlmResult<(Uuid, ToDeviceRequest, Vec<Session>)> {
        let mut messages = BTreeMap::new();
        let mut changed_sessions = Vec::new();

        for device in devices {
            let encrypted = device.encrypt(EventType::RoomKey, content.clone()).await;

            let (used_session, encrypted) = match encrypted {
                Ok(c) => c,
                Err(OlmError::MissingSession)
                | Err(OlmError::EventError(EventError::MissingSenderKey)) => {
                    continue;
                }
                Err(e) => return Err(e),
            };

            changed_sessions.push(used_session);

            messages
                .entry(device.user_id().clone())
                .or_insert_with(BTreeMap::new)
                .insert(
                    DeviceIdOrAllDevices::DeviceId(device.device_id().into()),
                    serde_json::value::to_raw_value(&encrypted)?,
                );
        }

        let id = Uuid::new_v4();

        let request = ToDeviceRequest {
            event_type: EventType::RoomEncrypted,
            txn_id: id,
            messages,
        };

        Ok((id, request, changed_sessions))
    }

    /// Given a list of user and an outbound session, return the list of users
    /// and their devices that this session should be shared with.
    ///
    /// Returns a boolean indicating whether the session needs to be rotated and
    /// the list of users/devices that should receive the session.
    pub async fn collect_session_recipients(
        &self,
        users: impl Iterator<Item = &UserId>,
        history_visibility: HistoryVisibility,
        outbound: &OutboundGroupSession,
    ) -> OlmResult<(bool, HashMap<UserId, Vec<Device>>)> {
        let users: HashSet<&UserId> = users.collect();
        let mut devices: HashMap<UserId, Vec<Device>> = HashMap::new();

        debug!(
            users = ?users,
            history_visibility = ?history_visibility,
            session_id = outbound.session_id(),
            "Calculating group session recipients"
        );

        let users_shared_with: HashSet<UserId> = outbound
            .shared_with_set
            .iter()
            .map(|k| k.key().clone())
            .collect();

        let users_shared_with: HashSet<&UserId> = users_shared_with.iter().collect();

        // A user left if a user is missing from the set of users that should
        // get the session but is in the set of users that received the session.
        let user_left = !users_shared_with
            .difference(&users)
            .collect::<HashSet<_>>()
            .is_empty();

        let visibility_changed = outbound.settings().history_visibility != history_visibility;

        // To protect the room history we need to rotate the session if either:
        //
        // 1. Any user left the room.
        // 2. Any of the users' devices got deleted or blacklisted.
        // 3. The history visibility changed.
        //
        // This is calculated in the following code and stored in this variable.
        let mut should_rotate = user_left || visibility_changed;

        for user_id in users {
            let user_devices = self.store.get_user_devices(&user_id).await?;
            let non_blacklisted_devices: Vec<Device> = user_devices
                .devices()
                .filter(|d| !d.is_blacklisted())
                .collect();

            // If we haven't already concluded that the session should be
            // rotated for other reasons, we also need to check whether any
            // of the devices in the session got deleted or blacklisted in the
            // meantime. If so, we should also rotate the session.
            if !should_rotate {
                // Device IDs that should receive this session
                let non_blacklisted_device_ids: HashSet<&DeviceId> = non_blacklisted_devices
                    .iter()
                    .map(|d| d.device_id())
                    .collect();

                if let Some(shared) = outbound.shared_with_set.get(user_id) {
                    #[allow(clippy::map_clone)]
                    // Devices that received this session
                    let shared: HashSet<DeviceIdBox> =
                        shared.iter().map(|d| d.key().clone()).collect();
                    let shared: HashSet<&DeviceId> = shared.iter().map(|d| d.as_ref()).collect();

                    // The set difference between
                    //
                    // 1. Devices that had previously received the session, and
                    // 2. Devices that would now receive the session
                    //
                    // represents newly deleted or blacklisted devices. If this
                    // set is non-empty, we must rotate.
                    let newly_deleted_or_blacklisted = shared
                        .difference(&non_blacklisted_device_ids)
                        .collect::<HashSet<_>>();

                    if !newly_deleted_or_blacklisted.is_empty() {
                        should_rotate = true;
                    }
                };
            }

            devices
                .entry(user_id.clone())
                .or_insert_with(Vec::new)
                .extend(non_blacklisted_devices);
        }

        debug!(
            should_rotate = should_rotate,
            session_id = outbound.session_id(),
            "Done calculating group session recipients"
        );

        Ok((should_rotate, devices))
    }

    /// Get to-device requests to share a group session with users in a room.
    ///
    /// # Arguments
    ///
    /// `room_id` - The room id of the room where the group session will be
    /// used.
    ///
    /// `users` - The list of users that should receive the group session.
    ///
    /// `encryption_settings` - The settings that should be used for the group
    /// session.
    pub async fn share_group_session(
        &self,
        room_id: &RoomId,
        users: impl Iterator<Item = &UserId>,
        encryption_settings: impl Into<EncryptionSettings>,
    ) -> OlmResult<Vec<Arc<ToDeviceRequest>>> {
        debug!(
            room_id = room_id.as_str(),
            "Checking if a group session needs to be shared for room {}", room_id
        );

        let encryption_settings = encryption_settings.into();
        let history_visibility = encryption_settings.history_visibility.clone();
        let mut changes = Changes::default();

        let (outbound, inbound) = self
            .get_or_create_outbound_session(room_id, encryption_settings.clone())
            .await?;

        if let Some(inbound) = inbound {
            changes.outbound_group_sessions.push(outbound.clone());
            changes.inbound_group_sessions.push(inbound);
        }

        let (should_rotate, devices) = self
            .collect_session_recipients(users, history_visibility, &outbound)
            .await?;

        let outbound = if should_rotate {
            let old_session_id = outbound.session_id();

            let (outbound, inbound) = self
                .create_outbound_group_session(room_id, encryption_settings)
                .await?;
            changes.outbound_group_sessions.push(outbound.clone());
            changes.inbound_group_sessions.push(inbound);

            debug!(
                room_id = room_id.as_str(),
                old_session_id = old_session_id,
                session_id = outbound.session_id(),
                "A user/device has left the group since we last sent a message, \
                   rotating the outbound session.",
            );

            outbound
        } else {
            outbound
        };

        let devices: Vec<Device> = devices
            .into_iter()
            .map(|(_, d)| {
                d.into_iter().filter(|d| {
                    matches!(
                        outbound.is_shared_with(d.user_id(), d.device_id()),
                        ShareState::NotShared
                    )
                })
            })
            .flatten()
            .collect();

        let key_content = outbound.as_json().await;
        let message_index = outbound.message_index().await;

        if !devices.is_empty() {
            let users = devices.iter().fold(BTreeMap::new(), |mut acc, d| {
                acc.entry(d.user_id())
                    .or_insert_with(BTreeSet::new)
                    .insert(d.device_id());
                acc
            });

            info!(
                index = message_index,
                users = ?users,
                room_id = room_id.as_str(),
                "Sharing an outbound group session",
            );
        }

        for device_map_chunk in devices.chunks(Self::MAX_TO_DEVICE_MESSAGES) {
            let (id, request, used_sessions) = self
                .encrypt_session_for(key_content.clone(), device_map_chunk)
                .await?;

            if !request.messages.is_empty() {
                outbound.add_request(id, request.into(), message_index);
                self.outbound_sessions_being_shared
                    .insert(id, outbound.clone());
            }

            changes.sessions.extend(used_sessions);
        }

        let requests = outbound.pending_requests();

        if requests.is_empty() {
            debug!(
                room_id = room_id.as_str(),
                session_id = outbound.session_id(),
                "The outbound group session doesn't need to be shared with \
                    anyone, marking as shared",
            );
            outbound.mark_as_shared();
        }

        self.store.save_changes(changes).await?;

        Ok(requests)
    }
}