zfsdu/src/app.rs

use crate::config::Config;
use crate::zfs::{self, ZfsEntry, ZfsSnapshot};
use anyhow::Result;
use std::collections::HashSet;

#[derive(Debug, Clone)]
pub enum ListItem {
    Dataset(ZfsEntry),
    Snapshot(ZfsSnapshot),
}

impl ListItem {
    pub fn name(&self) -> &str {
        match self {
            ListItem::Dataset(d) => d.short_name(),
            ListItem::Snapshot(s) => s.short_name(),
        }
    }

    pub fn full_name(&self) -> &str {
        match self {
            ListItem::Dataset(d) => &d.full_name,
            ListItem::Snapshot(s) => &s.full_name,
        }
    }

    pub fn used(&self) -> u64 {
        match self {
            ListItem::Dataset(d) => d.used,
            ListItem::Snapshot(s) => s.used,
        }
    }

    pub fn is_snapshot(&self) -> bool {
        matches!(self, ListItem::Snapshot(_))
    }

    pub fn is_dataset(&self) -> bool {
        matches!(self, ListItem::Dataset(_))
    }

    pub fn creation(&self) -> &str {
        match self {
            ListItem::Dataset(d) => &d.creation,
            ListItem::Snapshot(s) => &s.creation,
        }
    }
}

#[derive(Debug, Clone)]
pub struct TreeItem {
    pub item: ListItem,
    pub depth: usize,
    pub is_last: bool,
    pub parent_is_last: Vec<bool>,
    pub visible: bool,
    pub parent_path: String,  // For tracking which parent this belongs to
}

#[derive(Debug, Clone, PartialEq)]
pub enum InputMode {
    Normal,
    Search,
    Filter,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum SortMode {
    Size,
    Date,
    Name,
}

impl SortMode {
    pub fn next(self) -> Self {
        match self {
            SortMode::Size => SortMode::Date,
            SortMode::Date => SortMode::Name,
            SortMode::Name => SortMode::Size,
        }
    }

    pub fn label(&self) -> &'static str {
        match self {
            SortMode::Size => "Size",
            SortMode::Date => "Date",
            SortMode::Name => "Name",
        }
    }
}

pub struct App {
    pub root_entries: Vec<ZfsEntry>,
    pub tree_items: Vec<TreeItem>,
    pub selected: usize,
    pub marked: Vec<String>,
    pub expanded: HashSet<String>,  // Names of expanded datasets
    pub show_help: bool,
    pub message: Option<String>,
    pub confirm_delete: bool,
    pub page_size: usize,
    pub scroll_offset: usize,
    // Search & Filter
    pub input_mode: InputMode,
    pub search_query: String,
    pub filter_query: String,
    pub search_matches: Vec<usize>,
    pub current_match: usize,
    // Sort mode
    pub sort_mode: SortMode,
    // Dry-run mode
    pub dry_run_mode: bool,
    // Auto-refresh
    pub auto_refresh: bool,
    pub last_refresh: std::time::Instant,
    pub auto_refresh_interval: u64,
    // Config
    pub config: Config,
}

impl App {
    pub fn new() -> Result<Self> {
        let config = Config::load();
        let root_entries = zfs::get_zfs_tree()?;
        let expanded = HashSet::new();
        let tree_items = Self::build_tree_view(&root_entries, 0, &[], &expanded, "");
        let auto_refresh_interval = config.auto_refresh_interval;

        Ok(Self {
            root_entries,
            tree_items,
            selected: 0,
            marked: Vec::new(),
            expanded,
            show_help: false,
            message: None,
            confirm_delete: false,
            page_size: 20,
            scroll_offset: 0,
            input_mode: InputMode::Normal,
            search_query: String::new(),
            filter_query: String::new(),
            search_matches: Vec::new(),
            current_match: 0,
            sort_mode: SortMode::Size,
            dry_run_mode: false,
            auto_refresh: false,
            last_refresh: std::time::Instant::now(),
            auto_refresh_interval,
            config,
        })
    }

