mod dir_helper;
mod include_static_helper;

use std::{collections::HashMap, fmt, ops::ControlFlow, sync::Arc};

use anyhow::{anyhow, ensure, Context};
use codespan_reporting::diagnostic::Diagnostic;
use dir_helper::DirHelper;
use handlebars::{handlebars_helper, Handlebars};
use include_static_helper::IncludeStaticHelper;
use log::{debug, error, info};
use serde::Serialize;

use crate::{
    config::Config,
    dirs::Dirs,
    fun::seasons::Season,
    html::{breadcrumbs::breadcrumbs_to_html, navmap::NavigationMap, tree::branches_to_html},
    import_map::ImportMap,
    parse::parse_tree_with_diagnostics,
    state::{report_diagnostics, Source},
    tree::SemaRoots,
    vfs::{
        self, Cd, ContentCache, Dir, DirEntry, DynDir, EditPath, ImageSize, MemDir, Overlay,
        ToDynDir, VPath, VPathBuf,
    },
};

use crate::state::{FileId, Treehouse};

#[derive(Debug, Clone)]
pub struct ParsedTree {
    root_key: String,
    file_id: FileId,
}

#[derive(Serialize)]
struct Page {
    title: String,
    thumbnail: Option<Thumbnail>,
    scripts: Vec<String>,
    styles: Vec<String>,
    breadcrumbs: String,
    tree_path: Option<String>,
    tree: String,
}

#[derive(Serialize)]
struct Thumbnail {
    url: String,
    alt: Option<String>,
}

#[derive(Serialize)]
struct BaseTemplateData<'a> {
    config: &'a Config,
    import_map: String,
    season: Option<Season>,
    dev: bool,
}

#[derive(Serialize)]
struct PageTemplateData<'a> {
    #[serde(flatten)]
    base: &'a BaseTemplateData<'a>,
    page: Page,
}

fn create_handlebars(site: &str, static_: DynDir) -> Handlebars<'static> {
    let mut handlebars = Handlebars::new();

    handlebars_helper!(cat: |a: String, b: String| a + &b);

    handlebars.register_helper("cat", Box::new(cat));
    handlebars.register_helper("asset", Box::new(DirHelper::new(site, static_.clone())));
    handlebars.register_helper(
        "include_static",
        Box::new(IncludeStaticHelper::new(static_)),
    );

    handlebars
}

fn load_templates(handlebars: &mut Handlebars, dir: &dyn Dir) {
    vfs::walk_dir_rec(dir, VPath::ROOT, &mut |path| {
        if path.extension() == Some("hbs") {
            if let Some(content) = dir.content(path).and_then(|b| String::from_utf8(b).ok()) {
                if let Err(err) = handlebars.register_template_string(path.as_str(), content) {
                    error!("in template: {err}");
                }
            }
        }
        ControlFlow::Continue(())
    });
}

fn parse_tree(
    treehouse: &mut Treehouse,
    config: &Config,
    source: String,
    source_path: VPathBuf,
    target_path: VPathBuf,
    tree_path: String,
) -> anyhow::Result<(Option<ParsedTree>, Vec<Diagnostic<FileId>>)> {
    let file_id = treehouse.add_file(
        source_path.as_str().to_owned(),
        Source::Tree {
            input: source,
            target_path: target_path.clone(),
            tree_path: tree_path.clone(),
        },
    );

    match parse_tree_with_diagnostics(treehouse, file_id) {
        Ok(roots) => {
            let mut diagnostics = vec![];
            let roots = SemaRoots::from_roots(treehouse, &mut diagnostics, config, file_id, roots);

            let root_key = tree_path.clone();
            treehouse.roots.insert(root_key.clone(), roots);

            Ok((Some(ParsedTree { root_key, file_id }), diagnostics))
        }
        Err(diagnostics) => Ok((None, diagnostics)),
    }
}

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

    vfs::walk_dir_rec(&*dirs.content, VPath::ROOT, &mut |path| {
        if path.extension() == Some("tree") {
            if let Some(source) = dirs
                .content
                .content(path)
                .and_then(|b| String::from_utf8(b).ok())
            {
                let tree_path = path.with_extension("");
                let target_path = path.with_extension("html");

                debug!("tree file: {path}");

                match parse_tree(
                    &mut treehouse,
                    config,
                    source,
                    path.to_owned(),
                    target_path,
                    tree_path.as_str().to_owned(),
                ) {
                    Ok((parsed_tree, mut parse_diagnostics)) => {
                        diagnostics.append(&mut parse_diagnostics);
                        if let Some(parsed_tree) = parsed_tree {
                            parsed_trees.insert(tree_path, parsed_tree);
                        }
                    }
                    Err(err) => {
                        error!("failed to parse tree {path}: {err:?}")
                    }
                }
            }
        }

        ControlFlow::Continue(())
    });

    report_diagnostics(&treehouse.files, &diagnostics)?;

    Ok((treehouse, parsed_trees))
}

