macltok

use std::{

    collections::{HashMap, HashSet, VecDeque},

    sync::Arc,

};

use matrix_sdk::{

    Client, Room,

    ruma::{

        Int,

        api::client::room::{

            Visibility as SdkVisibility,

            create_room::v3::{Request as CreateRoomRequest, RoomPreset as SdkPreset},

        },

        events::{

            AnyGlobalAccountDataEventContent, GlobalAccountDataEventType, MessageLikeEventType,

            StateEventType, room::power_levels::RoomPowerLevelsEventContent,

        },

        serde::Raw,

    },

};

use ruma::{OwnedRoomId, OwnedUserId};

use tokio::{sync::mpsc, task::JoinSet};

use super::commands::{Job, guarded};

use crate::{

    CommandOutput, CreateRoomOptions, Error, ErrorKind, PowerLevelChanges, RoomCommand, RoomPreset,

    RoomVisibility,

};

/// state changes that must not race each other. same room, same alias, same dm

/// partner or same account data type run in order, everything else runs at once

pub(super) enum Mutation {

    Room(RoomCommand),

    GlobalAccountData {

        event_type: GlobalAccountDataEventType,

        content: Raw<AnyGlobalAccountDataEventContent>,

    },

}

pub(super) type MutationJob = Job<Mutation>;

#[derive(Clone, Debug, Eq, Hash, PartialEq)]

enum Key {

    Room(OwnedRoomId),

    Alias(String),

    Direct(OwnedUserId),

    AccountData(String),

    Independent(u64),

}

pub(super) async fn run_scheduler(

    client: Client,

    mut jobs: mpsc::Receiver<MutationJob>,

    queue_limit: usize,

) {

    let fresh_dms = FreshDms::default();

    let mut waiting = HashMap::<Key, VecDeque<MutationJob>>::new();

    let mut running = HashSet::<Key>::new();

    let mut tasks = JoinSet::<Key>::new();

    let mut queued = 0_usize;

    let mut independent = 0_u64;

    let mut receiving = true;

    while receiving || !tasks.is_empty() {

        tokio::select! {

            biased;

            Some(finished) = tasks.join_next(), if !tasks.is_empty() => {

                // a job's own panic is caught inside it, so this is only ever a key

                let Ok(key) = finished else { continue };

                let next = waiting.get_mut(&key).and_then(VecDeque::pop_front);

                if waiting.get(&key).is_some_and(VecDeque::is_empty) {

                    waiting.remove(&key);

                }

                match next {

                    Some(job) => {

                        queued -= 1;

                        spawn(&mut tasks, &client, &fresh_dms, key, job);

                    }

                    None => {

                        running.remove(&key);

                    }

                }

            }

            job = jobs.recv(), if receiving => {

                let Some(job) = job else {

                    receiving = false;

                    continue;

                };

                let key = key(&job.command, client.user_id().map(ruma::UserId::server_name)).unwrap_or_else(|| {

                    independent += 1;

                    Key::Independent(independent)

                });

                if running.insert(key.clone()) {

                    spawn(&mut tasks, &client, &fresh_dms, key, job);

                } else if queued < queue_limit {

                    waiting.entry(key).or_default().push_back(job);

                    queued += 1;

                } else {

                    job.reply.done(Err(Error::new(

                        ErrorKind::QueueFull,

                        "room command queue is full",

                    )));

                }

            }

        }

    }

}

/// `server` is ours, since a new alias lives there

