use crate::{
    ByteFormat, CountMode, InodeFilter, Throttle, WalkOptions, WalkResult, crossdev, token_count,
};
use anyhow::Result;
use filesize::PathExt;
use jwalk::rayon::prelude::*;
use owo_colors::{AnsiColors as Color, OwoColorize};
use std::fmt;
use std::time::Duration;
use std::{io, path::Path};

/// Aggregate the given `paths` and write information about them to `out` in a human-readable format.
/// If `compute_total` is set, it will write an additional line with the total size across all given `paths`.
/// If `sort_by_size_in_bytes` is set, we will sort all sizes (ascending) before outputting them.
pub fn aggregate(
    mut out: impl io::Write,
    mut err: Option<impl io::Write>,
    walk_options: WalkOptions,
    compute_total: bool,
    sort_by_size_in_bytes: bool,
    byte_format: ByteFormat,
    paths: impl IntoIterator<Item = impl AsRef<Path>>,
) -> Result<(WalkResult, Statistics)> {
    let paths: Vec<_> = paths
        .into_iter()
        .map(|path| path.as_ref().to_owned())
        .collect();

    let mut res = WalkResult::default();
    let mut stats = Statistics::for_mode(walk_options.count_mode);
    let mut total = 0;
    let num_roots = paths.len();
    let mut aggregates = Vec::new();
    let mut inodes = InodeFilter::default();
    let progress = Throttle::new(Duration::from_millis(100), Duration::from_secs(1).into());

    let mut precomputed_file_root_sizes: Vec<Option<(u128, u64)>> = vec![None; paths.len()];
    if walk_options.count_mode == CountMode::Tokens {
        if let Some(token_mode) = walk_options.token_mode.as_ref() {
            if token_mode.config.dry_run {
                for (path_idx, path) in paths.iter().enumerate() {
                    if !path.is_file() {
                        continue;
                    }
                    match path.symlink_metadata() {
                        Ok(metadata) => {
                            let _ = walk_options.count_hard_links || inodes.add(&metadata);
                            precomputed_file_root_sizes[path_idx] = Some((0, 0));
                        }
                        Err(_) => {
                            precomputed_file_root_sizes[path_idx] = Some((0, 1));
                        }
                    }
                }
            } else {
                let mut token_jobs = Vec::new();
                for (path_idx, path) in paths.iter().enumerate() {
                    if !path.is_file() {
                        continue;
                    }

                    match path.symlink_metadata() {
                        Ok(metadata) => {
                            if walk_options.count_hard_links || inodes.add(&metadata) {
                                token_jobs.push((path_idx, path.clone()));
                            } else {
                                precomputed_file_root_sizes[path_idx] = Some((0, 0));
                            }
                        }
                        Err(_) => {
                            precomputed_file_root_sizes[path_idx] = Some((0, 1));
                        }
                    }
                }

                let tokenizer = token_mode.tokenizer.clone();
                let token_results: Vec<(usize, Option<u128>)> = token_jobs
                    .par_iter()
                    .map(|(path_idx, path)| {
                        (
                            *path_idx,
                            token_count::count_tokens_in_text_file(path, tokenizer.as_ref())
                                .map(u128::from),
                        )
                    })
                    .collect();

                for (path_idx, maybe_file_size) in token_results {
                    precomputed_file_root_sizes[path_idx] = Some((maybe_file_size.unwrap_or(0), 0));
                }
            }
        } else {
            for (path_idx, path) in paths.iter().enumerate() {
                if path.is_file() {
                    precomputed_file_root_sizes[path_idx] = Some((0, 1));
                }
            }
        }
    }

    for (path_idx, path) in paths.iter().enumerate() {
        let mut num_bytes = 0u128;
        let mut num_errors = 0u64;

        // Fast path for explicit file roots avoids jwalk setup cost per file.
        if path.is_file() {
            stats.entries_traversed += 1;

            if let Some((file_size, file_errors)) = precomputed_file_root_sizes[path_idx] {
                num_errors += file_errors;
                stats.record_file_size(file_size, walk_options.count_mode);
                num_bytes += file_size;

                if sort_by_size_in_bytes {
                    aggregates.push((path.to_owned(), num_bytes, num_errors));
                } else {
                    output_colored_path(
                        &mut out,
                        path,
                        num_bytes,
                        num_errors,
                        path_color_of(path),
                        byte_format,
                    )?;
                }
                total += num_bytes;
                res.num_errors += num_errors;
                continue;
            }

            match path.as_path().symlink_metadata() {
                Ok(metadata) => {
                    let file_size = if walk_options.count_hard_links || inodes.add(&metadata) {
                        match walk_options.count_mode {
                            CountMode::Bytes => {
                                if walk_options.apparent_size {
                                    Some(u128::from(metadata.len()))
                                } else {
                                    Some(u128::from(
                                        path.as_path().size_on_disk_fast(&metadata).unwrap_or_else(
                                            |_| {
                                                num_errors += 1;
                                                0
                                            },
                                        ),
                                    ))
                                }
                            }
                            CountMode::Tokens => {
                                if let Some(token_mode) = walk_options.token_mode.as_ref() {
                                    if token_mode.config.dry_run {
                                        Some(0)
                                    } else {
                                        token_count::count_tokens_in_text_file(
                                            path.as_path(),
                                            token_mode.tokenizer.as_ref(),
                                        )
                                        .map(u128::from)
                                    }
                                } else {
                                    num_errors += 1;
                                    None
                                }
                            }
                        }
                    } else {
                        Some(0)
                    };

                    if let Some(file_size) = file_size {
                        stats.record_file_size(file_size, walk_options.count_mode);
                        num_bytes += file_size;
                    }
                }
                Err(_) => {
                    num_errors += 1;
                }
            }

            if sort_by_size_in_bytes {
                aggregates.push((path.to_owned(), num_bytes, num_errors));
            } else {
                output_colored_path(
                    &mut out,
                    path,
                    num_bytes,
                    num_errors,
                    path_color_of(path),
                    byte_format,
                )?;
            }
            total += num_bytes;
            res.num_errors += num_errors;
            continue;
        }

        let device_id = match crossdev::init(path.as_ref()) {
            Ok(id) => id,
            Err(_) => {
                num_errors += 1;
                res.num_errors += 1;
                aggregates.push((path.to_owned(), num_bytes, num_errors));
                continue;
            }
        };
        for entry in walk_options.iter_from_path(path.as_ref(), device_id, false) {
            stats.entries_traversed += 1;
            progress.throttled(|| {
                if let Some(err) = err.as_mut() {
                    write!(err, "Enumerating {} items\r", stats.entries_traversed).ok();
                }
            });
            match entry {
                Ok(entry) => {
                    let file_size: Option<u128> = match entry.client_state {
                        Some(Ok(ref m))
                            if (walk_options.count_hard_links || inodes.add(m))
                                && (walk_options.cross_filesystems
                                    || crossdev::is_same_device(device_id, m)) =>
                        {
                            match walk_options.count_mode {
                                CountMode::Bytes => {
                                    if walk_options.apparent_size {
                                        Some(u128::from(m.len()))
                                    } else {
                                        Some(u128::from(
                                            entry.path().size_on_disk_fast(m).unwrap_or_else(
                                                |_| {
                                                    num_errors += 1;
                                                    0
                                                },
                                            ),
                                        ))
                                    }
                                }
                                CountMode::Tokens => {
                                    if !m.is_file() {
                                        Some(0)
                                    } else {
                                        let Some(token_mode) = walk_options.token_mode.as_ref()
                                        else {
                                            num_errors += 1;
                                            continue;
                                        };
                                        if token_mode.config.dry_run {
                                            Some(0)
                                        } else {
                                            token_count::count_tokens_in_text_file(
                                                &entry.path(),
                                                token_mode.tokenizer.as_ref(),
                                            )
                                            .map(u128::from)
                                        }
                                    }
                                }
                            }
                        }
                        Some(Ok(_)) => Some(0),
                        Some(Err(_)) => {
                            num_errors += 1;
                            Some(0)
                        }
                        None => Some(0), // ignored directory
                    };

                    let Some(file_size) = file_size else {
                        continue;
                    };

                    stats.record_file_size(file_size, walk_options.count_mode);
                    num_bytes += file_size;
                }
                Err(_) => num_errors += 1,
            }
        }

        if let Some(err) = err.as_mut() {
            write!(err, "\x1b[2K\r").ok();
        }

        if sort_by_size_in_bytes {
            aggregates.push((path.to_owned(), num_bytes, num_errors));
        } else {
            output_colored_path(
                &mut out,
                path,
                num_bytes,
                num_errors,
                path_color_of(path),
                byte_format,
            )?;
        }
        total += num_bytes;
        res.num_errors += num_errors;
    }

    stats.finish(walk_options.count_mode);

    if sort_by_size_in_bytes {
        aggregates.sort_by_key(|&(_, num_bytes, _)| num_bytes);
        for (path, num_bytes, num_errors) in aggregates.into_iter() {
            output_colored_path(
                &mut out,
                &path,
                num_bytes,
                num_errors,
                path_color_of(&path),
                byte_format,
            )?;
        }
    }

    if num_roots > 1 && compute_total {
        output_colored_path(
            &mut out,
            Path::new("total"),
            total,
            res.num_errors,
            None,
            byte_format,
        )?;
    }
    Ok((res, stats))
}

