// Copyright (C) Moondance Labs Ltd.

// This file is part of Tanssi.

// Tanssi is free software: you can redistribute it and/or modify

// it under the terms of the GNU General Public License as published by

// the Free Software Foundation, either version 3 of the License, or

// (at your option) any later version.

// Tanssi is distributed in the hope that it will be useful,

// but WITHOUT ANY WARRANTY; without even the implied warranty of

// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License

// along with Tanssi.  If not, see <http://www.gnu.org/licenses/>

use {

    alloc::{

        collections::{btree_map::BTreeMap, btree_set::BTreeSet},

        vec::Vec,

    },

    core::{cmp, marker::PhantomData, mem},

    dp_collator_assignment::AssignedCollators,

    frame_support::traits::Get,

    sp_runtime::Saturating,

    tp_traits::{

        FullRotationMode, FullRotationModes, ParaId, RemoveInvulnerables as RemoveInvulnerablesT,

    },

};

// Separate import of `alloc::vec!` macro, which cause issues with rustfmt if grouped

// with `alloc::vec::Vec`.

use alloc::vec;

/// Helper methods to implement collator assignment algorithm

pub struct Assignment<T>(PhantomData<T>);

impl<T> Assignment<T>

where

    T: crate::Config,

{

    /// Assign new collators to missing container_chains.

    /// Old collators always have preference to remain on the same chain.

    /// If there are no missing collators, nothing is changed.

    ///

    /// `chains` should be shuffled or at least rotated on every session to ensure

    /// a fair distribution, because the order of that list affects container chain priority:

    /// the first chain on that list will be the first one to get new collators.

    ///

    /// Similarly, in the `collators` list order means priority, the first collators will be more

    /// likely to get assigned. Unlike the list of `chains` which should already be shuffled,

    /// collators will be shuffled using the `shuffle` callback when needed. This allows the

    /// algorithm to truncate the list of collators and only shuffle the first N. This ensures that

    /// shuffling doesn't cause a collator with low priority to be assigned instead of a collator

    /// with higher priority.

    {

        // The rest of this function mostly treats orchestrator chain as another container chain, so move it into

        // `old_assigned.container_chains`

        // Orchestrator chain must be the first one in the list because it always has priority

        );

        // Remove some previously assigned collators to allow new collators to take their place.

        // Based on full_rotation_mode. In regular sessions this is FullRotationMode::KeepAll, a no-op.

            } else {

            };

        }

        // Ensure the first `min_orchestrator_collators` of orchestrator chain are invulnerables

        // Invulnerables can be unassigned by `keep_collator_subset`, but here we will assign other

        // invulnerables again. The downside is that the new invulnerables can be different.

        // Add container chains with 0 collators so that they are shown in UI

        // The rest of this function mostly treats orchestrator chain as another container chain, remove it from

        // container chains before returning the final assignment.

        // Sanity check to avoid bricking orchestrator chain

    /// Keep a subset of collators instead of rotating all of them.

    {

        };

            FullRotationMode::KeepAll => {

            }

        };

            // Remove all

        // Less than N collators, no need to shuffle

        // Shuffle and keep first N

    /// Select which container chains will be assigned collators and how many collators, but do not specify which

    /// collator goes to which chain.

    ///

    /// Each chain has a min and max number of collators. If the number of collators is not enough to reach the min,

    /// no collators are assigned to that chain.

    ///

    /// If the available number of collators is:

    /// * lower than the min of the first chain: we assign all the collators to the first chain. This is the

    ///   orchestrator chain and we always want it to have collators.

    /// * lower than the sum of all the min: we cannot assign collators to all the chains. So remove chains until

    ///   we can. The order is important, the first chains will be assigned collators and the last ones will not.

    /// * lower than the sum of all the max: we can assign the min value to all the chains, and have some leftover.

    ///   We use the same order to decide where this extra collators will go, by filling the max of the first chain,

    ///   then the max of the second chain, and so on.

    /// * greater than the sum of all the max: all the chains will be assigned their max number of collators.

    ///

    /// # Params

    ///

    /// The first item of `chains` should be the orchestrator chain, because it will be the first one to be assigned

    /// collators.

    ///

    /// # Returns

    ///

    /// A list of `(para_id, num_collators)`.

            // Avoid panic if chains is empty

        // Let's count how many container chains we can support with the current number of collators

        // Handle orchestrator chain in a special way, we always want to assign collators to it, even if we don't

        // reach the min.

        // Skipping orchestrator chain because it was handled above

                // Do not break if there are still some available collators. Even if they were not enough to reach the

                // `min` of this chain, it is possible that one of the chains with less priority has a lower `min`, so

                // that chain should be assigned collators.

        }

            // Edge case: num collators less than min orchestrator collators: fill as much as we can

        } else {

            // After assigning the min to all the chains we have this remainder. The remainder will be assigned until

            // all the chains reach the max value.

        }

    /// Same as `prioritize_invulnerables` but return the invulnerables instead of inserting them into `old_assigned`.

    ///

    /// Mutates `old_assigned` by removing invulnerables from their old chain, even if they will later be assigned to

    /// the same chain.

        // TODO: clean this up, maybe change remove_invulnerables trait into something more ergonomic

        );

            // We already had invulnerables, we will just move them to the front of the list if they weren't already

        // Not enough invulnerables currently assigned, get rest from new_collators

        }

        );

            // Got invulnerables from new_collators, and maybe some were already assigned

        // Still not enough invulnerables, try to get an invulnerable that is currently assigned somewhere else

            // Remove collators already selected

            // If at this point we still do not have enough invulnerables, it means that there are no

            // enough invulnerables, so no problem, but return the invulnerables

        // In this case we must delete the old assignment of the invulnerables

        // old_assigned.remove_collators_in_set

        }

    /// Ensure orchestrator chain has `min_orchestrator` invulnerables. If that's not possible, it tries to add as

    /// many invulnerables as possible.

    ///

    /// Get invulnerables from:

    /// * old_assigned in orchestrator

    /// * new collators

    /// * old_assigned elsewhere

    ///

    /// In that order.

    ///

    /// Mutates `old_assigned` because invulnerables will be inserted there, and if invulnerables were already

    /// assigned to some other chain, they will be removed from that other chain as well.

    ///

    /// # Params

    ///

    /// * `old_assigned` must be a subset of `collators`

    /// * `old_assigned` must not have duplicate collators.

    ///

    /// # Returns

    ///

    /// The number of invulnerables assigned to the orchestrator chain, capped to `min_collators`.

    /// Assign collators assuming that the number of collators is greater than or equal to the required.

    /// The order of both container chains and collators is important to ensure randomness when `old_assigned` is

    /// empty.

    ///

    /// # Params

    ///

    /// * `old_assigned` does not need to be a subset of `collators`: collators are checked and removed.

    /// * `old_assigned` does not need to be a subset of `chains`, unused para ids are removed. Collators

    ///   assigned to a para_id not present in `chains` may be reassigned to another para_id.

    /// * `chains` `num_collators` can be 0. In that case an empty vec is returned for that para id.

    /// * `old_assigned` must not have duplicate collators.

    /// * `shuffle` is used to shuffle the list collators. The list will be truncated to only have

    ///   the number of required collators, to ensure that shuffling doesn't cause a collator with low

    ///   priority to be assigned instead of a collator with higher priority.

    ///

    /// # Returns

    ///

    /// The collator assigment, a map from `ParaId` to `Vec<T>`.

    ///

    /// Or an error if the number of collators is not enough to fill all the chains, or if the required number

    /// of collators overflows a `u32`.

    {

        }

        // This check is necessary to ensure priority: if the number of collators is less than required, it is

        // possible that the chain with the least priority could be assigned collators (since they are in

        // old_assigned), while some chains with higher priority might have no collators.

        // We checked that the sum of all `num_collators` fits in `usize`, so we can safely use `as usize`.

        // Remove invalid collators and para ids from `old_assigned`

        // Truncate num collators to required

        // Count number of needed new collators. This is equivalent to:

        // `required_collators - old_assigned.iter().map(|cs| cs.len()).sum()`.

        // Truncate list of new_collators to `needed_new_collators` and shuffle it.

        // This has the effect of keeping collator priority (the first collator of that list is more

        // likely to be assigned to a chain than the last collator of that list), while also

        // ensuring randomness (the original order does not directly affect which chain the

        // collators are assigned to).

                // Keep collators not already assigned

        // Fill missing collators

                // This error should never happen because we calculated `needed_new_collators`

                // using the same algorithm

            }

        }

    /// Insert invulnerables ensuring that they are always the first in the list.

    /// The order of both lists is preserved.

    /// `assigned` may already contain the invulnerables, in that case they are only moved to the front.

    ///

    /// Invulnerables need to be the first of the list because we may truncate the list of collators if the number of

    /// collators changes, and in that case we want invulnerables to stay assigned there.

    /// Removes invalid entries from `old_assigned`:

    ///

    /// * para ids not in `chains_with_collators`

    /// * collators not in `collators`

        // old_assigned.remove_container_chains_not_in_set

        // old_assigned.remove_collators_not_in_set

        }

}

/// Errors than can happen during collator assignment

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

pub enum AssignmentError {

    /// An empty list of collators was passed to `assign_collators_always_keep_old`

    ZeroCollators,

    /// The required number of collators for `assign_full` is greater than the provided number of collators.

    /// Also includes possible overflows in number of collators.

    NotEnoughCollators,

    /// No collators were assigned to orchestrator chain

    EmptyOrchestrator,

}

/// A `ParaId` and a range of collators that need to be assigned to it.

/// This can be a container chain, a parathread, or the orchestrator chain.

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

pub struct ChainNumCollators {

    /// Para id.

    pub para_id: ParaId,

    /// Min collators.

    pub min_collators: u32,

    /// Max collators. This will only be filled if all the other chains have reached min_collators

    pub max_collators: u32,

    /// True if this a parathread. False means parachain or orchestrator.

    pub parathread: bool,

}