ssukka 0.2.0

use crate::config::{JsStringEncoding, ObfuscationConfig};
use crate::css;
use crate::error::Result;
use crate::honeypot;
use crate::html::entities;
use crate::html::tags;
use crate::html::whitespace;
use crate::js;
use crate::js_ast;
use crate::structural;
use crate::symbol_map::SymbolMap;
use lol_html::html_content::ContentType;
use lol_html::{doc_comments, element, text, HtmlRewriter, Settings};
use rand::rngs::StdRng;
use rand::SeedableRng;
use std::cell::RefCell;
use std::rc::Rc;

/// ID-referencing attributes that need consistent renaming.
const ID_REF_ATTRS: &[&str] = &[
    "for",
    "aria-labelledby",
    "aria-describedby",
    "aria-controls",
    "aria-owns",
    "aria-activedescendant",
    "aria-flowto",
    "form",
    "headers",
    "list",
    "popovertarget",
];

/// Attributes that should not be entity-encoded (URLs, IDs, etc.)
fn should_skip_attr_encoding(name: &str) -> bool {
    matches!(
        name,
        "class" | "id" | "src" | "href" | "action" | "style" | "type" | "name" | "value"
    ) || ID_REF_ATTRS.contains(&name)
}

/// Pass 2: Apply all obfuscation transformations.
pub fn transform(html: &str, symbols: &SymbolMap, config: &ObfuscationConfig) -> Result<String> {
    let output = RefCell::new(Vec::with_capacity(html.len()));
    let rng = RefCell::new(match config.seed {
        Some(s) => StdRng::seed_from_u64(s),
        None => StdRng::from_rng(&mut rand::rng()),
    });

    let remove_comments = config.remove_comments;
    let collapse_ws = config.collapse_whitespace;
    let encode_text = config.encode_text_entities;
    let encode_attrs = config.encode_attr_entities;
    let shuffle_attrs = config.shuffle_attributes;
    let randomize_case = config.randomize_tag_case;
    let rename_classes = config.rename_classes;
    let rename_ids = config.rename_ids;
    let minify_css = config.minify_css;
    let unicode_escape = config.unicode_escape_selectors;
    // `Array` encoding needs the AST engine; without it, degrade to escapes.
    let js_encoding = match config.js_string_encoding {
        JsStringEncoding::Array if !config.js_ast => JsStringEncoding::Escapes,
        other => other,
    };
    let minify_js_opt = config.minify_js;
    let inject_honeypots = config.inject_honeypots;
    let honeypot_count = config.honeypot_count;
    let structural_obf = config.structural_obfuscation;
    let wants_ast = config.wants_ast();

    // Track preserved-whitespace context (Rc for sharing with end_tag_handlers)
    let preserved_depth = Rc::new(RefCell::new(0u32));

    // Stack of open element tag names (innermost last), used to find the direct
    // parent of a text node for structural obfuscation. Only elements with an
    // end tag are pushed, so void elements never unbalance it.
    let tag_stack: Rc<RefCell<Vec<String>>> = Rc::new(RefCell::new(Vec::new()));

    // Track whether we're inside a <style> or <script> RAWTEXT element.
    // Text inside these must NOT be entity-encoded - browsers parse them as raw text.
    let in_raw_text = Rc::new(RefCell::new(false));

    // Track whether the current <script> is actual JavaScript (not JSON, template, etc.)
    let script_is_js = RefCell::new(true);

    // Accumulation buffers for style/script text (lol_html may split text chunks)
    let style_buf: RefCell<String> = RefCell::new(String::new());
    let script_buf: RefCell<String> = RefCell::new(String::new());

    let mut element_handlers: Vec<_> = Vec::new();
    let mut document_handlers: Vec<_> = Vec::new();

    if remove_comments {
        document_handlers.push(doc_comments!(|comment| {
            comment.remove();
            Ok(())
        }));
    }

    element_handlers.push(element!("style", |el| {
        *style_buf.borrow_mut() = String::new();
        *in_raw_text.borrow_mut() = true;
        if let Some(handlers) = el.end_tag_handlers() {
            let flag = Rc::clone(&in_raw_text);
            let handler: lol_html::EndTagHandler<'static> = Box::new(move |_end| {
                *flag.borrow_mut() = false;
                Ok(())
            });
            handlers.push(handler);
        }
        Ok(())
    }));

    element_handlers.push(element!("script", |el| {
        *script_buf.borrow_mut() = String::new();
        *in_raw_text.borrow_mut() = true;

        let is_js = match el.get_attribute("type") {
            Some(t) => {
                let t = t.to_ascii_lowercase();
                t.is_empty() || t.contains("javascript") || t.contains("ecmascript")
            },
            None => true,
        };
        *script_is_js.borrow_mut() = is_js;

        if let Some(handlers) = el.end_tag_handlers() {
            let flag = Rc::clone(&in_raw_text);
            let handler: lol_html::EndTagHandler<'static> = Box::new(move |_end| {
                *flag.borrow_mut() = false;
                Ok(())
            });
            handlers.push(handler);
        }
        Ok(())
    }));

    element_handlers.push(element!("*", |el| {
        let tag_lower = el.tag_name().to_ascii_lowercase();

        // Maintain the open-element stack (only for elements that have an end
        // tag, keeping it balanced regardless of void/self-closing elements).
        if structural_obf {
            if let Some(handlers) = el.end_tag_handlers() {
                tag_stack.borrow_mut().push(tag_lower.clone());
                let stack = Rc::clone(&tag_stack);
                let handler: lol_html::EndTagHandler<'static> = Box::new(move |_end| {
                    stack.borrow_mut().pop();
                    Ok(())
                });
                handlers.push(handler);
            }
        }

        if whitespace::is_preserved_tag(&tag_lower) {
            *preserved_depth.borrow_mut() += 1;
            if let Some(handlers) = el.end_tag_handlers() {
                let depth = Rc::clone(&preserved_depth);
                let handler: lol_html::EndTagHandler<'static> = Box::new(move |_end| {
                    let mut d = depth.borrow_mut();
                    if *d > 0 {
                        *d -= 1;
                    }
                    Ok(())
                });
                handlers.push(handler);
            }
        }

        if rename_classes {
            if let Some(class_attr) = el.get_attribute("class") {
                let new_classes: Vec<&str> = class_attr
                    .split_whitespace()
                    .map(|c| symbols.get_class(c).unwrap_or(c))
                    .collect();
                el.set_attribute("class", &new_classes.join(" "))?;
            }
        }

        if rename_ids {
            if let Some(id_attr) = el.get_attribute("id") {
                if let Some(new_id) = symbols.get_id(&id_attr) {
                    el.set_attribute("id", new_id)?;
                }
            }

            for &attr_name in ID_REF_ATTRS {
                if let Some(value) = el.get_attribute(attr_name) {
                    let new_value: Vec<&str> = value
                        .split_whitespace()
                        .map(|id| symbols.get_id(id).unwrap_or(id))
                        .collect();
                    el.set_attribute(attr_name, &new_value.join(" "))?;
                }
            }

            if let Some(href) = el.get_attribute("href") {
                if let Some(id) = href.strip_prefix('#') {
                    if let Some(new_id) = symbols.get_id(id) {
                        el.set_attribute("href", &format!("#{new_id}"))?;
                    }
                }
            }
        }

        // Skip style/script for further obfuscation (entity encoding, shuffle, case)
        if tag_lower == "style" || tag_lower == "script" {
            return Ok(());
        }

        // Encode attribute values (skip functional attrs like IDs, URLs)
        if encode_attrs {
            let mut rng = rng.borrow_mut();
            let attrs: Vec<(String, String)> = el.attributes().iter().map(|a| (a.name(), a.value())).collect();
            for (name, value) in &attrs {
                if should_skip_attr_encoding(name) {
                    continue;
                }
                let encoded = entities::encode_attr_value(value, &mut rng);
                el.set_attribute(name, &encoded)?;
            }
        }

        if shuffle_attrs {
            let mut rng = rng.borrow_mut();
            let mut attrs: Vec<(String, String)> = el.attributes().iter().map(|a| (a.name(), a.value())).collect();
            tags::shuffle_attributes(&mut attrs, &mut rng);

            let attr_names: Vec<String> = el.attributes().iter().map(|a| a.name()).collect();
            for name in &attr_names {
                el.remove_attribute(name);
            }
            for (name, value) in &attrs {
                el.set_attribute(name, value)?;
            }
        }

        if randomize_case {
            let mut rng = rng.borrow_mut();
            let new_tag = tags::randomize_tag_case(&tag_lower, &mut rng);
            el.set_tag_name(&new_tag)?;
        }

        Ok(())
    }));

    // General text handler - applies to ALL text nodes, but skips RAWTEXT context
    element_handlers.push(text!("*", |text| {
        // Skip text inside <style> and <script> - handled by dedicated handlers
        if *in_raw_text.borrow() {
            return Ok(());
        }

        let is_preserved = *preserved_depth.borrow() > 0;
        let content = text.as_str().to_owned();

        if content.is_empty() {
            return Ok(());
        }

        let mut processed = content;

        if collapse_ws && !is_preserved {
            processed = whitespace::collapse_whitespace(&processed);
        }

        // Structural obfuscation: relocate non-blank text inside safe flow
        // elements into an encoded data-attribute, restored client-side.
        if structural_obf && !is_preserved && !processed.trim().is_empty() {
            let parent_is_safe = tag_stack
                .borrow()
                .last()
                .map(|t| structural::is_safe_tag(t))
                .unwrap_or(false);
            if parent_is_safe {
                text.replace(&structural::encode_text_node(&processed), ContentType::Html);
                return Ok(());
            }
        }

        if encode_text {
            let mut rng = rng.borrow_mut();
            processed = entities::encode_entities(&processed, &mut rng);
        }

        text.replace(&processed, ContentType::Html);
        Ok(())
    }));

    // Inject honeypots and the structural-restore script at the end of <body>.
    if inject_honeypots || structural_obf {
        element_handlers.push(element!("body", |el| {
            if inject_honeypots {
                let mut rng = rng.borrow_mut();
                let decoys = honeypot::generate(honeypot_count, &mut rng);
                el.append(&decoys, ContentType::Html);
            }
            if structural_obf {
                el.append(structural::restore_script(), ContentType::Html);
            }
            Ok(())
        }));
    }

    element_handlers.push(text!("style", |text| {
        style_buf.borrow_mut().push_str(text.as_str());

        if text.last_in_text_node() {
            let css_text = style_buf.borrow().clone();
            if !css_text.is_empty() {
                if let Ok(transformed) = css::transform_css(
                    &css_text,
                    symbols,
                    minify_css,
                    unicode_escape,
                    rename_classes,
                    rename_ids,
                ) {
                    text.replace(&transformed, ContentType::Html);
                }
            }
            *style_buf.borrow_mut() = String::new();
        } else {
            text.remove();
        }
        Ok(())
    }));

    element_handlers.push(text!("script", |text| {
        script_buf.borrow_mut().push_str(text.as_str());

        if text.last_in_text_node() {
            let content = script_buf.borrow().clone();
            if !content.is_empty() {
                if *script_is_js.borrow() {
                    // JavaScript: AST engine when requested, else the token path.
                    // The AST path falls back to the token path on parse failure.
                    let mut rng_ref = rng.borrow_mut();
                    let token_path = |rng: &mut StdRng| {
                        js::transform_js(
                            &content,
                            symbols,
                            js_encoding,
                            minify_js_opt,
                            rename_classes,
                            rename_ids,
                            rng,
                        )
                    };
                    let transformed = if wants_ast {
                        js_ast::transform(&content, symbols, config, &mut rng_ref)
                            .unwrap_or_else(|| token_path(&mut rng_ref))
                    } else {
                        token_path(&mut rng_ref)
                    };
                    text.replace(&transformed, ContentType::Html);
                } else if rename_classes || rename_ids {
                    // Non-JS (JSON, etc.): only rename class/ID references
                    let transformed = js::replace_symbols_word_boundary(&content, symbols, rename_classes, rename_ids);
                    text.replace(&transformed, ContentType::Html);
                }
            }
            *script_buf.borrow_mut() = String::new();
        } else {
            text.remove();
        }
        Ok(())
    }));

    {
        let mut settings = Settings::new();
        for handler in element_handlers {
            settings = settings.append_element_content_handler(handler);
        }
        for handler in document_handlers {
            settings = settings.append_document_content_handler(handler);
        }
        let mut rewriter = HtmlRewriter::new(settings, |chunk: &[u8]| {
            output.borrow_mut().extend_from_slice(chunk);
        });

        rewriter.write(html.as_bytes())?;
        rewriter.end()?;
    }

    let bytes = output.into_inner();
    String::from_utf8(bytes).map_err(|e| crate::error::SsukkaError::Rewrite(e.to_string()))
}

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

    fn default_config() -> ObfuscationConfig {
        ObfuscationConfig {
            seed: Some(42),
            ..Default::default()
        }
    }

    #[test]
    fn removes_comments() {
        let html = "<div><!-- comment -->text</div>";
        let symbols = SymbolMap::new(Some(42));
        let result = transform(html, &symbols, &default_config()).unwrap();
        assert!(!result.contains("comment"));
    }

    #[test]
    fn preserves_pre_whitespace() {
        let html = "<pre>  code  \n  here  </pre>";
        let config = ObfuscationConfig {
            seed: Some(42),
            encode_text_entities: false,
            randomize_tag_case: false,
            shuffle_attributes: false,
            ..Default::default()
        };
        let symbols = SymbolMap::new(Some(42));
        let result = transform(html, &symbols, &config).unwrap();
        assert!(result.contains("  code  \n  here  "));
    }

    #[test]
    fn renames_classes_consistently() {
        let html = r#"<div class="foo">text</div>"#;
        let config = default_config();
        let mut symbols = SymbolMap::new(Some(42));
        symbols.register_class("foo");
        let obf = symbols.get_class("foo").unwrap().to_owned();
        let result = transform(html, &symbols, &config).unwrap();
        assert!(result.contains(&obf));
    }

    #[test]
    fn renames_href_id_refs() {
        let html = r##"<a href="#sec">link</a><div id="sec">content</div>"##;
        let config = default_config();
        let mut symbols = SymbolMap::new(Some(42));
        symbols.register_id("sec");
        let obf = symbols.get_id("sec").unwrap().to_owned();
        let result = transform(html, &symbols, &config).unwrap();
        assert!(result.contains(&format!("#{obf}")));
    }

    #[test]
    fn style_not_entity_encoded() {
        let html = "<style>.foo { color: red; }</style>";
        let config = ObfuscationConfig {
            seed: Some(42),
            rename_classes: false,
            minify_css: false,
            unicode_escape_selectors: false,
            randomize_tag_case: false,
            shuffle_attributes: false,
            ..Default::default()
        };
        let symbols = SymbolMap::new(Some(42));
        let result = transform(html, &symbols, &config).unwrap();
        // CSS content should NOT contain HTML entities
        assert!(
            !result.contains("&#"),
            "Style content should not be entity-encoded. Got: {result}"
        );
        assert!(result.contains("color"));
    }

    #[test]
    fn script_not_entity_encoded() {
        let html = r#"<script>var x = 1 + 2;</script>"#;
        let config = ObfuscationConfig {
            seed: Some(42),
            js_string_encoding: JsStringEncoding::None,
            minify_js: false,
            randomize_tag_case: false,
            shuffle_attributes: false,
            ..Default::default()
        };
        let symbols = SymbolMap::new(Some(42));
        let result = transform(html, &symbols, &config).unwrap();
        assert!(
            result.contains("var x = 1 + 2"),
            "Script content should not be entity-encoded. Got: {result}"
        );
    }
}