use super::*;
use codec::{Decode, DecodeWithMemTracking, Encode};
use frame_support::{storage_alias, traits::Get, weights::Weight};
use scale_info::TypeInfo;
use scale_info::prelude::string::String;
use sp_std::collections::btree_set::BTreeSet;
use sp_std::vec::Vec;

pub const MIGRATION_NAME: &[u8] = b"migrate_cleanup_staking_hotkeys";

/// Persistent progress for the bounded `StakingHotkeys` cleanup.
///
/// The current row's hotkeys are stored as independent candidates so a pass can remove a
/// bounded subset from the latest on-chain vector. This avoids overwriting hotkeys added by
/// normal staking operations while the migration is running.
#[derive(Encode, Decode, DecodeWithMemTracking, Clone, PartialEq, Eq, Debug, TypeInfo)]
pub struct StakingHotkeysCleanupProgress {
    /// Hashed key of the last fully processed `StakingHotkeys` row.
    pub after: Option<Vec<u8>>,
    /// SCALE-encoded coldkey for a row that is only partially processed.
    pub coldkey: Option<Vec<u8>>,
    /// SCALE-encoded hotkeys from that row which still need classification.
    pub remaining_hotkeys: Vec<Vec<u8>>,
    pub rows_scanned: u64,
    pub relationships_scanned: u64,
    pub relationships_removed: u64,
    pub vector_writes: u64,
}

#[storage_alias]
pub type StakingHotkeysCleanupMigration<T: Config> =
    StorageValue<Pallet<T>, StakingHotkeysCleanupProgress, OptionQuery>;

fn candidate_weight<T: Config>() -> Weight {
    // One prefix probe in each share map plus one BasketClaimed read. The predicate may
    // short-circuit, but charging all three reads keeps each candidate conservatively bounded.
    T::DbWeight::get().reads(3)
}

fn vector_rewrite_weight<T: Config>() -> Weight {
    T::DbWeight::get().reads_writes(1, 1)
}

fn row_load_weight<T: Config>() -> Weight {
    // Iterating the next StakingHotkeys row performs one backend read.
    T::DbWeight::get().reads(1)
}

/// A conservative keep predicate for one hotkey/coldkey relationship.
///
/// Any stored share row is retained, including a zero-valued legacy row. The storage-bloat
/// migration runs first and clears zero legacy rows, while treating an unexpected remaining row
/// as live makes this cleanup fail safe. A basket watermark is also sufficient to retain the
/// relationship because zero-root-stake claimants still need to be discoverable by claims and
/// coldkey swaps.
fn relationship_must_remain<T: Config>(hotkey: &T::AccountId, coldkey: &T::AccountId) -> bool {
    AlphaV2::<T>::iter_prefix((hotkey, coldkey))
        .next()
        .is_some()
        || Alpha::<T>::iter_prefix((hotkey, coldkey)).next().is_some()
        || BasketClaimed::<T>::get(hotkey, coldkey) != 0
}

/// Starts the cleanup without scanning `StakingHotkeys` in the runtime-upgrade block.
pub fn kickoff_staking_hotkeys_cleanup<T: Config>() -> Weight {
    let mut weight = T::DbWeight::get().reads(2);
    if HasMigrationRun::<T>::get(MIGRATION_NAME) || StakingHotkeysCleanupMigration::<T>::exists() {
        return weight;
    }

    StakingHotkeysCleanupMigration::<T>::put(StakingHotkeysCleanupProgress {
        after: None,
        coldkey: None,
        remaining_hotkeys: Vec::new(),
        rows_scanned: 0,
        relationships_scanned: 0,
        relationships_removed: 0,
        vector_writes: 0,
    });
    weight.saturating_accrue(T::DbWeight::get().writes(1));
    log::info!(
        "Migration '{}' scheduled for bounded on_idle execution",
        String::from_utf8_lossy(MIGRATION_NAME)
    );
    weight
}

