use std::{collections::HashMap, ops::ControlFlow};

use anyhow::{Context, anyhow};
use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
use tracing::{error, info_span, instrument};

use crate::{
    config::Config,
    dirs::Dirs,
    doc::Doc,
    html::navmap::NavigationMap,
    import_map::ImportMap,
    parse::parse_tree_with_diagnostics,
    state::{Source, Treehouse, report_diagnostics},
    tree::SemaRoots,
    vfs::{self, Cd, Content, VPath, VPathBuf},
};

pub struct Sources {
    pub config: Config,
    pub treehouse: Treehouse,
    pub navigation_map: NavigationMap,
    pub import_map: ImportMap,
}

impl Sources {
    pub fn load(dirs: &Dirs) -> anyhow::Result<Self> {
        let config = {
            let _span = info_span!("load_config").entered();
            let mut config: Config = toml_edit::de::from_str(
                &vfs::query::<Content>(&dirs.root, VPath::new_const("treehouse.toml"))
                    .map(Content::string)
                    .ok_or_else(|| anyhow!("config file does not exist"))??,
            )
            .context("failed to deserialize config")?;
            config.site = std::env::var("TREEHOUSE_SITE").unwrap_or(config.site);
            config.autopopulate_emoji(&*dirs.emoji)?;
            config.autopopulate_pics(&*dirs.pic)?;
            config.load_syntaxes(dirs.syntax.clone())?;
            config
        };

        let treehouse = load_trees(&config, dirs)?;
        let navigation_map = NavigationMap::build(
            &treehouse,
            treehouse.files_by_tree_path[VPath::new("index")],
        );
        let import_map = ImportMap::generate(
            &config.site,
            &Cd::new(dirs.static_.clone(), VPathBuf::new("js")),
            &config.build.javascript.import_roots,
        );

        Ok(Sources {
            config,
            treehouse,
            navigation_map,
            import_map,
        })
    }
}

#[instrument(skip(config, dirs))]
fn load_trees(config: &Config, dirs: &Dirs) -> anyhow::Result<Treehouse> {
    let mut treehouse = Treehouse::new();
    let mut diagnostics = vec![];
    let mut parsed_trees = HashMap::new();

    let mut paths = vec![];
    let mut doc_paths = vec![];

    vfs::walk_dir_rec(&*dirs.content, VPath::ROOT, &mut |path| {
        match path.extension() {
            Some("tree") => paths.push(path.to_owned()),
            Some("dj") => doc_paths.push(path.to_owned()),
            _ => (),
        }
        ControlFlow::Continue(())
    });

    // Trees

    // NOTE: Sources are filled in later; they can be left out until a call to report_diagnostics.
    let file_ids: Vec<_> = paths
        .iter()
        .map(|path| treehouse.add_file(path.clone(), Source::Other(String::new())))
        .collect();

    let parse_results: Vec<_> = {
        let _span = info_span!("load_trees::parse").entered();
        paths
            .into_par_iter()
            .zip(&file_ids)
            .flat_map(|(path, &file_id)| {
                vfs::query::<Content>(&dirs.content, &path)
                    .and_then(|c| c.string().ok())
                    .map(|input| {
                        let parse_result = parse_tree_with_diagnostics(file_id, &input);
                        (path, file_id, input, parse_result)
                    })
            })
            .collect()
    };

    for (path, file_id, input, _) in &parse_results {
        let tree_path = path.with_extension("");
        treehouse
            .files_by_tree_path
            .insert(tree_path.clone(), *file_id);
        treehouse.set_source(
            *file_id,
            Source::Tree {
                input: input.clone(),
                tree_path,
            },
        );
    }

    {
        let _span = info_span!("load_trees::sema").entered();
        for (path, file_id, _, result) in parse_results {
            match result {
                Ok(roots) => {
                    let roots = SemaRoots::from_roots(
                        &mut treehouse,
                        &mut diagnostics,
                        config,
                        file_id,
                        roots,
                    );
                    treehouse.roots.insert(file_id, roots);
                    parsed_trees.insert(path, file_id);
                }
                Err(mut parse_diagnostics) => diagnostics.append(&mut parse_diagnostics),
            }
        }
    }

    // Docs

    let mut doc_file_ids = vec![];

    for path in &doc_paths {
        if let Some(input) =
            vfs::query::<Content>(&dirs.content, path).and_then(|c| c.string().ok())
        {
            let file_id = treehouse.add_file(path.clone(), Source::Other(input));
            treehouse.files_by_doc_path.insert(path.clone(), file_id);
            doc_file_ids.push(file_id);
        } else {
            error!(
                "doc {path} does not exist in content directory even though it was enumerated via walk_dir_rec"
            );
        }
    }

    for file_id in doc_file_ids {
        let (doc, mut doc_diagnostics) = Doc::parse(&mut treehouse, config, file_id);
        treehouse.docs.insert(file_id, doc);
        diagnostics.append(&mut doc_diagnostics);
    }

    // Tags

    for file_id in treehouse.files_by_tree_path.values() {
        let roots = &treehouse.roots[file_id];
        for tag_name in &roots.attributes.tags {
            let tag = treehouse.tags.entry(tag_name.clone()).or_default();
            tag.files.push(*file_id);
        }
    }

    for file_id in treehouse.files_by_doc_path.values() {
        let doc = &treehouse.docs[file_id];
        for tag_name in &doc.attributes.tags {
            let tag = treehouse.tags.entry(tag_name.clone()).or_default();
            tag.files.push(*file_id);
        }
    }

    // Diagnostics

    report_diagnostics(&treehouse, &diagnostics)?;

    Ok(treehouse)
}