vsrg 0.6.0

use super::*;
use crate::notes::{NoteGroup, NoteGroups};
use crate::rhythm::{ScrollSpeedTracks, TempoTracks};

impl<Note: LongNoteData> Default for LongNoteTypeStorage<Note> {
    fn default() -> Self {
        Self::new()
    }
}

impl<Note: LongNoteData> LongNoteTypeStorage<Note> {
    pub fn new() -> Self {
        Self {
            note_id_counter: 0,
            note_id_lookup: HashMap::new(),
            time_sorted_groups: vec![],
            scroll_sorted_groups: vec![],
        }
    }

    pub fn with_notes<G: NoteGroup>(
        note_groups: &NoteGroups<G>,
        tempo_tracks: &TempoTracks,
        scroll_tracks: &ScrollSpeedTracks,
        notes: impl Iterator<Item = Note>,
    ) -> Self {
        let mut storage = Self::new();
        for note in notes {
            let id = storage.advance_to_next_id();
            storage.insert(
                note_groups,
                tempo_tracks,
                scroll_tracks,
                id,
                note,
                Note::RuntimeData::default(),
            );
        }
        storage.sort_and_link();
        for group in &mut storage.time_sorted_groups {
            let note_group = note_groups.get_group_or_main(group.group_id);
            let tempo_track = tempo_tracks.get_track_or_main(note_group.tempo_track_id());
            let scroll_track = scroll_tracks.get_track_or_main(note_group.scroll_track_id());
            for index in 0..group.note_id.len() {
                apply_event_track_time_at(group, index, tempo_track, scroll_track);
            }
        }
        for group in &mut storage.scroll_sorted_groups {
            let note_group = note_groups.get_group_or_main(group.group_id);
            let tempo_track = tempo_tracks.get_track_or_main(note_group.tempo_track_id());
            let scroll_track = scroll_tracks.get_track_or_main(note_group.scroll_track_id());
            for index in 0..group.note_id.len() {
                apply_event_track_scroll_at(group, index, tempo_track, scroll_track);
            }
        }
        storage.health_check();
        storage
    }

    pub(super) fn advance_to_next_id(&mut self) -> u64 {
        loop {
            self.note_id_counter = self.note_id_counter.wrapping_add(1);
            if self.note_id_counter != 0 && !self.note_id_lookup.contains_key(&self.note_id_counter)
            {
                return self.note_id_counter;
            }
        }
    }

    fn insert<G: NoteGroup>(
        &mut self,
        groups: &NoteGroups<G>,
        tempos: &TempoTracks,
        scrolls: &ScrollSpeedTracks,
        id: u64,
        note: Note,
        runtime: Note::RuntimeData,
    ) {
        let (group_id, scroll_id, tempo_id) = get_ids(&note, groups);

        let start_time = note.start_time();
        let duration = note.duration();
        let clock = to_clock_interval(tempos, start_time, duration, tempo_id);
        let scroll_track = scrolls.get_track_or_main(scroll_id);
        let scroll = to_scroll_pos_on_track(scroll_track, clock);
        let runtime = runtime;
        let g = self
            .time_sorted_groups
            .iter()
            .position(|g| g.group_id == group_id)
            .unwrap_or_else(|| {
                self.time_sorted_groups
                    .push(TimeIndexedLongNoteGroup::new(group_id));
                self.scroll_sorted_groups
                    .push(ScrollIndexedLongNoteGroup::new(group_id));
                self.time_sorted_groups.len() - 1
            });
        let value = note.value_data();
        {
            let group = &mut self.time_sorted_groups[g];
            group.note_id.push(id);
            group.start_time.push(start_time);
            group.duration.push(duration);
            group.clock_time.push(clock);
            group.scroll_position.push(scroll);
            group.value_data.push(value.clone());
            group.runtime_data.push(runtime.clone());
        }
        {
            let group = &mut self.scroll_sorted_groups[g];
            group.note_id.push(id);
            group.start_time.push(start_time);
            group.duration.push(duration);
            group.clock_time.push(clock);
            group.scroll_position.push(scroll);
            group.value_data.push(value.clone());
            group.runtime_data.push(runtime.clone());
        }
    }

