use regex::{Captures, Regex};
use std::collections::HashMap;
use std::sync::LazyLock;

/// Post-comrak pass: inject `id` attributes into `<h1>`–`<h6>` tags that
/// don't already have one. IDs are generated by slugifying the text content.
/// Duplicate slugs get `-1`, `-2`, etc. appended.
pub fn inject_heading_ids(html: &str) -> String {
    static HEADING_RE: LazyLock<Regex> =
        LazyLock::new(|| Regex::new(r"<(h[1-6])([^>]*)>(.*?)</h[1-6]>").unwrap());
    static TAG_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"<[^>]+>").unwrap());
    static ID_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r#"\bid\s*=\s*""#).unwrap());

    let mut seen: HashMap<String, usize> = HashMap::new();

    HEADING_RE
        .replace_all(html, |caps: &Captures| {
            let tag = &caps[1];
            let attrs = &caps[2];
            let inner = &caps[3];

            // Skip if there's already an id attribute
            if ID_RE.is_match(attrs) {
                return caps[0].to_string();
            }

            // Strip HTML tags from inner content to get plain text
            let text = TAG_RE.replace_all(inner, "");
            let base_slug = slug::slugify(text.trim());

            let id = if base_slug.is_empty() {
                "heading".to_string()
            } else {
                base_slug
            };

            // Handle duplicates
            let count = seen.entry(id.clone()).or_insert(0);
            let final_id = if *count == 0 {
                *count += 1;
                id
            } else {
                let suffixed = format!("{}-{}", id, count);
                *count += 1;
                suffixed
            };

            format!("<{} id=\"{}\"{}>{}</{}>", tag, final_id, attrs, inner, tag)
        })
        .into_owned()
}

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

    #[test]
    fn basic_heading_gets_id() {
        let html = "<h2>Installation Guide</h2>";
        let result = inject_heading_ids(html);
        assert_eq!(result, r#"<h2 id="installation-guide">Installation Guide</h2>"#);
    }

    #[test]
    fn heading_with_existing_id_not_modified() {
        let html = r#"<h2 id="custom-id">Title</h2>"#;
        let result = inject_heading_ids(html);
        assert_eq!(result, html);
    }

    #[test]
    fn duplicate_headings_get_suffixed_ids() {
        let html = "<h2>Foo</h2>\n<h2>Foo</h2>\n<h2>Foo</h2>";
        let result = inject_heading_ids(html);
        assert!(result.contains(r#"id="foo""#));
        assert!(result.contains(r#"id="foo-1""#));
        assert!(result.contains(r#"id="foo-2""#));
    }

    #[test]
    fn tags_inside_headings_stripped_for_slug() {
        let html = "<h2>Hello <code>world</code></h2>";
        let result = inject_heading_ids(html);
        assert!(result.contains(r#"id="hello-world""#));
        // Inner HTML preserved
        assert!(result.contains("<code>world</code>"));
    }

    #[test]
    fn h1_through_h6_all_handled() {
        let html = "<h1>A</h1><h3>B</h3><h6>C</h6>";
        let result = inject_heading_ids(html);
        assert!(result.contains(r#"<h1 id="a">"#));
        assert!(result.contains(r#"<h3 id="b">"#));
        assert!(result.contains(r#"<h6 id="c">"#));
    }

    #[test]
    fn heading_with_em_tag() {
        let html = "<h2>Getting <em>Started</em></h2>";
        let result = inject_heading_ids(html);
        assert!(result.contains(r#"id="getting-started""#));
    }

    #[test]
    fn preserves_existing_classes() {
        let html = r#"<h2 class="special">Title</h2>"#;
        let result = inject_heading_ids(html);
        assert!(result.contains(r#"id="title""#));
        assert!(result.contains(r#"class="special""#));
    }
}