/// Continues the cleanup without exceeding `limit`.
///
/// Normal operations remain enabled. Each batch removes candidates from the latest stored
/// vector, so stake or hotkeys added between passes are never overwritten by an old snapshot.
pub fn continue_staking_hotkeys_cleanup<T: Config>(limit: Weight) -> Weight {
    // Cursor read plus either cursor persistence, or completion marker + cursor removal.
    let pass_overhead = T::DbWeight::get().reads_writes(1, 2);
    if !pass_overhead.all_lte(limit) {
        return Weight::zero();
    }

    let Some(mut progress) = StakingHotkeysCleanupMigration::<T>::get() else {
        return T::DbWeight::get().reads(1);
    };
    let work_limit = limit.saturating_sub(pass_overhead);
    let mut work_weight = Weight::zero();
    let mut finished = false;

    loop {
        if progress.coldkey.is_none() {
            // Reserve a possible write as well: an encoded empty vector is removed immediately
            // after it is read, without carrying an ambiguous empty-row cursor across blocks.
            let row_load_reserve =
                row_load_weight::<T>().saturating_add(T::DbWeight::get().writes(1));
            if !work_weight
                .saturating_add(row_load_reserve)
                .all_lte(work_limit)
            {
                break;
            }

            // Drop the iterator before this row can be mutated below. Altering a map while one
            // of its iterators is alive has undefined iteration results.
            let next = {
                let mut iter = match progress.after.as_ref() {
                    Some(raw) => StakingHotkeys::<T>::iter_from(raw.clone()),
                    None => StakingHotkeys::<T>::iter(),
                };
                iter.next()
            };
            work_weight.saturating_accrue(row_load_weight::<T>());

            let Some((coldkey, hotkeys)) = next else {
                finished = true;
                break;
            };

            if hotkeys.is_empty() {
                let empty_row_write = T::DbWeight::get().writes(1);
                StakingHotkeys::<T>::remove(&coldkey);
                work_weight.saturating_accrue(empty_row_write);
                progress.after = Some(StakingHotkeys::<T>::hashed_key_for(&coldkey));
                progress.rows_scanned = progress.rows_scanned.saturating_add(1);
                progress.vector_writes = progress.vector_writes.saturating_add(1);
                continue;
            }

            progress.coldkey = Some(coldkey.encode());
            progress.remaining_hotkeys =
                hotkeys.into_iter().map(|hotkey| hotkey.encode()).collect();
        }

        let Some(encoded_coldkey) = progress.coldkey.as_ref() else {
            continue;
        };
        let Ok(coldkey) = T::AccountId::decode(&mut &encoded_coldkey[..]) else {
            log::error!(
                "Migration '{}' found an undecodable coldkey cursor; stopping this pass",
                String::from_utf8_lossy(MIGRATION_NAME)
            );
            break;
        };

        // Reserve enough room to rewrite the current vector if any processed candidate is stale.
        let rewrite_weight = vector_rewrite_weight::<T>();
        let mut removals: BTreeSet<T::AccountId> = BTreeSet::new();
        let mut processed_any = false;

        while let Some(encoded_hotkey) = progress.remaining_hotkeys.last() {
            if !work_weight
                .saturating_add(candidate_weight::<T>())
                .saturating_add(rewrite_weight)
                .all_lte(work_limit)
            {
                break;
            }

            let encoded_hotkey = encoded_hotkey.clone();
            progress.remaining_hotkeys.pop();
            let Ok(hotkey) = T::AccountId::decode(&mut &encoded_hotkey[..]) else {
                // The candidate came from a successfully decoded StakingHotkeys vector. If the
                // persisted cursor is ever corrupted, preserving this relationship is safer than
                // deleting an account we cannot identify.
                log::error!(
                    "Migration '{}' found an undecodable hotkey candidate; preserving it",
                    String::from_utf8_lossy(MIGRATION_NAME)
                );
                processed_any = true;
                progress.relationships_scanned = progress.relationships_scanned.saturating_add(1);
                work_weight.saturating_accrue(candidate_weight::<T>());
                continue;
            };

            if !relationship_must_remain::<T>(&hotkey, &coldkey) {
                removals.insert(hotkey);
            }
            processed_any = true;
            progress.relationships_scanned = progress.relationships_scanned.saturating_add(1);
            work_weight.saturating_accrue(candidate_weight::<T>());
        }

        if !removals.is_empty() {
            let mut removed = 0_u64;
            StakingHotkeys::<T>::mutate_exists(&coldkey, |maybe_hotkeys| {
                if let Some(hotkeys) = maybe_hotkeys {
                    let before = hotkeys.len();
                    hotkeys.retain(|hotkey| !removals.contains(hotkey));
                    removed =
                        u64::try_from(before.saturating_sub(hotkeys.len())).unwrap_or(u64::MAX);
                    if hotkeys.is_empty() {
                        *maybe_hotkeys = None;
                    }
                }
            });
            work_weight.saturating_accrue(rewrite_weight);
            progress.relationships_removed = progress.relationships_removed.saturating_add(removed);
            progress.vector_writes = progress.vector_writes.saturating_add(1);
        }

        if progress.remaining_hotkeys.is_empty() {
            progress.after = Some(StakingHotkeys::<T>::hashed_key_for(&coldkey));
            progress.coldkey = None;
            progress.rows_scanned = progress.rows_scanned.saturating_add(1);
            continue;
        }

        if !processed_any {
            break;
        }

        // The current row was only partially classified. Persist its remaining candidates and
        // resume from the latest vector on the next idle pass.
        break;
    }

    if finished {
        HasMigrationRun::<T>::insert(MIGRATION_NAME, true);
        StakingHotkeysCleanupMigration::<T>::kill();
        log::info!(
            "Migration '{}' completed: scanned {} rows / {} relationships, removed {} relationships with {} vector writes",
            String::from_utf8_lossy(MIGRATION_NAME),
            progress.rows_scanned,
            progress.relationships_scanned,
            progress.relationships_removed,
            progress.vector_writes,
        );
    } else {
        StakingHotkeysCleanupMigration::<T>::put(progress);
    }

    pass_overhead.saturating_add(work_weight)
}