Cagire/src/state/sample_browser.rs

use std::fs;
use std::path::{Path, PathBuf};

use cagire_ratatui::{TreeLine, TreeLineKind};

const AUDIO_EXTENSIONS: &[&str] = &["wav", "flac", "ogg", "aiff", "aif", "mp3"];

pub enum SampleNode {
    Root {
        label: String,
        children: Vec<SampleNode>,
        expanded: bool,
    },
    Folder {
        name: String,
        children: Vec<SampleNode>,
        expanded: bool,
    },
    File {
        name: String,
    },
}

impl SampleNode {
    fn expanded(&self) -> bool {
        match self {
            SampleNode::Root { expanded, .. } | SampleNode::Folder { expanded, .. } => *expanded,
            SampleNode::File { .. } => false,
        }
    }

    fn set_expanded(&mut self, val: bool) {
        match self {
            SampleNode::Root { expanded, .. } | SampleNode::Folder { expanded, .. } => {
                *expanded = val;
            }
            SampleNode::File { .. } => {}
        }
    }

    fn is_expandable(&self) -> bool {
        !matches!(self, SampleNode::File { .. })
    }

    fn children(&self) -> &[SampleNode] {
        match self {
            SampleNode::Root { children, .. } | SampleNode::Folder { children, .. } => children,
            SampleNode::File { .. } => &[],
        }
    }

    fn label(&self) -> &str {
        match self {
            SampleNode::Root { label, .. } => label,
            SampleNode::Folder { name, .. } | SampleNode::File { name, .. } => name,
        }
    }

    fn flatten(&self, depth: u8, parent_folder: &str, file_index: usize, out: &mut Vec<TreeLine>) {
        let kind = match self {
            SampleNode::Root { expanded, .. } => TreeLineKind::Root { expanded: *expanded },
            SampleNode::Folder { expanded, .. } => TreeLineKind::Folder { expanded: *expanded },
            SampleNode::File { .. } => TreeLineKind::File,
        };
        out.push(TreeLine {
            depth,
            kind,
            label: self.label().to_string(),
            folder: parent_folder.to_string(),
            index: file_index,
        });
        if self.expanded() {
            let folder_name = self.label();
            let mut idx = 0;
            for child in self.children() {
                let child_idx = if matches!(child, SampleNode::File { .. }) {
                    let i = idx;
                    idx += 1;
                    i
                } else {
                    0
                };
                child.flatten(depth + 1, folder_name, child_idx, out);
            }
        }
    }
}

pub struct SampleTree {
    roots: Vec<SampleNode>,
}

impl SampleTree {
    pub fn from_paths(paths: &[PathBuf]) -> Self {
        if paths.len() == 1 {
            // Single path: show its contents directly at root level
            let roots = Self::scan_children(&paths[0]);
            return Self { roots };
        }
        let mut roots = Vec::new();
        for path in paths {
            if let Some(root) = Self::scan_root(path) {
                roots.push(root);
            }
        }
        Self { roots }
    }

    fn scan_children(path: &Path) -> Vec<SampleNode> {
        let entries = match fs::read_dir(path) {
            Ok(e) => e,
            Err(_) => return Vec::new(),
        };

        let mut files: Vec<String> = Vec::new();
        let mut folders: Vec<(String, PathBuf)> = Vec::new();

        for entry in entries.flatten() {
            let ft = entry.file_type().ok();
            let name = entry.file_name().to_string_lossy().into_owned();
            if ft.is_some_and(|t| t.is_dir()) {
                folders.push((name, entry.path()));
            } else if is_audio_file(&name) {
                files.push(name);
            }
        }

        folders.sort_by_key(|a| a.0.to_lowercase());
        files.sort_by_key(|a| a.to_lowercase());

        let mut nodes = Vec::new();
        for (name, folder_path) in folders {
            if let Some(folder) = Self::scan_folder(&name, &folder_path) {
                nodes.push(folder);
            }
        }
        for name in files {
            nodes.push(SampleNode::File { name });
        }
        nodes
    }

    fn scan_root(path: &Path) -> Option<SampleNode> {
        let entries = fs::read_dir(path).ok()?;
        let label = path
            .file_name()
            .map(|n| n.to_string_lossy().into_owned())
            .unwrap_or_else(|| path.display().to_string());

        let mut children = Vec::new();
        let mut files: Vec<String> = Vec::new();
        let mut folders: Vec<(String, PathBuf)> = Vec::new();

        for entry in entries.flatten() {
            let ft = entry.file_type().ok();
            let name = entry.file_name().to_string_lossy().into_owned();
            if ft.is_some_and(|t| t.is_dir()) {
                folders.push((name, entry.path()));
            } else if is_audio_file(&name) {
                files.push(name);
            }
        }

        folders.sort_by_key(|a| a.0.to_lowercase());
        files.sort_by_key(|a| a.to_lowercase());

        for (name, folder_path) in folders {
            if let Some(folder) = Self::scan_folder(&name, &folder_path) {
                children.push(folder);
            }
        }
        for name in files {
            children.push(SampleNode::File { name });
        }

        Some(SampleNode::Root {
            label,
            children,
            expanded: false,
        })
    }

    fn scan_folder(name: &str, path: &Path) -> Option<SampleNode> {
        let entries = fs::read_dir(path).ok()?;
        let mut children: Vec<SampleNode> = Vec::new();

        let mut files: Vec<String> = Vec::new();
        for entry in entries.flatten() {
            let ft = entry.file_type().ok();
            let entry_name = entry.file_name().to_string_lossy().into_owned();
            if ft.is_some_and(|t| t.is_file()) && is_audio_file(&entry_name) {
                files.push(entry_name);
            }
        }
        files.sort_by_key(|a| a.to_lowercase());

        for file in files {
            children.push(SampleNode::File { name: file });
        }

        if children.is_empty() {
            return None;
        }

        Some(SampleNode::Folder {
            name: name.to_string(),
            children,
            expanded: false,
        })
    }