fn generate_simple_template(
    sources: &Sources,
    handlebars: &Handlebars,
    template_name: &str,
) -> anyhow::Result<String> {
    let base_template_data = BaseTemplateData {
        config: &sources.config,
        import_map: serde_json::to_string_pretty(&sources.import_map)
            .expect("import map should be serializable to JSON"),
        season: Season::current(),
        dev: cfg!(debug_assertions),
    };
    handlebars
        .render(template_name, &base_template_data)
        .context("failed to render template")
}

fn generate_simple_template_or_error(
    sources: &Sources,
    handlebars: &Handlebars,
    template_name: &str,
) -> String {
    match generate_simple_template(sources, handlebars, template_name) {
        Ok(html) => html,
        Err(error) => format!("error: {error:?}"),
    }
}

fn generate_tree(
    sources: &Sources,
    dirs: &Dirs,
    handlebars: &Handlebars,
    parsed_tree: &ParsedTree,
) -> anyhow::Result<String> {
    let breadcrumbs = breadcrumbs_to_html(
        &sources.config,
        &sources.navigation_map,
        &parsed_tree.root_key,
    );

    let mut tree = String::new();
    let roots = sources
        .treehouse
        .roots
        .get(&parsed_tree.root_key)
        .expect("tree should have been added to the treehouse");
    branches_to_html(
        &mut tree,
        &sources.treehouse,
        &sources.config,
        dirs,
        parsed_tree.file_id,
        &roots.branches,
    );

    let base_template_data = BaseTemplateData {
        config: &sources.config,
        import_map: serde_json::to_string_pretty(&sources.import_map)
            .expect("import map should be serializable to JSON"),
        season: Season::current(),
        dev: cfg!(debug_assertions),
    };

    let template_data = PageTemplateData {
        base: &base_template_data,
        page: Page {
            title: roots.attributes.title.clone(),
            thumbnail: roots
                .attributes
                .thumbnail
                .as_ref()
                .map(|thumbnail| Thumbnail {
                    url: sources.config.pic_url(&*dirs.pic, &thumbnail.id),
                    alt: thumbnail.alt.clone(),
                }),
            scripts: roots.attributes.scripts.clone(),
            styles: roots.attributes.styles.clone(),
            breadcrumbs,
            tree_path: sources
                .treehouse
                .tree_path(parsed_tree.file_id)
                .map(|s| s.to_owned()),
            tree,
        },
    };
    let template_name = roots
        .attributes
        .template
        .clone()
        .unwrap_or_else(|| "_tree.hbs".into());

    ensure!(
        handlebars.has_template(&template_name),
        "template {template_name} does not exist"
    );

    handlebars
        .render(&template_name, &template_data)
        .context("template rendering failed")
}

fn generate_tree_or_error(
    sources: &Sources,
    dirs: &Dirs,
    handlebars: &Handlebars,
    parsed_tree: &ParsedTree,
) -> String {
    match generate_tree(sources, dirs, handlebars, parsed_tree) {
        Ok(html) => html,
        Err(error) => format!("error: {error:?}"),
    }
}

pub struct Sources {
    pub config: Config,
    pub treehouse: Treehouse,
    pub parsed_trees: HashMap<VPathBuf, ParsedTree>,
    pub navigation_map: NavigationMap,
    pub import_map: ImportMap,
}

