lla 0.5.5

use super::FileFormatter;
use crate::error::Result;
use crate::plugin::PluginManager;
use crate::utils::color::*;
use crate::utils::icons::format_with_icon;
use colored::Colorize;
use lla_plugin_interface::proto::DecoratedEntry;
use std::collections::{HashMap, HashSet};
use std::io::{self, Write};
use std::path::{Path, PathBuf};

const BUFFER_SIZE: usize = 16384;

pub struct TreeFormatter {
    pub show_icons: bool,
}

impl TreeFormatter {
    pub fn new(show_icons: bool) -> Self {
        Self { show_icons }
    }

    fn format_entry(&self, path: &Path) -> String {
        let colored_name = colorize_file_name(path).to_string();
        if self.show_icons {
            format_with_icon(path, colored_name, true)
        } else {
            colored_name
        }
    }

    fn build_tree(
        &self,
        entries: &[DecoratedEntry],
    ) -> (Vec<PathBuf>, HashMap<PathBuf, Vec<PathBuf>>) {
        let mut tree: HashMap<PathBuf, Vec<PathBuf>> = HashMap::with_capacity(entries.len());
        let mut path_set: HashSet<PathBuf> = HashSet::with_capacity(entries.len());
        let mut child_paths = HashSet::new();

        for entry in entries {
            path_set.insert(PathBuf::from(&entry.path));
        }

        for path in path_set.iter() {
            if let Some(parent) = path.parent() {
                if path_set.contains(parent) {
                    tree.entry(parent.to_path_buf())
                        .or_insert_with(Vec::new)
                        .push(path.clone());
                    child_paths.insert(path.clone());
                }
            }
        }
        for children in tree.values_mut() {
            children.sort_unstable();
        }
        let mut root_paths: Vec<_> = path_set
            .into_iter()
            .filter(|path| !child_paths.contains(path))
            .collect();
        root_paths.sort_unstable();

        (root_paths, tree)
    }

    fn write_tree_recursive(
        &self,
        path: &Path,
        prefix: &str,
        is_last: bool,
        tree: &HashMap<PathBuf, Vec<PathBuf>>,
        writer: &mut impl Write,
        current_depth: usize,
        max_depth: Option<usize>,
    ) -> io::Result<()> {
        if let Some(max) = max_depth {
            if current_depth > max {
                return Ok(());
            }
        }

        let node_prefix = if is_last { "└── " } else { "├── " };
        let child_prefix = if is_last { "    " } else { "│   " };

        let formatted_name = self.format_entry(path);
        write!(
            writer,
            "{}{}{}\n",
            prefix.bright_black(),
            node_prefix.bright_black(),
            formatted_name
        )?;

        if let Some(children) = tree.get(path) {
            let new_prefix = format!("{}{}", prefix, child_prefix);
            let last_idx = children.len().saturating_sub(1);
            for (i, child) in children.iter().enumerate() {
                let is_last_child = i == last_idx;
                self.write_tree_recursive(
                    child,
                    &new_prefix,
                    is_last_child,
                    tree,
                    writer,
                    current_depth + 1,
                    max_depth,
                )?;
            }
        }
        Ok(())
    }
}

impl FileFormatter for TreeFormatter {
    fn format_files(
        &self,
        files: &[DecoratedEntry],
        _plugin_manager: &mut PluginManager,
        depth: Option<usize>,
    ) -> Result<String> {
        if files.is_empty() {
            return Ok(String::new());
        }

        if depth == Some(0) {
            return Ok(String::new());
        }

        let (root_paths, tree) = self.build_tree(files);
        let mut buffer = Vec::with_capacity(BUFFER_SIZE);

        let last_idx = root_paths.len().saturating_sub(1);
        for (i, path) in root_paths.iter().enumerate() {
            let is_last = i == last_idx;
            self.write_tree_recursive(path, "", is_last, &tree, &mut buffer, 0, depth)?;
        }

        Ok(String::from_utf8_lossy(&buffer).into_owned())
    }
}