    /// Add a note to the storage.
    ///
    /// # Returns
    /// The id of the added note.
    pub fn add_note<G: NoteGroup>(
        &mut self,
        groups: &NoteGroups<G>,
        tempos: &TempoTracks,
        scrolls: &ScrollSpeedTracks,
        note: Note,
    ) -> u64 {
        self.add_notes(groups, tempos, scrolls, std::iter::once(note))[0]
    }

    /// Add multiple notes to the storage.
    ///
    /// # Returns
    /// The list of ids of added notes.
    pub fn add_notes<G: NoteGroup>(
        &mut self,
        groups: &NoteGroups<G>,
        tempos: &TempoTracks,
        scrolls: &ScrollSpeedTracks,
        notes: impl Iterator<Item = Note>,
    ) -> Vec<u64> {
        let mut ids = vec![];
        for note in notes {
            let id = self.advance_to_next_id();
            let group_id = note.group_id();
            self.insert(
                groups,
                tempos,
                scrolls,
                id,
                note,
                Note::RuntimeData::default(),
            );
            let note_group = groups.get_group_or_main(group_id);
            let tempo_track = tempos.get_track_or_main(note_group.tempo_track_id());
            let scroll_track = scrolls.get_track_or_main(note_group.scroll_track_id());
            let group_index = self
                .time_sorted_groups
                .iter()
                .position(|group| group.group_id == group_id)
                .unwrap();
            let time_group = &mut self.time_sorted_groups[group_index];
            let time_index = time_group.note_id.len() - 1;
            apply_event_track_time_at(time_group, time_index, tempo_track, scroll_track);
            let scroll_group = &mut self.scroll_sorted_groups[group_index];
            let scroll_index = scroll_group.note_id.len() - 1;
            apply_event_track_scroll_at(scroll_group, scroll_index, tempo_track, scroll_track);
            ids.push(id);
        }
        self.sort_and_link();
        self.health_check();
        ids
    }

    /// Remove a note from the storage, identified by its id.
    pub fn remove_note(&mut self, id: u64) {
        self.remove_notes(std::iter::once(id));
    }

    /// Remove multiple notes from the storage, identified by their ids.
    pub fn remove_notes(&mut self, ids: impl Iterator<Item = u64>) {
        let ids: std::collections::HashSet<_> = ids.collect();
        for group in &mut self.time_sorted_groups {
            for i in (0..group.note_id.len()).rev() {
                if ids.contains(&group.note_id[i]) {
                    group.remove(i);
                }
            }
        }
        for group in &mut self.scroll_sorted_groups {
            for i in (0..group.note_id.len()).rev() {
                if ids.contains(&group.note_id[i]) {
                    group.remove(i);
                }
            }
        }
        self.sort_and_link();
        self.health_check();
    }

    /// Replace a note while preserving its id and runtime state.
    pub fn replace_note<G: NoteGroup>(
        &mut self,
        groups: &NoteGroups<G>,
        tempos: &TempoTracks,
        scrolls: &ScrollSpeedTracks,
        id: u64,
        note: Note,
    ) {
        self.replace_notes(groups, tempos, scrolls, std::iter::once((id, note)));
    }