    fn build_tree_view(
        entries: &[ZfsEntry],
        depth: usize,
        parent_is_last: &[bool],
        expanded: &HashSet<String>,
        parent_path: &str,
    ) -> Vec<TreeItem> {
        let mut items = Vec::new();
        let len = entries.len();

        for (i, entry) in entries.iter().enumerate() {
            let is_last_entry = i == len - 1;
            let has_children = !entry.children.is_empty() || !entry.snapshots.is_empty();
            let is_expanded = expanded.contains(&entry.full_name);

            items.push(TreeItem {
                item: ListItem::Dataset(entry.clone()),
                depth,
                is_last: is_last_entry && (!is_expanded || !has_children),
                parent_is_last: parent_is_last.to_vec(),
                visible: true,
                parent_path: parent_path.to_string(),
            });

            // Only show children if expanded
            if is_expanded && has_children {
                let mut child_parent_is_last = parent_is_last.to_vec();
                child_parent_is_last.push(is_last_entry);

                // Add child datasets
                let child_items = Self::build_tree_view(
                    &entry.children,
                    depth + 1,
                    &child_parent_is_last,
                    expanded,
                    &entry.full_name,
                );
                items.extend(child_items);

                // Add snapshots (rendered after children, so last snapshot is always last item)
                let snap_len = entry.snapshots.len();
                for (j, snap) in entry.snapshots.iter().enumerate() {
                    items.push(TreeItem {
                        item: ListItem::Snapshot(snap.clone()),
                        depth: depth + 1,
                        is_last: j == snap_len - 1,
                        parent_is_last: child_parent_is_last.clone(),
                        visible: true,
                        parent_path: entry.full_name.clone(),
                    });
                }
            }
        }

        items
    }

    pub fn rebuild_tree(&mut self) {
        self.tree_items = Self::build_tree_view(&self.root_entries, 0, &[], &self.expanded, "");
        self.apply_filter();
        self.refresh_search_matches();
        self.ensure_valid_selection();
    }

    pub fn refresh(&mut self) -> Result<()> {
        self.root_entries = zfs::get_zfs_tree()?;
        self.resort_entries();
        self.rebuild_tree();
        self.last_refresh = std::time::Instant::now();
        self.message = Some("Refreshed".to_string());
        Ok(())
    }

    fn ensure_valid_selection(&mut self) {
        let count = self.visible_count();
        if count == 0 {
            self.selected = 0;
        } else if self.selected >= count {
            self.selected = count - 1;
        }
        self.adjust_scroll();
    }

    pub fn visible_items(&self) -> Vec<(usize, &TreeItem)> {
        self.tree_items
            .iter()
            .enumerate()
            .filter(|(_, ti)| ti.visible)
            .collect()
    }

    pub fn visible_count(&self) -> usize {
        self.tree_items.iter().filter(|ti| ti.visible).count()
    }

    fn nth_visible(&self, n: usize) -> Option<usize> {
        self.tree_items
            .iter()
            .enumerate()
            .filter(|(_, ti)| ti.visible)
            .nth(n)
            .map(|(i, _)| i)
    }

    fn visible_index_of(&self, real_index: usize) -> Option<usize> {
        self.tree_items
            .iter()
            .enumerate()
            .filter(|(_, ti)| ti.visible)
            .position(|(i, _)| i == real_index)
    }

    fn adjust_scroll(&mut self) {
        // Keep selected item visible
        if self.selected < self.scroll_offset {
            self.scroll_offset = self.selected;
        } else if self.selected >= self.scroll_offset + self.page_size {
            self.scroll_offset = self.selected - self.page_size + 1;
        }
    }

    pub fn next(&mut self) {
        let count = self.visible_count();
        if count > 0 && self.selected < count - 1 {
            self.selected += 1;
            self.adjust_scroll();
        }
    }

    pub fn previous(&mut self) {
        if self.selected > 0 {
            self.selected -= 1;
            self.adjust_scroll();
        }
    }

    pub fn first(&mut self) {
        self.selected = 0;
        self.scroll_offset = 0;
    }

    pub fn last(&mut self) {
        let count = self.visible_count();
        if count > 0 {
            self.selected = count - 1;
            self.adjust_scroll();
        }
    }

    pub fn page_up(&mut self) {
        self.selected = self.selected.saturating_sub(self.page_size);
        self.adjust_scroll();
    }

    pub fn page_down(&mut self) {
        let count = self.visible_count();
        if count > 0 {
            self.selected = (self.selected + self.page_size).min(count - 1);
            self.adjust_scroll();
        }
    }

    pub fn toggle_expand(&mut self) {
        if let Some(real_idx) = self.nth_visible(self.selected) {
            if let Some(tree_item) = self.tree_items.get(real_idx) {
                if let ListItem::Dataset(entry) = &tree_item.item {
                    let name = entry.full_name.clone();
                    if self.expanded.contains(&name) {
                        self.expanded.remove(&name);
                    } else {
                        self.expanded.insert(name);
                    }
                    self.rebuild_tree();
                }
            }
        }
    }

