proteus_lib/container/

prot.rs

//! Container model and play settings parsing for `.prot`/`.mka`.

use matroska::Matroska;
use rand::Rng;
use std::collections::HashMap;

use log::{error, info, warn};

use crate::container::info::*;
use crate::container::play_settings::{PlaySettingsFile, PlaySettingsLegacy, SettingsTrack};
use crate::dsp::effects::convolution_reverb::{
    parse_impulse_response_spec, parse_impulse_response_tail_db, ImpulseResponseSpec,
};
use crate::dsp::effects::AudioEffect;

/// Parsed `.prot` container with resolved tracks and playback metadata.
#[derive(Debug, Clone)]
pub struct Prot {
    pub info: Info,
    file_path: Option<String>,
    file_paths: Option<Vec<PathsTrack>>,
    file_paths_dictionary: Option<Vec<String>>,
    track_ids: Option<Vec<u32>>,
    track_paths: Option<Vec<String>>,
    duration: f64,
    play_settings: Option<PlaySettingsFile>,
    impulse_response_spec: Option<ImpulseResponseSpec>,
    impulse_response_tail_db: Option<f32>,
    effects: Option<Vec<AudioEffect>>,
}

impl Prot {
    /// Load a single container file and resolve tracks.
    pub fn new(file_path: &String) -> Self {
        let info = Info::new(file_path.clone());

        println!("Info: {:?}", info);

        let mut this = Self {
            info,
            file_path: Some(file_path.clone()),
            file_paths: None,
            file_paths_dictionary: None,
            track_ids: None,
            track_paths: None,
            duration: 0.0,
            play_settings: None,
            impulse_response_spec: None,
            impulse_response_tail_db: None,
            effects: None,
        };

        this.load_play_settings();
        this.refresh_tracks();

        this
    }

    /// Build a container from multiple standalone file path sets.
    pub fn new_from_file_paths(file_paths: Vec<PathsTrack>) -> Self {
        let mut file_paths_dictionary = Vec::new();
        // Add all file paths to file_paths_dictionary
        // but do not add duplicates
        for track in file_paths.clone() {
            for path in &track.file_paths {
                if !file_paths_dictionary.contains(path) {
                    file_paths_dictionary.push(path.clone());
                }
            }
        }

        let info = Info::new_from_file_paths(file_paths_dictionary.clone());

        let mut this = Self {
            info,
            file_path: None,
            file_paths: Some(file_paths),
            file_paths_dictionary: Some(file_paths_dictionary),
            track_ids: None,
            track_paths: None,
            duration: 0.0,
            play_settings: None,
            impulse_response_spec: None,
            impulse_response_tail_db: None,
            effects: None,
        };

        this.refresh_tracks();

        this
    }

    /// Legacy constructor for backwards compatibility.
    pub fn new_from_file_paths_legacy(file_paths: &Vec<Vec<String>>) -> Self {
        let mut paths_track_list = Vec::new();
        for track in file_paths {
            paths_track_list.push(PathsTrack::new_from_file_paths(track.clone()));
        }
        Self::new_from_file_paths(paths_track_list)
    }

    // fn get_duration_from_file_path(file_path: &String) -> f64 {
    //     let file = std::fs::File::open(file_path).unwrap();
    //     let symphonia: Symphonia = Symphonia::open(file).expect("Could not open file");
    // }