impl Sources {
    pub fn load(dirs: &Dirs) -> anyhow::Result<Self> {
        info!("loading config");
        let mut config: Config = toml_edit::de::from_str(
            &dirs
                .root
                .content(VPath::new("treehouse.toml"))
                .map(String::from_utf8)
                .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)?;

        info!("parsing tree files");
        let (treehouse, parsed_trees) = parse_trees(&config, dirs)?;

        info!("constructing navigation map");
        let navigation_map = NavigationMap::build(&treehouse, "index");

        info!("constructing import map");
        let import_map = ImportMap::generate(
            &config.site,
            &Cd::new(dirs.static_.clone(), VPathBuf::new("js")),
            &config.build.javascript.import_roots,
        );

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

/// Acceleration structure for `dir` operations on [`TreehouseDir`]s.
#[derive(Debug, Default)]
struct DirIndex {
    full_path: VPathBuf,
    children: HashMap<VPathBuf, DirIndex>,
}

impl DirIndex {
    pub fn new<'a>(paths: impl Iterator<Item = &'a VPath>) -> Self {
        let mut root = DirIndex::default();

        for path in paths {
            let mut parent = &mut root;
            let mut full_path = VPath::ROOT.to_owned();
            for segment in path.segments() {
                full_path.push(segment);
                let child = parent
                    .children
                    .entry(segment.to_owned())
                    .or_insert_with(|| DirIndex {
                        full_path: full_path.clone(),
                        children: HashMap::new(),
                    });
                parent = child;
            }
        }

        root
    }
}

struct TreehouseDir {
    dirs: Arc<Dirs>,
    sources: Arc<Sources>,
    dir_index: DirIndex,
    handlebars: Handlebars<'static>,
}

impl TreehouseDir {
    fn new(dirs: Arc<Dirs>, sources: Arc<Sources>, dir_index: DirIndex) -> Self {
        let mut handlebars = create_handlebars(&sources.config.site, dirs.static_.clone());
        load_templates(&mut handlebars, &dirs.template);

        Self {
            dirs,
            sources,
            dir_index,
            handlebars,
        }
    }
}

impl Dir for TreehouseDir {
    fn dir(&self, path: &VPath) -> Vec<DirEntry> {
        // NOTE: This does not include simple templates, because that's not really needed right now.

        let mut index = &self.dir_index;
        for component in path.segments() {
            if let Some(child) = index.children.get(component) {
                index = child;
            } else {
                // There cannot possibly be any entries under an invalid path.
                // Bail early.
                return vec![];
            }
        }

        index
            .children
            .values()
            .map(|child| DirEntry {
                path: child.full_path.clone(),
            })
            .collect()
    }

    fn content(&self, path: &VPath) -> Option<Vec<u8>> {
        debug!("content({path})");

        let path = if path.is_root() {
            VPath::new_const("index")
        } else {
            path
        };

        self.sources
            .parsed_trees
            .get(path)
            .map(|parsed_tree| {
                generate_tree_or_error(&self.sources, &self.dirs, &self.handlebars, parsed_tree)
                    .into()
            })
            .or_else(|| {
                if path.file_name().is_some_and(|s| !s.starts_with('_')) {
                    let template_name = path.with_extension("hbs");
                    if self.handlebars.has_template(template_name.as_str()) {
                        return Some(
                            generate_simple_template_or_error(
                                &self.sources,
                                &self.handlebars,
                                template_name.as_str(),
                            )
                            .into(),
                        );
                    }
                }
                None
            })
    }

    fn content_version(&self, _path: &VPath) -> Option<String> {
        None
    }
}

impl fmt::Debug for TreehouseDir {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("TreehouseDir")
    }
}

struct HtmlCanonicalize<T> {
    inner: T,
}

impl<T> HtmlCanonicalize<T> {
    pub fn new(inner: T) -> Self {
        Self { inner }
    }
}

impl<T> Dir for HtmlCanonicalize<T>
where
    T: Dir,
{
    fn dir(&self, path: &VPath) -> Vec<DirEntry> {
        self.inner.dir(path)
    }

    fn content(&self, path: &VPath) -> Option<Vec<u8>> {
        let mut path = path.to_owned();
        if path.extension() == Some("html") {
            path.set_extension("");
        }

        self.inner.content(&path)
    }

    fn content_version(&self, path: &VPath) -> Option<String> {
        self.inner.content_version(path)
    }

    fn image_size(&self, path: &VPath) -> Option<ImageSize> {
        self.inner.image_size(path)
    }

    fn anchor(&self, path: &VPath) -> Option<VPathBuf> {
        self.inner.anchor(path)
    }

    fn edit_path(&self, path: &VPath) -> Option<EditPath> {
        self.inner.edit_path(path)
    }
}

impl<T> fmt::Debug for HtmlCanonicalize<T>
where
    T: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "HtmlCanonicalize({:?})", self.inner)
    }
}

pub fn target(dirs: Arc<Dirs>, sources: Arc<Sources>) -> DynDir {
    let mut root = MemDir::new();
    root.add(VPath::new("static"), dirs.static_.clone());

    let dir_index = DirIndex::new(sources.parsed_trees.keys().map(|x| &**x));
    let tree_view = TreehouseDir::new(dirs, sources, dir_index);

    let tree_view = ContentCache::new(tree_view);
    tree_view.warm_up();
    let tree_view = HtmlCanonicalize::new(tree_view);

    Overlay::new(tree_view.to_dyn(), root.to_dyn()).to_dyn()
}