    pub fn is_expanded(&self, name: &str) -> bool {
        self.expanded.contains(name)
    }

    pub fn toggle_mark(&mut self) {
        if let Some(real_idx) = self.nth_visible(self.selected) {
            if let Some(tree_item) = self.tree_items.get(real_idx) {
                if tree_item.item.is_snapshot() {
                    let name = tree_item.item.full_name().to_string();
                    if let Some(pos) = self.marked.iter().position(|x| x == &name) {
                        self.marked.remove(pos);
                    } else {
                        self.marked.push(name);
                    }
                }
            }
        }
        self.next();
    }

    pub fn is_marked(&self, name: &str) -> bool {
        self.marked.iter().any(|x| x == name)
    }

    // Reset confirm_delete state — call on any action that isn't 'd'
    pub fn reset_confirm(&mut self) {
        if self.confirm_delete {
            self.confirm_delete = false;
            self.message = Some("Delete cancelled".to_string());
        }
    }

    pub fn delete_marked(&mut self) -> Result<()> {
        if self.marked.is_empty() {
            self.confirm_delete = false;
            self.message = Some("No snapshots marked".to_string());
            return Ok(());
        }

        // In dry-run mode, just show what would happen
        if self.dry_run_mode {
            self.confirm_delete = false;
            self.show_dry_run_preview();
            return Ok(());
        }

        if !self.confirm_delete {
            self.confirm_delete = true;
            let total_size = self.total_marked_size();
            self.message = Some(format!(
                "Delete {} snapshot(s) (~{})? Press 'd' to confirm",
                self.marked.len(),
                zfs::format_size(total_size)
            ));
            return Ok(());
        }

        let mut deleted = 0;
        let mut errors = Vec::new();

        for name in &self.marked.clone() {
            match zfs::delete_snapshot(name) {
                Ok(_) => deleted += 1,
                Err(e) => errors.push(format!("{}: {}", name, e)),
            }
        }

        self.marked.clear();
        self.confirm_delete = false;

        if errors.is_empty() {
            self.message = Some(format!("Deleted {} snapshot(s)", deleted));
        } else {
            self.message = Some(format!(
                "Deleted {}, {} errors",
                deleted,
                errors.len()
            ));
        }

        self.refresh()?;
        Ok(())
    }

    pub fn toggle_help(&mut self) {
        self.show_help = !self.show_help;
    }

    pub fn total_marked_size(&self) -> u64 {
        self.tree_items
            .iter()
            .filter(|ti| self.is_marked(ti.item.full_name()))
            .map(|ti| ti.item.used())
            .sum()
    }

    pub fn max_size(&self) -> u64 {
        self.tree_items
            .iter()
            .filter(|ti| ti.visible)
            .map(|ti| ti.item.used())
            .max()
            .unwrap_or(1)
    }

    // Search & Filter functions
    pub fn start_search(&mut self) {
        self.input_mode = InputMode::Search;
        self.search_query.clear();
        self.search_matches.clear();
        self.current_match = 0;
    }

    pub fn start_filter(&mut self) {
        self.input_mode = InputMode::Filter;
    }

    pub fn cancel_input(&mut self) {
        let was_search = self.input_mode == InputMode::Search;
        self.input_mode = InputMode::Normal;
        if was_search {
            self.search_query.clear();
            self.search_matches.clear();
        }
    }

    pub fn confirm_input(&mut self) {
        match self.input_mode {
            InputMode::Search => {
                self.execute_search();
                self.input_mode = InputMode::Normal;
            }
            InputMode::Filter => {
                self.apply_filter();
                self.input_mode = InputMode::Normal;
            }
            InputMode::Normal => {}
        }
    }

    pub fn input_char(&mut self, c: char) {
        match self.input_mode {
            InputMode::Search => self.search_query.push(c),
            InputMode::Filter => self.filter_query.push(c),
            InputMode::Normal => {}
        }
    }

    pub fn input_backspace(&mut self) {
        match self.input_mode {
            InputMode::Search => { self.search_query.pop(); }
            InputMode::Filter => { self.filter_query.pop(); }
            InputMode::Normal => {}
        }
    }

