use std::{
    collections::HashMap,
    io::{IsTerminal, Write},
    path::{Path, PathBuf},
};

use clap::Parser;
use indicatif::{ProgressBar, ProgressStyle};
use serde::Serialize;
use walkdir::WalkDir;

fn main() {
    if let Err(err) = fallible_main() {
        eprintln!("ERROR: {err}");
        let mut underlying = err.source();
        while let Some(err) = underlying {
            eprintln!("caused by: {err}");
            underlying = err.source();
        }
        std::process::exit(1);
    }
}

fn fallible_main() -> Result<(), Box<dyn std::error::Error>> {
    let args = Args::parse();
    let progress = args.progress_bar();

    let mut filecounts = FileCounts::default();
    for dir in args.dirs() {
        filecounts.filecount(dir, &progress, args.ignore_errors)?;
    }

    progress.done();

    if args.json() {
        filecounts.to_json(std::io::stdout())?;
    } else {
        filecounts.to_text(std::io::stdout())?;
    }
    Ok(())
}

trait VisualProgress {
    fn inc(&self, _filename: &Path) {}
    fn done(&self) {}
}

struct NoProgress {}

impl VisualProgress for NoProgress {}

struct Progress {
    bar: ProgressBar,
}

impl Progress {
    fn new() -> Self {
        let bar = ProgressBar::no_length();
        bar.set_style(ProgressStyle::with_template("{pos} {msg}").unwrap());
        Self { bar }
    }
}

impl VisualProgress for Progress {
    fn inc(&self, filename: &Path) {
        self.bar.set_message(filename.to_string_lossy().to_string());
        self.bar.inc(1);
    }

    fn done(&self) {
        self.bar.finish_and_clear();
    }
}

#[derive(Default, Serialize)]
struct FileCounts {
    counts: HashMap<PathBuf, usize>,
}

impl FileCounts {
    #[allow(clippy::borrowed_box)]
    fn filecount<P: AsRef<Path>>(
        &mut self,
        dir: P,
        progress: &Box<dyn VisualProgress>,
        ignore_errors: bool,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let dir = dir.as_ref();
        let mut count = 0;
        for entry in WalkDir::new(dir) {
            match entry {
                Ok(entry) => {
                    count += 1;
                    progress.inc(entry.path());
                }
                Err(err) if ignore_errors => {
                    count += 1;
                    progress.inc(Path::new("error reading directory"));
                    eprintln!("ERROR: {err}");
                }
                Err(err) => {
                    Err(err)?;
                }
            }
        }
        self.counts.insert(dir.to_path_buf(), count);
        Ok(())
    }

    fn to_text(&self, mut writer: impl Write) -> Result<(), Box<dyn std::error::Error>> {
        let mut dirs: Vec<&PathBuf> = self.counts.keys().collect();
        dirs.sort();
        for dir in dirs {
            let count = self.counts.get(dir).unwrap();
            writer.write_all(format!("{:8} {}\n", *count, dir.display()).as_bytes())?;
        }
        Ok(())
    }

    fn to_json(&self, mut writer: impl Write) -> Result<(), Box<dyn std::error::Error>> {
        serde_json::to_writer(&mut writer, &self.counts)?;
        writer.write_all(b"\n")?;
        Ok(())
    }
}

/// Count files in directories.
#[derive(Debug, Parser)]
#[clap(version)]
struct Args {
    /// Count files in this directory. Default is current working directory.
    #[clap(default_value = ".")]
    dirs: Vec<PathBuf>,

    /// Output JSON in one line. Use something like `jq` to format it.
    #[clap(long)]
    json: bool,

    /// Ignore errors reading directories.
    #[clap(long, short)]
    ignore_errors: bool,

    /// Hide progress.
    #[clap(long)]
    quiet: bool,
}

impl Args {
    fn dirs(&self) -> &[PathBuf] {
        &self.dirs
    }

    fn json(&self) -> bool {
        self.json
    }

    fn progress_bar(&self) -> Box<dyn VisualProgress> {
        if self.quiet || !std::io::stderr().is_terminal() {
            Box::new(NoProgress {})
        } else {
            Box::new(Progress::new())
        }
    }
}