    /// Replace multiple notes while preserving their ids and runtime state.
    ///
    /// Replacement with unknown ids are skipped
    /// If multiple replacements repeat the same id, then the last replacement takes effect.
    pub fn replace_notes<G: NoteGroup>(
        &mut self,
        groups: &NoteGroups<G>,
        tempos: &TempoTracks,
        scrolls: &ScrollSpeedTracks,
        notes: impl Iterator<Item = (u64, Note)>,
    ) {
        for (id, note) in notes {
            let Some(loc) = self.note_id_lookup.get(&id).copied() else {
                continue;
            };
            let runtime =
                self.time_sorted_groups[loc.group_index].remove(loc.index_in_time_sorted_group);
            self.scroll_sorted_groups[loc.group_index].remove(loc.index_in_scroll_sorted_group);
            let group_id = note.group_id();
            self.insert(groups, tempos, scrolls, id, note, runtime);
            let ng = groups.get_group_or_main(group_id);
            let tempo_track = tempos.get_track_or_main(ng.tempo_track_id());
            let scroll_track = scrolls.get_track_or_main(ng.scroll_track_id());
            let group_index = self
                .time_sorted_groups
                .iter()
                .position(|group| group.group_id == group_id)
                .unwrap();
            let time_group = &mut self.time_sorted_groups[group_index];
            apply_event_track_time_at(
                time_group,
                time_group.note_id.len() - 1,
                tempo_track,
                scroll_track,
            );
            let scroll_group = &mut self.scroll_sorted_groups[group_index];
            apply_event_track_scroll_at(
                scroll_group,
                scroll_group.note_id.len() - 1,
                tempo_track,
                scroll_track,
            );
            self.relink();
        }
        self.sort_and_link();
        self.health_check();
    }

    /// Recalculate note's clock time reindex after edits to tempo tracks.
    pub fn rebuild_clock_time<G: NoteGroup>(
        &mut self,
        groups: &NoteGroups<G>,
        tempos: &TempoTracks,
        scrolls: &ScrollSpeedTracks,
    ) {
        for group in &mut self.time_sorted_groups {
            let ng = groups.get_group_or_main(group.group_id);
            let tempo_track_id = ng.tempo_track_id();
            let scroll_track_id = ng.scroll_track_id();
            let tempo_track = tempos.get_track_or_main(tempo_track_id);
            let scroll_track = scrolls.get_track_or_main(scroll_track_id);
            for i in 0..group.note_id.len() {
                let clock = to_clock_interval(
                    tempos,
                    group.start_time[i],
                    group.duration[i],
                    tempo_track_id,
                );
                group.clock_time[i] = clock;
                group.scroll_position[i] = to_scroll_pos_on_track(scroll_track, clock);
                apply_event_track_time_at(group, i, tempo_track, scroll_track);
            }
        }
        for group in &mut self.scroll_sorted_groups {
            let ng = groups.get_group_or_main(group.group_id);
            let tempo_track_id = ng.tempo_track_id();
            let scroll_track_id = ng.scroll_track_id();
            let tempo_track = tempos.get_track_or_main(tempo_track_id);
            let scroll_track = scrolls.get_track_or_main(scroll_track_id);
            for i in 0..group.note_id.len() {
                let clock = to_clock_interval(
                    tempos,
                    group.start_time[i],
                    group.duration[i],
                    tempo_track_id,
                );
                group.clock_time[i] = clock;
                group.scroll_position[i] = to_scroll_pos_on_track(scroll_track, clock);
                apply_event_track_scroll_at(group, i, tempo_track, scroll_track);
            }
        }
        self.sort_and_link();
        self.health_check();
    }