    fn execute_search(&mut self) {
        self.search_matches.clear();

        if self.search_query.is_empty() {
            return;
        }

        let query = self.search_query.to_lowercase();

        for (i, tree_item) in self.tree_items.iter().enumerate() {
            if tree_item.visible {
                let name = tree_item.item.full_name().to_lowercase();
                if name.contains(&query) {
                    self.search_matches.push(i);
                }
            }
        }

        self.current_match = 0;
        self.jump_to_current_match();

        let count = self.search_matches.len();
        self.message = Some(format!("Found {} match(es)", count));
    }

    pub fn next_match(&mut self) {
        if self.search_matches.is_empty() {
            return;
        }
        self.current_match = (self.current_match + 1) % self.search_matches.len();
        self.jump_to_current_match();
    }

    pub fn prev_match(&mut self) {
        if self.search_matches.is_empty() {
            return;
        }
        if self.current_match == 0 {
            self.current_match = self.search_matches.len() - 1;
        } else {
            self.current_match -= 1;
        }
        self.jump_to_current_match();
    }

    fn jump_to_current_match(&mut self) {
        if let Some(&real_idx) = self.search_matches.get(self.current_match) {
            if let Some(visible_idx) = self.visible_index_of(real_idx) {
                self.selected = visible_idx;
                self.adjust_scroll();
            }
        }
    }

    pub fn apply_filter(&mut self) {
        if self.filter_query.is_empty() {
            for ti in &mut self.tree_items {
                ti.visible = true;
            }
        } else {
            let query = self.filter_query.to_lowercase();

            for ti in &mut self.tree_items {
                let name = ti.item.full_name().to_lowercase();
                ti.visible = name.contains(&query);
            }
        }

        self.ensure_valid_selection();

        let visible = self.visible_count();
        let total = self.tree_items.len();
        if !self.filter_query.is_empty() {
            self.message = Some(format!("Filter: {} of {} shown", visible, total));
        }
    }

    pub fn clear_filter(&mut self) {
        self.filter_query.clear();
        self.apply_filter();
        self.message = Some("Filter cleared".to_string());
    }

    pub fn is_search_match(&self, real_index: usize) -> bool {
        self.search_matches.contains(&real_index)
    }

    fn refresh_search_matches(&mut self) {
        if self.search_query.is_empty() {
            return;
        }
        self.search_matches.clear();
        let query = self.search_query.to_lowercase();
        for (i, tree_item) in self.tree_items.iter().enumerate() {
            if tree_item.visible {
                let name = tree_item.item.full_name().to_lowercase();
                if name.contains(&query) {
                    self.search_matches.push(i);
                }
            }
        }
        if !self.search_matches.is_empty() {
            if self.current_match >= self.search_matches.len() {
                self.current_match = self.search_matches.len() - 1;
            }
        } else {
            self.current_match = 0;
        }
    }

    // Get the scope parent: if cursor is on a dataset, use its full_name;
    // if on a snapshot, use its parent_path (the owning dataset)
    fn current_scope_parent(&self) -> Option<String> {
        if let Some(real_idx) = self.nth_visible(self.selected) {
            self.tree_items.get(real_idx).map(|ti| match &ti.item {
                ListItem::Dataset(d) => d.full_name.clone(),
                ListItem::Snapshot(_) => ti.parent_path.clone(),
            })
        } else {
            None
        }
    }

    // Select/Unselect Group (like MC) - works in current scope
    pub fn select_group(&mut self) {
        let current_parent = self.current_scope_parent();

        let mut count = 0;
        for ti in &self.tree_items {
            if ti.visible && ti.item.is_snapshot() {
                let in_scope = match &current_parent {
                    Some(parent) => &ti.parent_path == parent,
                    None => true,
                };

                if in_scope {
                    let name = ti.item.full_name().to_string();
                    if !self.marked.contains(&name) {
                        self.marked.push(name);
                        count += 1;
                    }
                }
            }
        }
        self.message = Some(format!("Selected {} snapshot(s)", count));
    }

    pub fn unselect_group(&mut self) {
        let current_parent = self.current_scope_parent();

        let visible_names: Vec<String> = self.tree_items
            .iter()
            .filter(|ti| {
                if !ti.visible || !ti.item.is_snapshot() {
                    return false;
                }
                match &current_parent {
                    Some(parent) => &ti.parent_path == parent,
                    None => true,
                }
            })
            .map(|ti| ti.item.full_name().to_string())
            .collect();

        let before = self.marked.len();
        self.marked.retain(|name| !visible_names.contains(name));
        let removed = before - self.marked.len();
        self.message = Some(format!("Unselected {} snapshot(s)", removed));
    }

