macltok

use std::cmp::Ordering;

use ruma::OwnedRoomId;

const FIRST: u8 = 0x20;

const LAST: u8 = 0x7e;

const BASE: u64 = (LAST - FIRST + 1) as u64;

const MAX_LEN: usize = 50;

pub(crate) const EVENT_TYPE: &str = "org.matrix.msc3230.space_order";

pub(crate) fn valid(order: &str) -> bool {

    (1..=MAX_LEN).contains(&order.len()) && order.bytes().all(|b| (FIRST..=LAST).contains(&b))

}

/// how msc3230 clients sort spaces: ordered ones by (order, room id), then the rest by room id

pub(crate) fn compare(

    (a_id, a_order): (&OwnedRoomId, Option<&str>),

    (b_id, b_order): (&OwnedRoomId, Option<&str>),

) -> Ordering {

    match (a_order, b_order) {

        (Some(a), Some(b)) => a.cmp(b).then_with(|| a_id.cmp(b_id)),

        (Some(_), None) => Ordering::Less,

        (None, Some(_)) => Ordering::Greater,

        (None, None) => a_id.cmp(b_id),

    }

}

/// the orders to write so `spaces`, given with their current valid orders, sort in

/// exactly this sequence. keeps the longest run of orders that already sorts right

pub(crate) fn plan(spaces: &[(OwnedRoomId, Option<String>)]) -> Vec<(OwnedRoomId, String)> {

    let key = |index: usize| (&spaces[index].0, spaces[index].1.as_deref());

    let n = spaces.len();

    let mut unordered_tail = n;

    while unordered_tail > 0

        && spaces[unordered_tail - 1].1.is_none()

        && (unordered_tail == n

            || compare(key(unordered_tail - 1), key(unordered_tail)) == Ordering::Less)

    {

        unordered_tail -= 1;

    }

    let mut chain = vec![0usize; n];

    let mut previous = vec![None; n];

    for i in (0..n).filter(|&i| spaces[i].1.is_some()) {

        chain[i] = 1;

        for j in 0..i {

            if chain[j] > 0 && chain[j] + 1 > chain[i] && compare(key(j), key(i)) == Ordering::Less

            {

                chain[i] = chain[j] + 1;

                previous[i] = Some(j);

            }

        }

    }

    let mut longest: Option<usize> = None;

    let mut best: Option<(usize, usize, Option<usize>)> = None;

    for split in 0..=n {

        if split > 0 && chain[split - 1] > longest.map_or(0, |end| chain[end]) {

            longest = Some(split - 1);

        }

        if split < unordered_tail {

            continue;

        }

        let cost = split - longest.map_or(0, |end| chain[end]);

        if best.is_none_or(|(best_cost, _, _)| cost < best_cost) {

            best = Some((cost, split, longest));

        }

    }

    let Some((_, split, end)) = best else {

        return Vec::new();

    };

    let mut kept = vec![false; n];

    let mut cursor = end;

    while let Some(index) = cursor {

        kept[index] = true;

        cursor = previous[index];

    }

    let mut writes = Vec::new();

    let mut run_start = 0;

    for index in 0..=split {

        if index < split && !kept[index] {

            continue;

        }

        if index > run_start {

            let lo = run_start

                .checked_sub(1)

                .and_then(|i| spaces[i].1.as_deref());

            let hi = (index < split)

                .then(|| spaces[index].1.as_deref())

                .flatten();

            let Some(orders) = between(lo, hi, index - run_start) else {

                return rebalance(&spaces[..split]);

            };

            writes.extend(

                spaces[run_start..index]

                    .iter()

                    .map(|(id, _)| id.clone())

                    .zip(orders),

            );

        }

        run_start = index + 1;

    }

    writes

}

fn rebalance(spaces: &[(OwnedRoomId, Option<String>)]) -> Vec<(OwnedRoomId, String)> {

    let orders = between(None, None, spaces.len()).unwrap_or_default();

    spaces

        .iter()

        .zip(orders)

        .filter(|((_, current), order)| current.as_ref() != Some(order))

        .map(|((id, _), order)| (id.clone(), order))

        .collect()

}

/// `count` evenly spread orders strictly between `lo` and `hi`, as short as they can be

fn between(lo: Option<&str>, hi: Option<&str>, count: usize) -> Option<Vec<String>> {

    let lo = lo.unwrap_or("").as_bytes();

    (lo.len().max(1)..=MAX_LEN).find_map(|len| {

        let lo = digits(lo, len);

        let hi = match hi {

            Some(hi) => digits(&hi.as_bytes()[..hi.len().min(len)], len),

            None => {

                let mut one = vec![0; len + 1];

                one[0] = 1;

                one

            }

        };

        if hi <= lo {

            return None;

        }

        let step = divide(&subtract(&hi, &lo), count as u64 + 1);

        if step.iter().all(|&digit| digit == 0) {

            return None;

        }

        let mut point = lo;

        Some(

            (0..count)

                .map(|_| {

                    point = add(&point, &step);

                    let mut order: Vec<u8> = point[1..].iter().map(|&d| FIRST + d as u8).collect();

                    while order.last() == Some(&FIRST) {

                        order.pop();

                    }

                    String::from_utf8(order).expect("orders are ascii")

                })

                .collect(),

        )

    })

}