fn key(mutation: &Mutation, server: Option<&ruma::ServerName>) -> Option<Key> {

    let room_id = match mutation {

        Mutation::GlobalAccountData { event_type, .. } => {

            return Some(Key::AccountData(event_type.to_string()));

        }

        Mutation::Room(RoomCommand::Join { room, .. }) => {

            return Some(match OwnedRoomId::try_from(room.clone()) {

                Ok(room_id) => Key::Room(room_id),

                Err(alias) => Key::Alias(alias.to_string()),

            });

        }

        Mutation::Room(RoomCommand::CreateDirect(user_id)) => {

            return Some(Key::Direct(user_id.clone()));

        }

        Mutation::Room(RoomCommand::Create(options)) => {

            return match (&options.alias_localpart, options.invite.as_slice()) {

                (Some(localpart), _) => Some(Key::Alias(format!("#{localpart}:{}", server?))),

                (None, [user_id]) if options.direct => Some(Key::Direct(user_id.clone())),

                _ => None,

            };

        }

        Mutation::Room(

            RoomCommand::Leave(room_id)

            | RoomCommand::Forget(room_id)

            | RoomCommand::AcceptInvite(room_id)

            | RoomCommand::RejectInvite(room_id)

            | RoomCommand::InviteUser { room_id, .. }

            | RoomCommand::KickUser { room_id, .. }

            | RoomCommand::BanUser { room_id, .. }

            | RoomCommand::UnbanUser { room_id, .. }

            | RoomCommand::SetPowerLevels { room_id, .. }

            | RoomCommand::SetUserPowerLevel { room_id, .. }

            | RoomCommand::SetAccountData { room_id, .. },

        ) => room_id,

        Mutation::Room(_) => return None,

    };

    Some(Key::Room(room_id.clone()))

}

/// dms made here that the next sync hasn't told the sdk about yet, since `get_dm_room`

/// only knows `m.direct` once it comes back round

type FreshDms = Arc<std::sync::Mutex<HashMap<OwnedUserId, OwnedRoomId>>>;

fn spawn(

    tasks: &mut JoinSet<Key>,

    client: &Client,

    fresh_dms: &FreshDms,

    key: Key,

    job: MutationJob,

) {

    let client = client.clone();

    let fresh_dms = fresh_dms.clone();

    tasks.spawn(async move {

        let Job { command, reply } = job;

        reply.send(guarded("room command", execute(&client, &fresh_dms, command)).await);

        key

    });

}

async fn execute(

    client: &Client,

    fresh_dms: &FreshDms,

    mutation: Mutation,

) -> Result<CommandOutput, Error> {

    let command = match mutation {

        Mutation::GlobalAccountData {

            event_type,

            content,

        } => {

            client

                .account()

                .set_account_data_raw(event_type, content)

                .await

                .map_err(|source| Error::from_sdk("writing account data failed", source))?;

            return Ok(CommandOutput::Done);

        }

        Mutation::Room(command) => command,

    };

    let room = |room_id: &OwnedRoomId| client.get_room(room_id).ok_or_else(Error::room_unknown);

    let joined = |room: Room| CommandOutput::Room(room.room_id().to_owned());

    let sdk = |message: &'static str| move |source| Error::from_sdk(message, source);

    match command {

        RoomCommand::Join { room, via } => client

            .join_room_by_id_or_alias(&room, &via)

            .await

            .map(joined)

            .map_err(sdk("joining room failed")),

        RoomCommand::Leave(room_id) => room(&room_id)?

            .leave()

            .await

            .map(done)

            .map_err(sdk("leaving room failed")),

        RoomCommand::Forget(room_id) => room(&room_id)?

            .forget()

            .await

            .map(done)

            .map_err(sdk("forgetting room failed")),

        RoomCommand::AcceptInvite(room_id) => room(&room_id)?

            .join()

            .await

            .map(done)

            .map_err(sdk("accepting invite failed")),

        RoomCommand::RejectInvite(room_id) => room(&room_id)?

            .leave()

            .await

            .map(done)

            .map_err(sdk("rejecting invite failed")),

        RoomCommand::InviteUser { room_id, user_id } => room(&room_id)?

            .invite_user_by_id(&user_id)

            .await

            .map(done)

            .map_err(sdk("inviting user failed")),

        RoomCommand::KickUser {

            room_id,

            user_id,

            reason,

        } => room(&room_id)?

            .kick_user(&user_id, reason.as_deref())

            .await

            .map(done)

            .map_err(sdk("kicking user failed")),

        RoomCommand::BanUser {

            room_id,

            user_id,

            reason,

        } => room(&room_id)?

            .ban_user(&user_id, reason.as_deref())

            .await

            .map(done)

            .map_err(sdk("banning user failed")),

        RoomCommand::UnbanUser {

            room_id,

            user_id,

            reason,

        } => room(&room_id)?

            .unban_user(&user_id, reason.as_deref())

            .await

            .map(done)

            .map_err(sdk("unbanning user failed")),

        RoomCommand::SetPowerLevels { room_id, changes } => {

            apply_power_level_changes(&room(&room_id)?, changes)

                .await

                .map(done)

        }

        RoomCommand::SetUserPowerLevel {

            room_id,

            user_id,

            level,

        } => room(&room_id)?

            .update_power_levels(vec![(&user_id, power_level(level)?)])

            .await

            .map(done)

            .map_err(sdk("updating user power level failed")),

        RoomCommand::SetAccountData {

            room_id,

            event_type,

            content,

        } => room(&room_id)?

            .set_account_data_raw(event_type, content)

            .await

            .map(done)

            .map_err(sdk("writing room account data failed")),

        RoomCommand::Create(options) => client

            .create_room(create_room_request(options))

            .await

            .map(joined)

            .map_err(sdk("creating room failed")),

        RoomCommand::CreateDirect(user_id) => {

            if let Some(existing) = client.get_dm_room(&user_id) {

                return Ok(joined(existing));

            }

            let fresh = fresh_dms

                .lock()

                .expect("nothing panics holding this")

                .get(&user_id)

                .cloned();

            if let Some(room_id) = fresh.filter(|room_id| client.get_room(room_id).is_some()) {

                return Ok(CommandOutput::Room(room_id));

            }

            let room = client

                .create_dm(&user_id)

                .await

                .map_err(sdk("creating direct room failed"))?;

            fresh_dms

                .lock()

                .expect("nothing panics holding this")

                .insert(user_id, room.room_id().to_owned());

            Ok(joined(room))

        }

        _ => Err(Error::new(ErrorKind::Other, "room command was misrouted")),

    }

}

