//! The labelled corpus: folder-to-tag mapping, and the scratch vault built from it.
//!
//! Two modes need the same thing before they can do anything: every file under
//! `training/` imported, analysed, and tagged from the folder it came in, so
//! `sample_features` holds a vector next to a label. `afcl` exports that state as
//! a layer; `layer-eval` cross-validates over it. The step is expensive enough
//! (decode plus full analysis per file) that it is worth sharing, and subtle
//! enough (a fresh vault every run, a fatal unmapped folder) that two copies
//! would drift.

use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::time::Instant;

use audiofiles_core::analysis::{self, config::AnalysisConfig};
use audiofiles_core::db::Database;
use audiofiles_core::starter_rules::DRUM_CLASSES;
use audiofiles_core::store::SampleStore;
use audiofiles_core::tags;
use rayon::prelude::*;

/// Corpus subdirectory holding the labelled one-shots, one folder per class.
pub(crate) const TRAINING_SUBDIR: &str = "training";

/// Map a corpus folder name to its canonical tag.
///
/// The folder names and [`DRUM_CLASSES`] are the same taxonomy, so this resolves
/// against that table rather than carrying a second copy of it. The only mismatch
/// is punctuation: the table labels the class `hi-hat` and the corpus folder is
/// `hihat`, so both sides are compared with `-` and spaces removed.
pub(crate) fn tag_for_folder(folder: &str) -> Option<&'static str> {
    fn squash(s: &str) -> String {
        s.chars()
            .filter(|c| !matches!(c, '-' | ' ' | '_'))
            .flat_map(char::to_lowercase)
            .collect()
    }
    let want = squash(folder);
    DRUM_CLASSES
        .iter()
        .find(|c| squash(c.label) == want)
        .map(|c| c.tag)
}

/// The short class label for a canonical tag (`instrument.drum.kick` -> `kick`).
///
/// Report tables are unreadable at full tag width once there is a column per
/// class, and the confusion matrix is square in the number of classes.
pub(crate) fn label_for_tag(tag: &str) -> &str {
    DRUM_CLASSES
        .iter()
        .find(|c| c.tag == tag)
        .map_or(tag, |c| c.label)
}

/// Collect the labelled files as (path, tag) pairs, grouped for a stable report.
pub(crate) fn collect_labelled(
    training: &Path,
) -> Result<BTreeMap<&'static str, Vec<PathBuf>>, String> {
    let entries = std::fs::read_dir(training)
        .map_err(|e| format!("cannot read {}: {e}", training.display()))?;

    let mut by_tag: BTreeMap<&'static str, Vec<PathBuf>> = BTreeMap::new();
    let mut unmapped = Vec::new();
    for entry in entries.flatten() {
        if !entry.path().is_dir() {
            continue;
        }
        let folder = entry.file_name().to_string_lossy().to_string();
        let Some(tag) = tag_for_folder(&folder) else {
            unmapped.push(folder);
            continue;
        };
        let mut files: Vec<PathBuf> = std::fs::read_dir(entry.path())
            .map_err(|e| format!("cannot read {}: {e}", entry.path().display()))?
            .flatten()
            .map(|f| f.path())
            .filter(|p| p.is_file())
            .collect();
        // Deterministic order so two runs over the same corpus produce the same
        // layer, which is what makes the checked-in artifact reviewable.
        files.sort();
        by_tag.entry(tag).or_default().extend(files);
    }

    if !unmapped.is_empty() {
        // Loud rather than silent: a folder nobody mapped is a class silently
        // missing from the shipped layer, which reads downstream as the
        // classifier being bad at that class rather than never having seen it.
        unmapped.sort();
        return Err(format!(
            "no tag mapping for corpus folder(s): {}. Add them to DRUM_CLASSES or move them out of {}",
            unmapped.join(", "),
            training.display()
        ));
    }
    if by_tag.is_empty() {
        return Err(format!(
            "no labelled class folders under {}",
            training.display()
        ));
    }
    Ok(by_tag)
}

/// A scratch vault holding the analysed, tagged corpus.
///
/// Just the database. It used to carry the per-class file lists and the analysed
/// count as well, for the `afcl` generator's manifest; that generator went with
/// the bundled layer (2026-08-08, wiki `af-likeness-web`) and both meters read
/// the corpus back out of `db` rather than off this struct.
pub(crate) struct LabelledVault {
    pub(crate) db: Database,
}