fn path_color_of(path: impl AsRef<Path>) -> Option<Color> {
    (!path.as_ref().is_file()).then_some(Color::Cyan)
}

fn output_colored_path(
    out: &mut impl io::Write,
    path: impl AsRef<Path>,
    value: u128,
    num_errors: u64,
    path_color: Option<Color>,
    byte_format: ByteFormat,
) -> std::result::Result<(), io::Error> {
    let size = byte_format.display(value).to_string();
    let size = size.green();
    let size_width = byte_format.width();
    let path = path.as_ref().display();

    let errors = if num_errors != 0 {
        format!(
            "  <{num_errors} IO Error{plural_s}>",
            plural_s = if num_errors > 1 { "s" } else { "" }
        )
    } else {
        "".into()
    };

    if let Some(color) = path_color {
        writeln!(out, "{size:>size_width$} {}{errors}", path.color(color))
    } else {
        writeln!(out, "{size:>size_width$} {path}{errors}")
    }
}

/// Statistics obtained during a filesystem walk
pub struct Statistics {
    /// The amount of entries we have seen during filesystem traversal
    pub entries_traversed: u64,
    /// The size of the smallest file encountered in bytes
    pub smallest_file_in_bytes: u128,
    /// The size of the largest file encountered in bytes
    pub largest_file_in_bytes: u128,
    /// The size of the smallest file encountered in tokens.
    pub smallest_file_in_tokens: Option<u128>,
    /// The size of the largest file encountered in tokens.
    pub largest_file_in_tokens: Option<u128>,
}