    /// Rebuild the active track list (e.g., after shuffle).
    pub fn refresh_tracks(&mut self) {
        let mut longest_duration = 0.0;

        if let Some(file_paths) = &self.file_paths {
            // Choose random file path from each file_paths array
            let mut track_paths: Vec<String> = Vec::new();
            for track in file_paths {
                if track.file_paths.is_empty() {
                    continue;
                }
                let selections = track.selections_count as usize;
                if selections == 0 {
                    continue;
                }
                for _ in 0..selections {
                    let random_number = rand::thread_rng().gen_range(0..track.file_paths.len());
                    let track_path = track.file_paths[random_number].clone();

                    let index_in_dictionary = self
                        .file_paths_dictionary
                        .as_ref()
                        .unwrap()
                        .iter()
                        .position(|x| *x == track_path)
                        .unwrap();
                    let duration = self.info.get_duration(index_in_dictionary as u32).unwrap();

                    if duration > longest_duration {
                        longest_duration = duration;
                        self.duration = longest_duration;
                    }

                    track_paths.push(track_path);
                }
            }

            self.track_paths = Some(track_paths);

            return;
        }

        if !self.file_path.is_some() {
            return;
        }

        let mut track_index_array: Vec<u32> = Vec::new();
        match self.play_settings.as_ref() {
            Some(play_settings) => match play_settings {
                PlaySettingsFile::Legacy(file) => {
                    collect_legacy_tracks(
                        file.settings.inner(),
                        &mut track_index_array,
                        &mut longest_duration,
                        &self.info,
                        &mut self.duration,
                    );
                }
                PlaySettingsFile::V1(file) => {
                    collect_tracks_from_ids(
                        &file.settings.inner().tracks,
                        &mut track_index_array,
                        &mut longest_duration,
                        &self.info,
                        &mut self.duration,
                    );
                }
                PlaySettingsFile::V2(file) => {
                    collect_tracks_from_ids(
                        &file.settings.inner().tracks,
                        &mut track_index_array,
                        &mut longest_duration,
                        &self.info,
                        &mut self.duration,
                    );
                }
                PlaySettingsFile::V3(file) => {
                    collect_tracks_from_ids(
                        &file.settings.inner().tracks,
                        &mut track_index_array,
                        &mut longest_duration,
                        &self.info,
                        &mut self.duration,
                    );
                }
                PlaySettingsFile::Unknown { .. } => {
                    error!("Unknown file format");
                }
            },
            None => {
                warn!("No play_settings.json found; no tracks resolved.");
            }
        }

        self.track_ids = Some(track_index_array);
    }

    /// Return effects parsed from play_settings, if any.
    pub fn get_effects(&self) -> Option<Vec<AudioEffect>> {
        self.effects.clone()
    }

    fn load_play_settings(&mut self) {
        println!("Loading play settings...");
        let Some(file_path) = self.file_path.as_ref() else {
            return;
        };

        let file = std::fs::File::open(file_path).unwrap();
        let mka: Matroska = Matroska::open(file).expect("Could not open file");

        let mut parsed = None;

        for attachment in &mka.attachments {
            if attachment.name == "play_settings.json" {
                match serde_json::from_slice::<PlaySettingsFile>(&attachment.data) {
                    Ok(play_settings) => {
                        parsed = Some(play_settings);
                        break;
                    }
                    Err(err) => {
                        error!("Failed to parse play_settings.json: {}", err);
                    }
                }
            }
        }

        let Some(play_settings) = parsed else {
            return;
        };

        info!("Parsed play_settings.json");

        self.impulse_response_spec = parse_impulse_response_spec(&play_settings);
        self.impulse_response_tail_db = parse_impulse_response_tail_db(&play_settings);

        match &play_settings {
            PlaySettingsFile::V1(file) => {
                self.effects = Some(file.settings.inner().effects.clone());
            }
            PlaySettingsFile::V2(file) => {
                self.effects = Some(file.settings.inner().effects.clone());
            }
            PlaySettingsFile::V3(file) => {
                self.effects = Some(file.settings.inner().effects.clone());
            }
            _ => {}
        }

        if let Some(effects) = self.effects.as_ref() {
            info!(
                "Loaded play_settings effects ({}): {:?}",
                effects.len(),
                effects
            );
        }

        self.play_settings = Some(play_settings);
    }

    /// Get the convolution impulse response spec, if configured.
    pub fn get_impulse_response_spec(&self) -> Option<ImpulseResponseSpec> {
        self.impulse_response_spec.clone()
    }

    /// Get the configured impulse response tail trim in dB, if any.
    pub fn get_impulse_response_tail_db(&self) -> Option<f32> {
        self.impulse_response_tail_db
    }

    /// Return the container path if this is a `.prot`/`.mka` file.
    pub fn get_container_path(&self) -> Option<String> {
        self.file_path.clone()
    }