    /// Recalculate scroll position and reindexes after edits to scroll speed tracks.
    pub fn rebuild_scroll_pos<G: NoteGroup>(
        &mut self,
        groups: &NoteGroups<G>,
        tempos: &TempoTracks,
        scrolls: &ScrollSpeedTracks,
    ) {
        for group in &mut self.time_sorted_groups {
            let ng = groups.get_group_or_main(group.group_id);
            let tempo_track = tempos.get_track_or_main(ng.tempo_track_id());
            let scroll_track = scrolls.get_track_or_main(ng.scroll_track_id());
            for i in 0..group.note_id.len() {
                group.scroll_position[i] =
                    to_scroll_pos_on_track(scroll_track, group.clock_time[i]);
                apply_event_track_time_at(group, i, tempo_track, scroll_track);
            }
        }
        for group in &mut self.scroll_sorted_groups {
            let ng = groups.get_group_or_main(group.group_id);
            let tempo_track = tempos.get_track_or_main(ng.tempo_track_id());
            let scroll_track = scrolls.get_track_or_main(ng.scroll_track_id());
            for i in 0..group.note_id.len() {
                group.scroll_position[i] =
                    to_scroll_pos_on_track(scroll_track, group.clock_time[i]);
                apply_event_track_scroll_at(group, i, tempo_track, scroll_track);
            }
        }
        self.sort_and_link();
        self.health_check();
    }

    pub(super) fn sort_and_link(&mut self) {
        for group in &mut self.time_sorted_groups {
            group.rebuild();
        }
        for group in &mut self.scroll_sorted_groups {
            group.rebuild();
        }
        self.relink();
    }

    pub(super) fn relink(&mut self) {
        self.time_sorted_groups.retain(|g| !g.note_id.is_empty());
        self.scroll_sorted_groups.retain(|g| !g.note_id.is_empty());
        self.note_id_lookup.clear();
        for (g, group) in self.time_sorted_groups.iter().enumerate() {
            for (i, &id) in group.note_id.iter().enumerate() {
                self.note_id_lookup.insert(
                    id,
                    NoteLocation {
                        group_index: g,
                        index_in_time_sorted_group: i,
                        index_in_scroll_sorted_group: 0,
                    },
                );
            }
        }
        for group in &self.scroll_sorted_groups {
            for (i, id) in group.note_id.iter().enumerate() {
                self.note_id_lookup
                    .get_mut(id)
                    .unwrap()
                    .index_in_scroll_sorted_group = i;
            }
        }
    }

    pub fn remove_group(&mut self, group_id: &GroupId, mut handle_removed: impl FnMut(Note)) {
        let ids = self
            .time_sorted_groups
            .iter()
            .filter(|g| &g.group_id == group_id)
            .flat_map(|g| g.note_id.iter().copied())
            .collect::<Vec<_>>();
        self.extract_notes(ids.into_iter(), |_, note| handle_removed(note));
    }

    pub fn extract_notes(
        &mut self,
        ids: impl Iterator<Item = u64> + Clone,
        mut handle_removed: impl FnMut(u64, Note),
    ) {
        let mut seen = std::collections::HashSet::new();
        for (id, note) in self
            .query_notes(ids.clone())
            .filter(|(id, _)| seen.insert(*id))
        {
            handle_removed(id, note);
        }
        self.remove_notes(ids);
    }

    pub(crate) fn health_check(&self) {
        debug_assert_eq!(
            self.time_sorted_groups.len(),
            self.scroll_sorted_groups.len()
        );
        debug_assert_eq!(
            self.note_id_lookup.len(),
            self.time_sorted_groups
                .iter()
                .map(|g| g.note_id.len())
                .sum::<usize>()
        );
        for (g, (time, scroll)) in self
            .time_sorted_groups
            .iter()
            .zip(&self.scroll_sorted_groups)
            .enumerate()
        {
            debug_assert_eq!(time.group_id, scroll.group_id);
            debug_assert_eq!(time.note_id.len(), scroll.note_id.len());
            time.health_check();
            scroll.health_check();
            for (i, id) in time.note_id.iter().enumerate() {
                let loc = &self.note_id_lookup[id];
                let j = loc.index_in_scroll_sorted_group;
                debug_assert_eq!(loc.group_index, g);
                debug_assert_eq!(loc.index_in_time_sorted_group, i);
                debug_assert_eq!(scroll.note_id[j], *id);
                debug_assert_eq!(time.start_time[i], scroll.start_time[j]);
                debug_assert_eq!(time.duration[i], scroll.duration[j]);
                debug_assert_eq!(time.clock_time[i], scroll.clock_time[j]);
                debug_assert_eq!(time.scroll_position[i], scroll.scroll_position[j]);
            }
        }
    }
}