    pub fn visible_entries(&self) -> Vec<TreeLine> {
        let mut out = Vec::new();
        for root in &self.roots {
            root.flatten(0, "", 0, &mut out);
        }
        out
    }

    fn node_at_mut(&mut self, visible_index: usize) -> Option<&mut SampleNode> {
        let mut count = 0;
        for root in &mut self.roots {
            if let Some(node) = Self::walk_mut(root, visible_index, &mut count) {
                return Some(node);
            }
        }
        None
    }

    fn walk_mut<'a>(
        node: &'a mut SampleNode,
        target: usize,
        count: &mut usize,
    ) -> Option<&'a mut SampleNode> {
        if *count == target {
            return Some(node);
        }
        *count += 1;
        if node.expanded() {
            let children = match node {
                SampleNode::Root { children, .. } | SampleNode::Folder { children, .. } => {
                    children
                }
                SampleNode::File { .. } => return None,
            };
            for child in children.iter_mut() {
                if let Some(found) = Self::walk_mut(child, target, count) {
                    return Some(found);
                }
            }
        }
        None
    }
}

pub struct SampleBrowserState {
    pub tree: SampleTree,
    pub cursor: usize,
    pub scroll_offset: usize,
    pub search_query: String,
    pub search_active: bool,
    filtered: Option<Vec<TreeLine>>,
}

impl SampleBrowserState {
    pub fn new(paths: &[PathBuf]) -> Self {
        Self {
            tree: SampleTree::from_paths(paths),
            cursor: 0,
            scroll_offset: 0,
            search_query: String::new(),
            search_active: false,
            filtered: None,
        }
    }

    pub fn entries(&self) -> Vec<TreeLine> {
        if let Some(ref filtered) = self.filtered {
            return filtered.clone();
        }
        self.tree.visible_entries()
    }

    pub fn current_entry(&self) -> Option<TreeLine> {
        let entries = self.entries();
        entries.into_iter().nth(self.cursor)
    }

    pub fn visible_count(&self) -> usize {
        if let Some(ref filtered) = self.filtered {
            return filtered.len();
        }
        self.tree.visible_entries().len()
    }

    pub fn move_up(&mut self) {
        if self.cursor > 0 {
            self.cursor -= 1;
            if self.cursor < self.scroll_offset {
                self.scroll_offset = self.cursor;
            }
        }
    }

    pub fn move_down(&mut self, visible_height: usize) {
        let count = self.visible_count();
        if count == 0 {
            return;
        }
        if self.cursor + 1 < count {
            self.cursor += 1;
            if self.cursor >= self.scroll_offset + visible_height {
                self.scroll_offset = self.cursor - visible_height + 1;
            }
        }
    }

    pub fn toggle_expand(&mut self) {
        if self.filtered.is_some() {
            return;
        }
        if let Some(node) = self.tree.node_at_mut(self.cursor) {
            if node.is_expandable() {
                let new_val = !node.expanded();
                node.set_expanded(new_val);
            }
        }
    }

    pub fn collapse_at_cursor(&mut self) {
        if self.filtered.is_some() {
            return;
        }
        if let Some(node) = self.tree.node_at_mut(self.cursor) {
            if node.expanded() {
                node.set_expanded(false);
            }
        }
    }

    pub fn activate_search(&mut self) {
        self.search_active = true;
    }

    pub fn update_search(&mut self) {
        if self.search_query.is_empty() {
            self.filtered = None;
        } else {
            let query = self.search_query.to_lowercase();
            let full = self.full_entries();
            let filtered: Vec<TreeLine> = full
                .into_iter()
                .filter(|line| line.label.to_lowercase().contains(&query))
                .collect();
            self.filtered = Some(filtered);
        }
        self.cursor = 0;
        self.scroll_offset = 0;
    }

    pub fn clear_search(&mut self) {
        self.search_query.clear();
        self.search_active = false;
        self.filtered = None;
        self.cursor = 0;
        self.scroll_offset = 0;
    }

    fn full_entries(&self) -> Vec<TreeLine> {
        let mut out = Vec::new();
        for root in &self.tree.roots {
            Self::flatten_all(root, 0, "", 0, &mut out);
        }
        out
    }

    fn flatten_all(
        node: &SampleNode,
        depth: u8,
        parent_folder: &str,
        file_index: usize,
        out: &mut Vec<TreeLine>,
    ) {
        let kind = match node {
            SampleNode::Root { .. } => TreeLineKind::Root { expanded: true },
            SampleNode::Folder { .. } => TreeLineKind::Folder { expanded: true },
            SampleNode::File { .. } => TreeLineKind::File,
        };
        out.push(TreeLine {
            depth,
            kind,
            label: node.label().to_string(),
            folder: parent_folder.to_string(),
            index: file_index,
        });
        let folder_name = node.label();
        let mut idx = 0;
        for child in node.children() {
            let child_idx = if matches!(child, SampleNode::File { .. }) {
                let i = idx;
                idx += 1;
                i
            } else {
                0
            };
            Self::flatten_all(child, depth + 1, folder_name, child_idx, out);
        }
    }
}

fn is_audio_file(name: &str) -> bool {
    let lower = name.to_lowercase();
    AUDIO_EXTENSIONS.iter().any(|ext| lower.ends_with(ext))
}