    /// Override the impulse response spec at runtime.
    pub fn set_impulse_response_spec(&mut self, spec: ImpulseResponseSpec) {
        self.impulse_response_spec = Some(spec);
    }

    /// Override the impulse response tail trim at runtime.
    pub fn set_impulse_response_tail_db(&mut self, tail_db: f32) {
        self.impulse_response_tail_db = Some(tail_db);
    }

    /// Return per-track keys for UI selection.
    pub fn get_keys(&self) -> Vec<u32> {
        // This should just be a range from 0 to the length of the track_paths or track_ids array
        if let Some(track_paths) = &self.track_paths {
            return (0..track_paths.len() as u32).collect();
        }

        if let Some(track_ids) = &self.track_ids {
            return (0..track_ids.len() as u32).collect();
        }

        Vec::new()
    }

    /// Return per-track identifiers or file paths for display.
    pub fn get_ids(&self) -> Vec<String> {
        if let Some(track_paths) = &self.track_paths {
            return track_paths.clone();
        }

        if let Some(track_ids) = &self.track_ids {
            return track_ids.into_iter().map(|id| format!("{}", id)).collect();
        }

        Vec::new()
    }

    /// Return a list of `(key, path, optional track_id)` for buffering.
    pub fn enumerated_list(&self) -> Vec<(u16, String, Option<u32>)> {
        let mut list: Vec<(u16, String, Option<u32>)> = Vec::new();
        if let Some(track_paths) = &self.track_paths {
            for (index, file_path) in track_paths.iter().enumerate() {
                list.push((index as u16, String::from(file_path), None));
            }

            return list;
        }

        if let Some(track_ids) = &self.track_ids {
            for (index, track_id) in track_ids.iter().enumerate() {
                list.push((
                    index as u16,
                    String::from(self.file_path.as_ref().unwrap()),
                    Some(*track_id),
                ));
            }

            return list;
        }

        list
    }

    /// Return container track entries for shared container streaming.
    pub fn container_track_entries(&self) -> Option<(String, Vec<(u16, u32)>)> {
        let file_path = self.file_path.as_ref()?;
        let track_ids = self.track_ids.as_ref()?;
        let mut entries = Vec::new();
        for (index, track_id) in track_ids.iter().enumerate() {
            entries.push((index as u16, *track_id));
        }
        Some((file_path.clone(), entries))
    }

    /// Get the longest selected duration (seconds).
    pub fn get_duration(&self) -> &f64 {
        &self.duration
    }

    /// Return per-track `(level, pan)` settings keyed by track key.
    pub fn get_track_mix_settings(&self) -> HashMap<u16, (f32, f32)> {
        let mut settings = HashMap::new();

        let tracks = match self.play_settings.as_ref() {
            Some(PlaySettingsFile::V1(file)) => Some(&file.settings.inner().tracks),
            Some(PlaySettingsFile::V2(file)) => Some(&file.settings.inner().tracks),
            _ => None,
        };

        if let Some(tracks) = tracks {
            for (index, track) in tracks.iter().enumerate() {
                settings.insert(index as u16, (track.level, track.pan));
            }
        }

        settings
    }

    /// Return the number of selected tracks.
    pub fn get_length(&self) -> usize {
        if let Some(track_paths) = &self.track_paths {
            return track_paths.len();
        }

        if let Some(file_paths) = &self.file_paths {
            return file_paths.len();
        }

        if let Some(track_ids) = &self.track_ids {
            return track_ids.len();
        }

        0
    }

    /// Return the number of possible unique selections based on track settings.
    pub fn count_possible_combinations(&self) -> Option<u128> {
        if let Some(file_paths) = &self.file_paths {
            return count_paths_track_combinations(file_paths);
        }

        let play_settings = self.play_settings.as_ref()?;
        match play_settings {
            PlaySettingsFile::Legacy(file) => {
                count_legacy_track_combinations(file.settings.inner())
            }
            PlaySettingsFile::V1(file) => {
                count_settings_track_combinations(&file.settings.inner().tracks)
            }
            PlaySettingsFile::V2(file) => {
                count_settings_track_combinations(&file.settings.inner().tracks)
            }
            PlaySettingsFile::V3(file) => {
                count_settings_track_combinations(&file.settings.inner().tracks)
            }
            PlaySettingsFile::Unknown { .. } => None,
        }
    }