fn get_ids<Note: LongNoteData, G: NoteGroup>(
    note: &Note,
    note_groups: &NoteGroups<G>,
) -> (GroupId, TrackId, TrackId) {
    let group_id = note.group_id();
    let group = note_groups
        .get_group(group_id)
        .unwrap_or(note_groups.main_group());
    let scroll_track_id = group.scroll_track_id();
    let tempo_track_id = group.tempo_track_id();
    (group_id, scroll_track_id, tempo_track_id)
}

fn apply_event_track_time_at<Note: LongNoteData>(
    group: &mut TimeIndexedLongNoteGroup<Note>,
    index: usize,
    tempo_track: &crate::rhythm::TempoTrack,
    scroll_track: &crate::rhythm::ScrollSpeedTrack,
) {
    let note = crate::notes::NoteMut {
        id: group.note_id[index],
        start_time: group.start_time[index],
        duration: group.duration[index],
        clock_time: group.clock_time[index],
        group_id: group.group_id,
        value_data: group.value_data.get(index).unwrap(),
        runtime_data: group.runtime_data.get_mut(index).unwrap(),
    };
    Note::apply_event_track(note, tempo_track, scroll_track);
}

fn apply_event_track_scroll_at<Note: LongNoteData>(
    group: &mut ScrollIndexedLongNoteGroup<Note>,
    index: usize,
    tempo_track: &crate::rhythm::TempoTrack,
    scroll_track: &crate::rhythm::ScrollSpeedTrack,
) {
    let note = crate::notes::NoteMut {
        id: group.note_id[index],
        start_time: group.start_time[index],
        duration: group.duration[index],
        clock_time: group.clock_time[index],
        group_id: group.group_id,
        value_data: group.value_data.get(index).unwrap(),
        runtime_data: group.runtime_data.get_mut(index).unwrap(),
    };
    Note::apply_event_track(note, tempo_track, scroll_track);
}

fn to_scroll_pos_on_track(
    scroll_track: &crate::rhythm::ScrollSpeedTrack,
    clock_time: Interval<ClockTime>,
) -> Interval<ScrollPosition> {
    let start = scroll_track.calculate_scroll_position(clock_time.start());
    let end = scroll_track.calculate_scroll_position(clock_time.end());
    Interval::new(start, end)
}

fn to_clock_interval(
    tempo_tracks: &TempoTracks,
    start_time: Time,
    duration: Time,
    tempo_track_id: TrackId,
) -> Interval<ClockTime> {
    let start_clock_time = to_clock_time(tempo_tracks, start_time, tempo_track_id);
    let end_clock_time = match duration {
        Time::Clock(duration) => start_clock_time + duration,
        Time::Beat(duration) => {
            let track = tempo_tracks.get_track_or_main(tempo_track_id);
            let start_beat = match start_time {
                Time::Clock(start_time) => track.clock_to_beats(start_time),
                Time::Beat(start_time) => start_time,
            };
            track.beat_to_clock(start_beat + duration)
        }
    };
    Interval::new(start_clock_time, end_clock_time)
}

fn to_clock_time(tempo_tracks: &TempoTracks, time: Time, tempo_track_id: TrackId) -> ClockTime {
    let track = tempo_tracks.get_track_or_main(tempo_track_id);
    match time {
        Time::Clock(clock_time) => clock_time,
        Time::Beat(beat_time) => track.beat_to_clock(beat_time),
    }
}

pub(super) fn soa_set_element<T: Soars>(soa: &mut Soa<T>, index: usize, data: T) {
    let len = soa.len();
    soa.swap_remove(index);
    soa.push(data);
    soa.swap(index, len - 1);
}