impl Default for Statistics {
    fn default() -> Self {
        Self {
            entries_traversed: 0,
            smallest_file_in_bytes: 0,
            largest_file_in_bytes: 0,
            smallest_file_in_tokens: None,
            largest_file_in_tokens: None,
        }
    }
}

impl Statistics {
    fn for_mode(mode: CountMode) -> Self {
        let mut stats = Self::default();
        match mode {
            CountMode::Bytes => {
                stats.smallest_file_in_bytes = u128::MAX;
            }
            CountMode::Tokens => {
                stats.smallest_file_in_tokens = Some(u128::MAX);
                stats.largest_file_in_tokens = Some(0);
            }
        }
        stats
    }

    fn record_file_size(&mut self, value: u128, mode: CountMode) {
        match mode {
            CountMode::Bytes => {
                self.smallest_file_in_bytes = self.smallest_file_in_bytes.min(value);
                self.largest_file_in_bytes = self.largest_file_in_bytes.max(value);
            }
            CountMode::Tokens => {
                if let Some(smallest) = self.smallest_file_in_tokens.as_mut() {
                    *smallest = (*smallest).min(value);
                }
                if let Some(largest) = self.largest_file_in_tokens.as_mut() {
                    *largest = (*largest).max(value);
                }
            }
        }
    }

    fn finish(&mut self, mode: CountMode) {
        if self.entries_traversed != 0 {
            return;
        }
        match mode {
            CountMode::Bytes => {
                self.smallest_file_in_bytes = 0;
            }
            CountMode::Tokens => {
                self.smallest_file_in_tokens = Some(0);
                self.largest_file_in_tokens = Some(0);
            }
        }
    }
}

impl fmt::Debug for Statistics {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if let (Some(smallest), Some(largest)) =
            (self.smallest_file_in_tokens, self.largest_file_in_tokens)
        {
            return f
                .debug_struct("Statistics")
                .field("entries_traversed", &self.entries_traversed)
                .field("smallest_file_in_tokens", &smallest)
                .field("largest_file_in_tokens", &largest)
                .finish();
        }

        f.debug_struct("Statistics")
            .field("entries_traversed", &self.entries_traversed)
            .field("smallest_file_in_bytes", &self.smallest_file_in_bytes)
            .field("largest_file_in_bytes", &self.largest_file_in_bytes)
            .finish()
    }
}