    /// Return the unique file paths used for a multi-file container.
    pub fn get_file_paths_dictionary(&self) -> Vec<String> {
        match &self.file_paths_dictionary {
            Some(dictionary) => dictionary.to_vec(),
            None => Vec::new(),
        }
    }
}

/// Standalone file-path track configuration.
#[derive(Debug, Clone)]
pub struct PathsTrack {
    /// Candidate file paths for this track.
    pub file_paths: Vec<String>,
    /// Track gain scalar.
    pub level: f32,
    /// Track pan position.
    pub pan: f32,
    /// Number of selections to pick per refresh.
    pub selections_count: u32,
}

fn count_settings_track_combinations(tracks: &[SettingsTrack]) -> Option<u128> {
    let mut total: u128 = 1;
    for track in tracks {
        let choices = track.ids.len() as u128;
        let selections = track.selections_count;
        let count = checked_pow(choices, selections)?;
        total = total.checked_mul(count)?;
    }
    Some(total)
}

fn count_paths_track_combinations(tracks: &[PathsTrack]) -> Option<u128> {
    let mut total: u128 = 1;
    for track in tracks {
        let choices = track.file_paths.len() as u128;
        let selections = track.selections_count;
        let count = checked_pow(choices, selections)?;
        total = total.checked_mul(count)?;
    }
    Some(total)
}

fn count_legacy_track_combinations(settings: &PlaySettingsLegacy) -> Option<u128> {
    let mut total: u128 = 1;
    for track in &settings.tracks {
        let choices = track.length.unwrap_or(0) as u128;
        let count = checked_pow(choices, 1)?;
        total = total.checked_mul(count)?;
    }
    Some(total)
}

fn checked_pow(base: u128, exp: u32) -> Option<u128> {
    if exp == 0 {
        return Some(1);
    }
    if base == 0 {
        return Some(1);
    }
    let mut result: u128 = 1;
    for _ in 0..exp {
        result = result.checked_mul(base)?;
    }
    Some(result)
}

impl PathsTrack {
    /// Create a new PathsTrack from a vector of file paths.
    pub fn new_from_file_paths(file_paths: Vec<String>) -> Self {
        PathsTrack {
            file_paths,
            level: 1.0,
            pan: 0.0,
            selections_count: 1,
        }
    }
}

fn collect_tracks_from_ids(
    tracks: &[SettingsTrack],
    track_index_array: &mut Vec<u32>,
    longest_duration: &mut f64,
    info: &Info,
    total_duration: &mut f64,
) {
    for track in tracks {
        if track.ids.is_empty() {
            continue;
        }
        let selections = track.selections_count as usize;
        if selections == 0 {
            continue;
        }
        for _ in 0..selections {
            let random_number = rand::thread_rng().gen_range(0..track.ids.len());
            let index = track.ids[random_number];
            if let Some(track_duration) = info.get_duration(index) {
                if track_duration > *longest_duration {
                    *longest_duration = track_duration;
                    *total_duration = *longest_duration;
                }
            }
            track_index_array.push(index);
        }
    }
}

fn collect_legacy_tracks(
    settings: &PlaySettingsLegacy,
    track_index_array: &mut Vec<u32>,
    longest_duration: &mut f64,
    info: &Info,
    total_duration: &mut f64,
) {
    for track in &settings.tracks {
        let (Some(starting_index), Some(length)) = (track.starting_index, track.length) else {
            continue;
        };
        let starting_index = starting_index + 1;
        let index = rand::thread_rng().gen_range(starting_index..(starting_index + length));
        if let Some(track_duration) = info.get_duration(index) {
            if track_duration > *longest_duration {
                *longest_duration = track_duration;
                *total_duration = *longest_duration;
            }
        }
        track_index_array.push(index);
    }
}