/// one state event carrying every requested change, since the sdk's own patcher is private

async fn apply_power_level_changes(room: &Room, changes: PowerLevelChanges) -> Result<(), Error> {

    let mut levels = room

        .power_levels()

        .await

        .map_err(|source| Error::from_sdk("reading power levels failed", source.into()))?;

    for (target, value) in [

        (&mut levels.ban, changes.ban),

        (&mut levels.invite, changes.invite),

        (&mut levels.kick, changes.kick),

        (&mut levels.redact, changes.redact),

        (&mut levels.events_default, changes.events_default),

        (&mut levels.state_default, changes.state_default),

        (&mut levels.users_default, changes.users_default),

        (&mut levels.notifications.room, changes.room_notification),

    ] {

        if let Some(value) = value {

            *target = power_level(value)?;

        }

    }

    for (event_type, value) in [

        (StateEventType::RoomName.into(), changes.room_name),

        (StateEventType::RoomAvatar.into(), changes.room_avatar),

        (StateEventType::RoomTopic.into(), changes.room_topic),

        (StateEventType::SpaceChild.into(), changes.space_child),

        (MessageLikeEventType::Beacon.into(), changes.beacon),

        (StateEventType::BeaconInfo.into(), changes.beacon_info),

    ] {

        if let Some(value) = value {

            levels.events.insert(event_type, power_level(value)?);

        }

    }

    let content = RoomPowerLevelsEventContent::try_from(levels).map_err(|source| {

        Error::with_source(

            ErrorKind::Invalid,

            "normalising power levels failed",

            source,

        )

    })?;

    room.send_state_event(content)

        .await

        .map(drop)

        .map_err(|source| Error::from_sdk("updating power levels failed", source))

}

fn done<T>(_: T) -> CommandOutput {

    CommandOutput::Done

}

fn power_level(value: i64) -> Result<Int, Error> {

    value.try_into().map_err(|source| {

        Error::with_source(

            ErrorKind::Invalid,

            "power level is outside the Matrix integer range",

            source,

        )

    })

}

fn create_room_request(options: CreateRoomOptions) -> CreateRoomRequest {

    let mut request = CreateRoomRequest::new();

    request.name = options.name;

    request.topic = options.topic;

    request.room_alias_name = options.alias_localpart;

    request.invite = options.invite;

    request.is_direct = options.direct;

    request.visibility = match options.visibility {

        RoomVisibility::Private => SdkVisibility::Private,

        RoomVisibility::Public => SdkVisibility::Public,

    };

    request.preset = options.preset.map(|preset| match preset {

        RoomPreset::PrivateChat => SdkPreset::PrivateChat,

        RoomPreset::TrustedPrivateChat => SdkPreset::TrustedPrivateChat,

        RoomPreset::PublicChat => SdkPreset::PublicChat,

    });

    request

}