/// Import, analyse and tag every labelled file into a fresh vault at `vault`.
///
/// Prints its own progress: both callers are long-running terminal modes and
/// this is the slow part of each.
pub(crate) fn build_vault(
    corpus: &Path,
    vault: &Path,
    config: &AnalysisConfig,
) -> Result<LabelledVault, String> {
    let by_tag = collect_labelled(&corpus.join(TRAINING_SUBDIR))?;

    let total: usize = by_tag.values().map(Vec::len).sum();
    println!(
        "  {} labelled file(s) across {} class(es):",
        total,
        by_tag.len()
    );
    for (tag, files) in &by_tag {
        println!("    {:<28} {:>5}", tag, files.len());
    }
    println!();

    // A fresh vault every run. What is built from it is a pure function of the
    // corpus, and leftovers from a previous run would silently widen it.
    if vault.exists() {
        std::fs::remove_dir_all(vault)
            .map_err(|e| format!("could not clear scratch vault {}: {e}", vault.display()))?;
    }
    let samples_dir = vault.join("samples");
    std::fs::create_dir_all(&samples_dir)
        .map_err(|e| format!("could not create {}: {e}", samples_dir.display()))?;

    let db = Database::open(vault.join("audiofiles.db"))
        .map_err(|e| format!("Database::open failed: {e}"))?;
    let store =
        SampleStore::new(&samples_dir).map_err(|e| format!("SampleStore::new failed: {e}"))?;

    // Import and tag. No VFS nodes: nothing here runs a UI query, and neither
    // the export nor the eval needs one; both read `sample_features` joined to
    // `tags`.
    let start = Instant::now();
    let mut to_analyze: Vec<(String, PathBuf)> = Vec::with_capacity(total);
    let mut import_failures = 0usize;
    for (tag, files) in &by_tag {
        for path in files {
            match store.import(path, &db) {
                Ok(hash) => {
                    if let Err(e) = tags::add_tag(&db, &hash, tag) {
                        eprintln!("  tag {tag} on {}: {e}", path.display());
                        import_failures += 1;
                        continue;
                    }
                    to_analyze.push((hash, path.clone()));
                }
                Err(e) => {
                    eprintln!("  import {}: {e}", path.display());
                    import_failures += 1;
                }
            }
        }
    }
    println!(
        "  imported {} file(s) in {:.1}s{}",
        to_analyze.len(),
        start.elapsed().as_secs_f64(),
        if import_failures > 0 {
            format!(", {import_failures} failed")
        } else {
            String::new()
        }
    );

    // Analyse everything. The ingest benchmark caps this because analysis is the
    // expensive stage and it only needs a sample; here every vector is the
    // payload, so there is no budget to apply.
    //
    // Parallel, but still deterministic: `analyze_sample` is a pure function of
    // the file, and a rayon `collect` into a Vec restores input order, so the
    // batch written to the DB is the same sequence a serial run would write.
    // That matters because the exported layer is a checked-in artifact.
    let start = Instant::now();
    let results: Vec<_> = to_analyze
        .par_iter()
        .filter_map(|(hash, path)| analysis::analyze_sample(hash, path, config).ok())
        .collect();
    let analyzed = results.len();
    analysis::save_analysis_batch(&db, &results)
        .map_err(|e| format!("save_analysis_batch failed: {e}"))?;
    println!(
        "  analysed {} file(s) in {:.1}s{}",
        analyzed,
        start.elapsed().as_secs_f64(),
        if analyzed < to_analyze.len() {
            format!(", {} failed to analyse", to_analyze.len() - analyzed)
        } else {
            String::new()
        }
    );

    Ok(LabelledVault { db })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn folder_names_map_to_canonical_tags() {
        assert_eq!(tag_for_folder("kick"), Some("instrument.drum.kick"));
        assert_eq!(tag_for_folder("snare"), Some("instrument.drum.snare"));
        assert_eq!(tag_for_folder("cymbal"), Some("instrument.drum.cymbal"));
        assert_eq!(tag_for_folder("clap"), Some("instrument.drum.clap"));
        assert_eq!(tag_for_folder("tom"), Some("instrument.drum.tom"));
        assert_eq!(tag_for_folder("percussion"), Some("instrument.percussion"));
    }

    #[test]
    fn hihat_folder_matches_the_hi_hat_label() {
        // The one place the corpus and DRUM_CLASSES disagree on spelling.
        assert_eq!(tag_for_folder("hihat"), Some("instrument.drum.hihat"));
        assert_eq!(tag_for_folder("hi-hat"), Some("instrument.drum.hihat"));
        assert_eq!(tag_for_folder("Hi Hat"), Some("instrument.drum.hihat"));
    }

    #[test]
    fn unknown_folder_has_no_tag() {
        assert_eq!(tag_for_folder("bass"), None);
        assert_eq!(tag_for_folder(""), None);
    }

    #[test]
    fn every_drum_class_is_reachable_from_some_folder_name() {
        // Guards the mapping against a DRUM_CLASSES entry whose label stops
        // resolving; without this, a renamed label silently drops a class.
        for class in DRUM_CLASSES {
            assert_eq!(
                tag_for_folder(class.label),
                Some(class.tag),
                "class {} no longer resolves from its own label",
                class.label
            );
        }
    }

    #[test]
    fn labels_round_trip_from_their_tag() {
        for class in DRUM_CLASSES {
            assert_eq!(label_for_tag(class.tag), class.label);
        }
        // An unknown tag prints as itself rather than vanishing.
        assert_eq!(label_for_tag("instrument.bass"), "instrument.bass");
    }

    #[test]
    fn collect_labelled_rejects_an_unmapped_folder() {
        let dir = tempfile::tempdir().unwrap();
        let training = dir.path().join(TRAINING_SUBDIR);
        std::fs::create_dir_all(training.join("kick")).unwrap();
        std::fs::create_dir_all(training.join("didgeridoo")).unwrap();
        let err = collect_labelled(&training).unwrap_err();
        assert!(err.contains("didgeridoo"), "{err}");
    }

    #[test]
    fn collect_labelled_groups_files_under_their_tag() {
        let dir = tempfile::tempdir().unwrap();
        let training = dir.path().join(TRAINING_SUBDIR);
        std::fs::create_dir_all(training.join("kick")).unwrap();
        std::fs::create_dir_all(training.join("snare")).unwrap();
        std::fs::write(training.join("kick/b.wav"), b"x").unwrap();
        std::fs::write(training.join("kick/a.wav"), b"x").unwrap();
        std::fs::write(training.join("snare/c.wav"), b"x").unwrap();

        let got = collect_labelled(&training).unwrap();
        assert_eq!(got["instrument.drum.kick"].len(), 2);
        assert_eq!(got["instrument.drum.snare"].len(), 1);
        // Sorted, so the artifact is reproducible across runs.
        assert!(got["instrument.drum.kick"][0].ends_with("a.wav"));
    }
}