/// big-endian base-95 digits with one leading digit of headroom, zero padded to `len`

fn digits(order: &[u8], len: usize) -> Vec<u64> {

    let mut out = vec![0; len + 1];

    for (slot, byte) in out[1..].iter_mut().zip(order) {

        *slot = u64::from(byte - FIRST);

    }

    out

}

fn subtract(a: &[u64], b: &[u64]) -> Vec<u64> {

    let mut out = vec![0; a.len()];

    let mut borrow = 0;

    for i in (0..a.len()).rev() {

        let (value, under) = (a[i] + BASE - b[i] - borrow, a[i] < b[i] + borrow);

        out[i] = value % BASE;

        borrow = u64::from(under);

    }

    out

}

fn add(a: &[u64], b: &[u64]) -> Vec<u64> {

    let mut out = vec![0; a.len()];

    let mut carry = 0;

    for i in (0..a.len()).rev() {

        let sum = a[i] + b[i] + carry;

        out[i] = sum % BASE;

        carry = sum / BASE;

    }

    out

}

fn divide(a: &[u64], divisor: u64) -> Vec<u64> {

    let mut remainder = 0;

    a.iter()

        .map(|&digit| {

            let current = remainder * BASE + digit;

            remainder = current % divisor;

            current / divisor

        })

        .collect()

}

#[cfg(test)]

mod tests {

    use ruma::owned_room_id;

    use super::*;

    fn apply(spaces: &[(OwnedRoomId, Option<String>)]) -> (Vec<OwnedRoomId>, usize) {

        let writes = plan(spaces);

        let mut after: Vec<_> = spaces

            .iter()

            .map(|(id, order)| {

                let written = writes.iter().find(|(w, _)| w == id).map(|(_, o)| o.clone());

                (id.clone(), written.or_else(|| order.clone()))

            })

            .collect();

        assert!(after.iter().all(|(_, o)| o.as_deref().is_none_or(valid)));

        after.sort_by(|a, b| compare((&a.0, a.1.as_deref()), (&b.0, b.1.as_deref())));

        (after.into_iter().map(|(id, _)| id).collect(), writes.len())

    }

    fn ids(spaces: &[(OwnedRoomId, Option<String>)]) -> Vec<OwnedRoomId> {

        spaces.iter().map(|(id, _)| id.clone()).collect()

    }

    #[test]

    fn already_sorted_spaces_need_no_writes() {

        let spaces = vec![

            (owned_room_id!("!z:x"), Some("a".to_owned())),

            (owned_room_id!("!a:x"), None),

            (owned_room_id!("!b:x"), None),

        ];

        assert_eq!(apply(&spaces), (ids(&spaces), 0));

    }

    #[test]

    fn moving_one_space_writes_one_order() {

        let spaces = vec![

            (owned_room_id!("!c:x"), Some("c".to_owned())),

            (owned_room_id!("!a:x"), Some("a".to_owned())),

            (owned_room_id!("!b:x"), Some("b".to_owned())),

            (owned_room_id!("!d:x"), Some("d".to_owned())),

        ];

        assert_eq!(apply(&spaces), (ids(&spaces), 1));

    }

    #[test]

    fn unordered_spaces_out_of_id_order_get_orders() {

        let spaces = vec![

            (owned_room_id!("!c:x"), None),

            (owned_room_id!("!a:x"), None),

            (owned_room_id!("!b:x"), None),

        ];

        assert_eq!(apply(&spaces), (ids(&spaces), 1));

    }

    #[test]

    fn orders_with_no_room_between_them_rebalance() {

        let spaces = vec![

            (owned_room_id!("!a:x"), Some("a".to_owned())),

            (owned_room_id!("!c:x"), None),

            (owned_room_id!("!b:x"), Some("a ".to_owned())),

        ];

        assert_eq!(apply(&spaces).0, ids(&spaces));

        let prefix = "~".repeat(MAX_LEN - 1);

        let spaces = vec![

            (owned_room_id!("!b:x"), Some(format!("{prefix}!"))),

            (owned_room_id!("!a:x"), None),

            (owned_room_id!("!c:x"), Some(format!("{prefix}\""))),

        ];

        assert_eq!(apply(&spaces).0, ids(&spaces));

    }

}