use std::collections::BTreeMap;

use rand::prelude::*;

use crate::types::Candidate;

/// Reorder tracks so that artists are evenly spread out.
/// Greedily picks from the artist with the most remaining tracks,
/// avoiding back-to-back repeats when possible.
/// Ported faithfully from drift's interleave_artists().
pub fn interleave_artists(tracks: Vec<Candidate>) -> Vec<Candidate> {
    let mut rng = rand::rng();

    let mut by_artist: BTreeMap<String, Vec<Candidate>> = BTreeMap::new();
    for track in tracks {
        by_artist
            .entry(track.artist.clone())
            .or_default()
            .push(track);
    }
    for group in by_artist.values_mut() {
        group.shuffle(&mut rng);
    }

    let mut result = Vec::new();
    let mut last_artist: Option<String> = None;

    while !by_artist.is_empty() {
        let mut artists: Vec<String> = by_artist.keys().cloned().collect();
        artists.sort_by(|a, b| by_artist[b].len().cmp(&by_artist[a].len()));

        let pick = artists
            .iter()
            .find(|a| last_artist.as_ref() != Some(a))
            .or(artists.first())
            .cloned()
            .unwrap();

        let group = by_artist.get_mut(&pick).unwrap();
        let track = group.pop().unwrap();
        if group.is_empty() {
            by_artist.remove(&pick);
        }

        last_artist = Some(pick);
        result.push(track);
    }

    result
}

/// Full random shuffle.
pub fn shuffle(mut tracks: Vec<Candidate>) -> Vec<Candidate> {
    let mut rng = rand::rng();
    tracks.shuffle(&mut rng);
    tracks
}