    pub fn invert_selection(&mut self) {
        let current_parent = self.current_scope_parent();

        let mut count = 0;
        for ti in &self.tree_items {
            // Only invert in current scope (same parent)
            if ti.visible && ti.item.is_snapshot() {
                let in_scope = match &current_parent {
                    Some(parent) => &ti.parent_path == parent,
                    None => true, // If no parent, work on all visible
                };

                if in_scope {
                    let name = ti.item.full_name().to_string();
                    if let Some(pos) = self.marked.iter().position(|x| x == &name) {
                        self.marked.remove(pos);
                    } else {
                        self.marked.push(name);
                        count += 1;
                    }
                }
            }
        }
        self.message = Some(format!("Inverted {} in scope", count));
    }

    // Expand all datasets (recursively from root_entries, not just visible tree)
    pub fn expand_all(&mut self) {
        Self::collect_expandable(&self.root_entries, &mut self.expanded);
        self.rebuild_tree();
        self.message = Some("Expanded all".to_string());
    }

    fn collect_expandable(entries: &[ZfsEntry], expanded: &mut HashSet<String>) {
        for entry in entries {
            if entry.has_children() {
                expanded.insert(entry.full_name.clone());
            }
            Self::collect_expandable(&entry.children, expanded);
        }
    }

    // Collapse all datasets
    pub fn collapse_all(&mut self) {
        self.expanded.clear();
        self.rebuild_tree();
        self.message = Some("Collapsed all".to_string());
    }

    // Toggle sort mode
    pub fn toggle_sort(&mut self) {
        self.sort_mode = self.sort_mode.next();
        self.resort_entries();
        self.rebuild_tree();
        self.message = Some(format!("Sort by: {}", self.sort_mode.label()));
    }

    fn resort_entries(&mut self) {
        Self::sort_entries_recursive(&mut self.root_entries, self.sort_mode);
    }

    fn sort_entries_recursive(entries: &mut [ZfsEntry], mode: SortMode) {
        match mode {
            SortMode::Size => entries.sort_by(|a, b| b.used.cmp(&a.used)),
            SortMode::Date => entries.sort_by(|a, b| b.creation.cmp(&a.creation)),
            SortMode::Name => entries.sort_by(|a, b| a.full_name.cmp(&b.full_name)),
        }

        for entry in entries.iter_mut() {
            Self::sort_entries_recursive(&mut entry.children, mode);
            // Also sort snapshots
            match mode {
                SortMode::Size => entry.snapshots.sort_by(|a, b| b.used.cmp(&a.used)),
                SortMode::Date => entry.snapshots.sort_by(|a, b| b.creation.cmp(&a.creation)),
                SortMode::Name => entry.snapshots.sort_by(|a, b| a.full_name.cmp(&b.full_name)),
            }
        }
    }

    // Toggle dry-run mode
    pub fn toggle_dry_run(&mut self) {
        self.dry_run_mode = !self.dry_run_mode;
        if self.dry_run_mode {
            self.message = Some("Dry-run mode ON".to_string());
        } else {
            self.message = Some("Dry-run mode OFF".to_string());
        }
    }

    // Show dry-run preview
    pub fn show_dry_run_preview(&mut self) {
        if self.marked.is_empty() {
            self.message = Some("No snapshots marked".to_string());
            return;
        }

        let total_size = self.total_marked_size();
        self.message = Some(format!(
            "Would delete {} snapshot(s), ~{} freed",
            self.marked.len(),
            zfs::format_size(total_size)
        ));
    }

    // Toggle auto-refresh
    pub fn toggle_auto_refresh(&mut self) {
        self.auto_refresh = !self.auto_refresh;
        if self.auto_refresh {
            self.message = Some(format!("Auto-refresh ON ({}s)", self.auto_refresh_interval));
        } else {
            self.message = Some("Auto-refresh OFF".to_string());
        }
    }

    // Check if auto-refresh is due
    pub fn check_auto_refresh(&mut self) -> Result<()> {
        if self.auto_refresh && self.last_refresh.elapsed() >= std::time::Duration::from_secs(self.auto_refresh_interval) {
            self.refresh()?;
        }
        Ok(())
    }

    // Get snapshot by name
    pub fn get_snapshot_date(&self, name: &str) -> Option<&str> {
        for ti in &self.tree_items {
            if let ListItem::Snapshot(s) = &ti.item {
                if s.full_name == name {
                    return Some(&s.creation);
                }
            }
        }
        None
    }
}