use std::ops::ControlFlow;

use crate::{
    parse::parse_tree_with_diagnostics,
    state::{report_diagnostics, Source, Treehouse},
    tree::ast::{Branch, Roots},
    vfs::{self, Content, Dir, VPath},
};

use super::WcArgs;

fn wc_branch(source: &str, branch: &Branch) -> usize {
    let word_count = source[branch.content.clone()].split_whitespace().count();
    word_count
        + branch
            .children
            .iter()
            .map(|branch| wc_branch(source, branch))
            .sum::<usize>()
}

fn wc_roots(source: &str, roots: &Roots) -> usize {
    roots
        .branches
        .iter()
        .map(|branch| wc_branch(source, branch))
        .sum()
}

pub fn wc_cli(content_dir: &dyn Dir, mut wc_args: WcArgs) -> anyhow::Result<()> {
    if wc_args.paths.is_empty() {
        vfs::walk_dir_rec(content_dir, VPath::ROOT, &mut |path| {
            if path.extension() == Some("tree") {
                wc_args.paths.push(path.to_owned());
            }
            ControlFlow::Continue(())
        });
    }

    let mut treehouse = Treehouse::new();

    let mut total = 0;

    for path in &wc_args.paths {
        if let Some(content) =
            vfs::query::<Content>(content_dir, path).and_then(|b| b.string().ok())
        {
            let file_id = treehouse.add_file(path.clone(), Source::Other(content.clone()));
            match parse_tree_with_diagnostics(file_id, &content) {
                Ok(parsed) => {
                    let source = treehouse.source(file_id);
                    let word_count = wc_roots(source.input(), &parsed);
                    println!("{word_count:>8} {}", treehouse.path(file_id));
                    total += word_count;
                }
                Err(diagnostics) => {
                    report_diagnostics(&treehouse, &diagnostics)?;
                }
            }
        }
    }

    println!("{total:>8} total");

    Ok(())
}
