forklaunch 1.15.0

use std::collections::HashMap;

use anyhow::Result;

use super::rendered_template::RenderedTemplate;

#[derive(Clone)]
struct CommentStyle {
    open: &'static str,
    line: &'static str,
    close: &'static str,
}

fn get_comment_styles() -> HashMap<&'static str, CommentStyle> {
    let mut styles = HashMap::new();

    let c_style = CommentStyle {
        open: "/**",
        line: " * ",
        close: " */",
    };
    for ext in [
        // JavaScript/TypeScript
        "ts", "tsx", "js", "jsx", "mjs", // Java
        "java", "scala", "kt", "kts", // C family
        "c", "h", "cpp", "hpp", "cc", "cxx", "cs", // CSS family
        "css", "scss", "less",  // Rust
        "rs",    // Swift
        "swift", // PHP
        "php",
    ] {
        styles.insert(ext, c_style.clone());
    }

    let py_style = CommentStyle {
        open: "\"\"\"",
        line: "",
        close: "\"\"\"",
    };
    for ext in ["py", "pyi", "pyx"] {
        styles.insert(ext, py_style.clone());
    }

    let bash_style = CommentStyle {
        open: "#",
        line: "# ",
        close: "#",
    };
    for ext in [
        // Shell scripts
        "sh",
        "bash",
        "zsh",
        "fish",
        // Ruby
        "rb",
        "rake",
        "gemspec",
        // Perl
        "pl",
        "pm",
        // Python build
        "pyproject.toml",
        // Config files
        "yaml",
        "yml",
        "toml",
        // R
        "r",
        "R",
        // Makefile
        "mk",
        "makefile",
        "Makefile",
        // HashiCorp
        "tf",
        "hcl",
    ] {
        styles.insert(ext, bash_style.clone());
    }

    styles.insert(
        "rb",
        CommentStyle {
            open: "=begin",
            line: "",
            close: "=end",
        },
    );

    let html_style = CommentStyle {
        open: "<!--",
        line: " ",
        close: "-->",
    };
    for ext in [
        "html", "htm", "xml", "svg", "xaml", "vue", "jsx", "tsx", "asp", "aspx", "cshtml", "vbhtml",
    ] {
        styles.insert(ext, html_style.clone());
    }

    styles.insert(
        "lua",
        CommentStyle {
            open: "--[[",
            line: "",
            close: "--]]",
        },
    );

    let lisp_style = CommentStyle {
        open: ";;",
        line: ";; ",
        close: ";;",
    };
    for ext in ["lisp", "cl", "el", "clj", "cljs", "cljc", "edn"] {
        styles.insert(ext, lisp_style.clone());
    }

    let haskell_style = CommentStyle {
        open: "{-",
        line: " ",
        close: "-}",
    };
    for ext in ["hs", "lhs"] {
        styles.insert(ext, haskell_style.clone());
    }

    let sql_style = CommentStyle {
        open: "/*",
        line: " ",
        close: "*/",
    };
    for ext in ["sql", "mysql", "pgsql", "sqlite"] {
        styles.insert(ext, sql_style.clone());
    }

    let ini_style = CommentStyle {
        open: ";",
        line: "; ",
        close: ";",
    };
    for ext in ["ini", "conf", "cfg"] {
        styles.insert(ext, ini_style.clone());
    }

    styles.insert(
        "bat",
        CommentStyle {
            open: "rem",
            line: "rem ",
            close: "rem",
        },
    );

    styles.insert(
        "ps1",
        CommentStyle {
            open: "<#",
            line: "",
            close: "#>",
        },
    );

    styles
}

fn get_comment_style(filename: &str) -> Option<CommentStyle> {
    let styles = get_comment_styles();
    let extension = filename.rsplit('.').next().unwrap_or("").to_lowercase();

    if extension.is_empty() || extension == "json" || filename.to_lowercase() == "readme.md" {
        return None;
    }

    styles.get(extension.as_str()).cloned()
}

pub(crate) fn generate_watermark(filename: &str) -> Option<String> {
    let Some(style) = get_comment_style(filename) else {
        return None;
    };

    let file_line = format!("File: {}", filename);
    let content = vec![
        "Generated by ForkLaunch",
        file_line.as_str(),
        "This is an auto-generated file. Modifications are encouraged but may inhibit automated upgrades.",
    ];

    if style.open == "#" || style.open == ";" {
        return Some(
            content
                .iter()
                .map(|line| format!("{}{}", style.line, line))
                .collect::<Vec<String>>()
                .join("\n")
                + "\n\n",
        );
    }

    let mut result = vec![style.open.to_string()];
    result.extend(content.iter().map(|line| format!("{}{}", style.line, line)));
    result.push(style.close.to_string());
    result.push(String::new());
    result.push(String::new());

    Some(result.join("\n"))
}

pub(crate) fn apply_watermark(rendered_template: &RenderedTemplate) -> Result<String> {
    let file_name = rendered_template
        .path
        .file_name()
        .unwrap()
        .to_str()
        .unwrap();
    let maybe_watermark = generate_watermark(file_name);

    let contents = rendered_template.content.clone();
    let watermark = maybe_watermark.clone().unwrap_or(String::new());

    Ok(format!(
        "{}{}",
        if maybe_watermark.is_some() {
            if contents.contains("Generated by ForkLaunch")
                && contents.contains("File:")
                && contents.contains("Modifications are encouraged")
            {
                ""
            } else {
                &watermark
            }
        } else {
            ""
        },
        contents,
    ))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_typescript_watermark() {
        let watermark = generate_watermark("example.ts").unwrap();
        assert!(watermark.starts_with("/**\n"));
        assert!(watermark.contains("Generated by ForkLaunch"));
        assert!(watermark.contains("File: example.ts"));
        assert!(watermark.contains("Modifications are encouraged"));
        assert!(watermark.ends_with("\n\n"));
    }

    #[test]
    fn test_python_watermark() {
        let watermark = generate_watermark("example.py").unwrap();
        assert!(watermark.starts_with("\"\"\"\n"));
        assert!(watermark.contains("Generated by ForkLaunch"));
        assert!(watermark.contains("File: example.py"));
        assert!(watermark.contains("Modifications are encouraged"));
        assert!(watermark.ends_with("\n\n"));
    }

    #[test]
    fn test_yaml_watermark() {
        let watermark = generate_watermark("example.yml").unwrap();
        assert!(watermark.starts_with("# "));
        assert!(watermark.contains("Generated by ForkLaunch"));
        assert!(watermark.contains("File: example.yml"));
        assert!(watermark.contains("Modifications are encouraged"));
        assert!(watermark.ends_with("\n\n"));
    }
}