use lol_html::{
    doc_comments, doctype, element,
    html_content::{Element, TextChunk},
    text, DocumentContentHandlers, ElementContentHandlers, HtmlRewriter, MemorySettings, Selector,
    Settings,
};
use magnus::{
    function, gc, method,
    r_hash::ForEach,
    scan_args,
    typed_data::Obj,
    value::{Opaque, ReprValue},
    DataTypeFunctions, Integer, IntoValue, Module, Object, RArray, RHash, RModule, Ruby, Symbol,
    TypedData, Value,
};

use std::{
    borrow::Cow,
    cell::{Ref, RefCell},
    mem,
    ops::Deref,
    primitive::str,
    rc::Rc,
};

use crate::{
    html::{element::SelmaHTMLElement, end_tag::SelmaHTMLEndTag, text_chunk::SelmaHTMLTextChunk},
    native_ref_wrap::NativeRefWrap,
    sanitizer::SelmaSanitizer,
    selector::SelmaSelector,
    tags::Tag,
};

#[derive(Clone)]
pub struct Handler {
    rb_handler: Opaque<Value>,
    // Store the selector's data by value (Rust-owned). The `#selector` method may return a
    // freshly-built `Selma::Selector` that nothing else on the Ruby side references; keeping a
    // `Obj`/`Opaque` handle to it would make its lifetime depend on GC (use-after-free if it is
    // collected while the Rewriter is alive). We only ever read Rust data off it, so clone it.
    selector: SelmaSelector,
    // total_element_handler_calls: usize,
    // total_elapsed_element_handlers: f64,

    // total_text_handler_calls: usize,
    // total_elapsed_text_handlers: f64,
}

struct RewriterOptions {
    max_allowed_memory_usage: usize,
    preallocated_parsing_buffer_size: usize,
}

pub struct Rewriter {
    sanitizer: Option<SelmaSanitizer>,
    handlers: Vec<Handler>,
    options: RewriterOptions,
    // total_elapsed: f64,
}

#[derive(TypedData)]
#[magnus(class = "Selma::Rewriter", free_immediately, mark)]
pub struct SelmaRewriter(std::cell::RefCell<Rewriter>);

impl DataTypeFunctions for SelmaRewriter {
    fn mark(&self, marker: &gc::Marker) {
        self.0.borrow().handlers.iter().for_each(|handler| {
            marker.mark(handler.rb_handler);
        });
    }
}

type RewriterValues = (
    Option<Option<Obj<SelmaSanitizer>>>,
    Option<RArray>,
    Option<RHash>,
);

impl SelmaRewriter {
    const SELMA_ON_END_TAG: &'static str = "on_end_tag";
    const SELMA_HANDLE_ELEMENT: &'static str = "handle_element";
    const SELMA_HANDLE_TEXT_CHUNK: &'static str = "handle_text_chunk";

    /// @yard
    /// @def new(sanitizer: Selma::Sanitizer.new(Selma::Sanitizer::Config::DEFAULT), handlers: [])
    /// @param sanitizer [Selma::Sanitizer] The sanitizer which performs the initial cleanup
    /// @param handlers  [Array<Selma::Selector>] The handlers to use to perform HTML rewriting
    /// @param options  [Hash] Any additional options to pass to the rewriter
    /// @return [Selma::Rewriter]
    fn new(args: &[Value]) -> Result<Self, magnus::Error> {
        let (rb_sanitizer, rb_handlers, rb_options) = Self::scan_parse_args(args)?;
        let ruby = Ruby::get().unwrap();

        let sanitizer = match rb_sanitizer {
            None => {
                // no `sanitizer:` kwarg provided, use default
                let default_sanitizer = SelmaSanitizer::new(&[])?;
                let wrapped_sanitizer = ruby.obj_wrap(default_sanitizer);
                // wrapped_sanitizer.funcall::<&str, (), Value>("setup", ())?;
                Some(wrapped_sanitizer.deref().to_owned())
            }
            Some(sanitizer_value) => sanitizer_value.map(|sanitizer| sanitizer.deref().to_owned()),
        };

        let handlers = match rb_handlers {
            None => vec![],
            Some(rb_handlers) => {
                let mut handlers: Vec<Handler> = vec![];

                for rb_handler in rb_handlers.into_iter() {
                    // prevents missing #selector from ruining things
                    if !rb_handler.respond_to("selector", true).unwrap() {
                        let classname = unsafe { rb_handler.classname() };
                        return Err(magnus::Error::new(
                            ruby.exception_no_method_error(),
                            format!(
                                "Could not call #selector on {classname:?}; is this an object that defines it?",

                            ),
                        ));
                    }

                    let rb_selector: Obj<SelmaSelector> = match rb_handler.funcall("selector", ()) {
                        Err(err) => {
                            return Err(magnus::Error::new(
                                ruby.exception_type_error(),
                                format!("Error instantiating selector: {err:?}"),
                            ));
                        }
                        Ok(rb_selector) => rb_selector,
                    };
                    let handler = Handler {
                        rb_handler: Opaque::from(rb_handler),
                        // clone the selector's data out of the Ruby object right away so the
                        // Handler no longer depends on that object surviving GC (see struct docs)
                        selector: (*rb_selector).clone(),
                        // total_element_handler_calls: 0,
                        // total_elapsed_element_handlers: 0.0,

                        // total_text_handler_calls: 0,
                        // total_elapsed_text_handlers: 0.0,
                    };
                    handlers.push(handler);
                }
                handlers
            }
        };

        if sanitizer.is_none() && handlers.is_empty() {
            return Err(magnus::Error::new(
                ruby.exception_arg_error(),
                "Must provide a sanitizer or a handler",
            ));
        }

        let mut rewriter_options = RewriterOptions::new();

        match rb_options {
            None => {}
            Some(options) => {
                options.foreach(|key: Symbol, value: RHash| {
                    let ruby = Ruby::get().unwrap();
                    let key = key.to_string();
                    match key.as_str() {
                        "memory" => {
                            if let Some(max_allowed_memory_usage) = value.get(ruby.to_symbol("max_allowed_memory_usage")) {
                                let max_allowed_memory_usage =
                                    Integer::from_value(max_allowed_memory_usage);
                                if let Some(max_allowed_memory_usage) = max_allowed_memory_usage {
                                    match max_allowed_memory_usage.to_u64() {
                                        Ok(max_allowed_memory_usage) => {
                                            rewriter_options.max_allowed_memory_usage =
                                                max_allowed_memory_usage as usize;
                                        }
                                        Err(_e) => {
                                            return Err(magnus::Error::new(
                                                ruby.exception_arg_error(),
                                                "max_allowed_memory_usage must be a positive integer",
                                            ));
                                        }
                                    }
                                } else {
                                    rewriter_options.max_allowed_memory_usage = RewriterOptions::DEFAULT_MAX_ALLOWED_MEMORY_USAGE;
                                }
                            }

                            if let Some(preallocated_parsing_buffer_size) = value.get(ruby.to_symbol("preallocated_parsing_buffer_size")) {
                                let preallocated_parsing_buffer_size =
                                    Integer::from_value(preallocated_parsing_buffer_size);
                                if let Some(preallocated_parsing_buffer_size) = preallocated_parsing_buffer_size {
                                    match preallocated_parsing_buffer_size.to_u64() {
                                        Ok(preallocated_parsing_buffer_size) => {
                                            rewriter_options.preallocated_parsing_buffer_size =
                                                preallocated_parsing_buffer_size as usize;
                                        }
                                        Err(_e) => {
                                            return Err(magnus::Error::new(
                                                ruby.exception_arg_error(),
                                                "preallocated_parsing_buffer_size must be a positive integer",
                                            ));
                                        }
                                    }
                                } else {
                                    rewriter_options.preallocated_parsing_buffer_size = RewriterOptions::DEFAULT_PREALLOCATED_PARSING_BUFFER_SIZE;
                                }
                            }
                        }
                        _ => {
                            return Err(magnus::Error::new(
                                ruby.exception_arg_error(),
                                format!("Unknown option: {key:?}"),
                            ));
                        }
                    }
                    Ok(ForEach::Continue)
                })?;
            }
        }

        if rewriter_options.preallocated_parsing_buffer_size
            > rewriter_options.max_allowed_memory_usage
        {
            return Err(magnus::Error::new(
                ruby.exception_arg_error(),
                "max_allowed_memory_usage must be greater than preallocated_parsing_buffer_size",
            ));
        }

        Ok(Self(std::cell::RefCell::new(Rewriter {
            sanitizer,
            handlers,
            options: rewriter_options,
            // total_elapsed: 0.0,
        })))
    }

    #[allow(clippy::let_unit_value)]
    fn scan_parse_args(args: &[Value]) -> Result<RewriterValues, magnus::Error> {
        let args = scan_args::scan_args(args)?;
        let _: () = args.required;
        let _: () = args.optional;
        let _: () = args.splat;
        let _: () = args.trailing;
        let _: () = args.block;

        let kwargs = scan_args::get_kwargs::<
            _,
            (),
            (
                Option<Option<Obj<SelmaSanitizer>>>,
                Option<RArray>,
                Option<RHash>,
            ),
            (),
        >(args.keywords, &[], &["sanitizer", "handlers", "options"])?;
        let (rb_sanitizer, rb_handlers, rb_options) = kwargs.optional;

        Ok((rb_sanitizer, rb_handlers, rb_options))
    }

    /// Perform HTML rewrite sequence.
    fn rewrite(&self, html: String) -> Result<String, magnus::Error> {
        let binding = self.0.borrow();

        let mut sanitizer_document_content_handlers: Vec<DocumentContentHandlers> = vec![];
        let mut sanitizer_element_content_handlers: Vec<(Cow<Selector>, ElementContentHandlers)> =
            vec![];

        match &binding.sanitizer {
            None => (),
            Some(sanitizer) => {
                if !sanitizer.get_allow_doctype() {
                    sanitizer_document_content_handlers.push(doctype!(|d| {
                        sanitizer.remove_doctype(d);
                        Ok(())
                    }));
                }
                if !sanitizer.get_allow_comments() {
                    sanitizer_document_content_handlers.push(doc_comments!(|c| {
                        sanitizer.remove_comment(c);
                        Ok(())
                    }));
                }
                sanitizer_element_content_handlers.push(element!("*", |el| {
                    sanitizer.try_remove_element(el);
                    if el.removed() {
                        return Ok(());
                    }
                    // if it was removed, there are no attributes to sanitize
                    match sanitizer.sanitize_attributes(el) {
                        Ok(_) => Ok(()),
                        Err(err) => Err(err.to_string().into()),
                    }
                }));
            }
        };

        let handlers: &Vec<Handler> = &binding.handlers;

        match Self::perform_handler_rewrite(
            self,
            sanitizer_document_content_handlers,
            sanitizer_element_content_handlers,
            handlers,
            html,
        ) {
            Ok(rewritten_html) => match &binding.sanitizer {
                None => match String::from_utf8(rewritten_html) {
                    Ok(output) => Ok(output),
                    Err(err) => Err(magnus::Error::new(
                        Ruby::get().unwrap().exception_runtime_error(),
                        format!("{err:?}"),
                    )),
                },
                Some(sanitizer) => {
                    Self::perform_final_sanitization(self, sanitizer, rewritten_html)
                }
            },
            Err(err) => Err(err),
        }
    }

    // to get rid of some really nasty edge cases with dangerous tags, we perform one more
    // sanitization pass at the end
    fn perform_final_sanitization(
        &self,
        sanitizer: &SelmaSanitizer,
        html: Vec<u8>,
    ) -> Result<String, magnus::Error> {
        // TODO: this should ideally be done ahead of time on `initialize`, not on every `#rewrite` call
        let mut element_content_handlers: Vec<(Cow<Selector>, ElementContentHandlers)> = vec![];

        if sanitizer.get_escape_tagfilter() {
            element_content_handlers.push(element!(Tag::ESCAPEWORTHY_TAGS_CSS, |el| {
                let should_remove = sanitizer.allow_element(el);
                if should_remove {
                    sanitizer.force_remove_element(el);
                }

                Ok(())
            }));
        }

        match Self::run_rewrite(self, vec![], element_content_handlers, html.as_slice()) {
            Ok(rewritten_html) => match String::from_utf8(rewritten_html) {
                Ok(output) => Ok(output),
                Err(err) => Err(magnus::Error::new(
                    Ruby::get().unwrap().exception_runtime_error(),
                    format!("{err:?}"),
                )),
            },
            Err(err) => Err(err),
        }
    }

    pub fn perform_handler_rewrite<'a>(
        &self,
        sanitizer_document_content_handlers: Vec<DocumentContentHandlers<'a>>,
        sanitizer_element_content_handlers: Vec<(Cow<Selector>, ElementContentHandlers<'a>)>,
        handlers: &'a [Handler],
        html: String,
    ) -> Result<Vec<u8>, magnus::Error> {
        // TODO: this should ideally be done ahead of time on `initialize`, not on every `#rewrite` call
        let mut element_content_handlers: Vec<(Cow<Selector>, ElementContentHandlers)> = vec![];

        // have sanitization happen first
        element_content_handlers.extend(sanitizer_element_content_handlers);

        handlers.iter().for_each(|handler| {
            let element_stack: Rc<RefCell<Vec<String>>> = Rc::new(RefCell::new(vec![]));

            let selector = &handler.selector;

            // TODO: test final raise by simulating errors
            if let Some(match_element) = selector.match_element() {
                let closure_element_stack = element_stack.clone();

                element_content_handlers.push(element!(match_element, move |el| {
                    match Self::process_element_handlers(
                        handler,
                        el,
                        &closure_element_stack.borrow(),
                    ) {
                        Ok(_) => Ok(()),
                        Err(err) => Err(err.to_string().into()),
                    }
                }));
            }

            if let Some(match_text_within) = selector.match_text_within() {
                let closure_element_stack = element_stack.clone();

                element_content_handlers.push(text!(match_text_within, move |text| {
                    let element_stack = closure_element_stack.as_ref().borrow();
                    // check if current tag is a tag we should be ignoring text within;
                    // also checks if tag is within an ancestery of ignored tags
                    if let Some(ignore_text_within) = selector.ignore_text_within() {
                        if ignore_text_within.iter().any(|t| element_stack.contains(t)) {
                            return Ok(());
                        }
                    }

                    match Self::process_text_handlers(handler, text) {
                        Ok(_) => Ok(()),
                        Err(err) => Err(err.to_string().into()),
                    }
                }));
            }

            // we need to check *every* element we iterate over, to create a stack of elements
            element_content_handlers.push(element!("*", move |el| {
                let tag_name = el.tag_name().to_lowercase();

                // no need to track self-closing tags
                if Tag::tag_from_tag_name(&tag_name).self_closing {
                    return Ok(());
                };

                element_stack.as_ref().borrow_mut().push(tag_name);

                let closure_element_stack = element_stack.clone();

                let handler: lol_html::EndTagHandler<'static> = Box::new(move |_end_tag| {
                    closure_element_stack.as_ref().borrow_mut().pop();
                    Ok(())
                });
                // ignore void elements (lol_html's void list may differ from selma's `self_closing`)
                let _ = el.on_end_tag(handler);

                Ok(())
            }));
        });

        Self::run_rewrite(
            self,
            sanitizer_document_content_handlers,
            element_content_handlers,
            html.as_bytes(),
        )
    }

    fn run_rewrite<'a>(
        &self,
        document_content_handlers: Vec<DocumentContentHandlers<'a>>,
        element_content_handlers: Vec<(Cow<Selector>, ElementContentHandlers<'a>)>,
        html: &[u8],
    ) -> Result<Vec<u8>, magnus::Error> {
        let binding = &self.0.borrow();
        let mut output = vec![];
        {
            let mut settings =
                Settings::new().with_memory_settings(Self::get_memory_options(binding));
            for handler in document_content_handlers {
                settings = settings.append_document_content_handler(handler);
            }
            for handler in element_content_handlers {
                settings = settings.append_element_content_handler(handler);
            }

            let mut rewriter = HtmlRewriter::new(settings, |c: &[u8]| output.extend_from_slice(c));
            match rewriter.write(html) {
                Ok(_) => {}
                Err(err) => {
                    return Err(magnus::Error::new(
                        Ruby::get().unwrap().exception_runtime_error(),
                        format!("{err:?}"),
                    ));
                }
            }
        }
        Ok(output)
    }

    fn process_element_handlers(
        handler: &Handler,
        element: &mut Element,
        ancestors: &[String],
    ) -> Result<(), magnus::Error> {
        let ruby = Ruby::get().unwrap();
        let rb_handler = handler.rb_handler.into_value_with(&ruby);

        // if `on_end_tag` function is defined, call it
        if rb_handler.respond_to(Self::SELMA_ON_END_TAG, true).unwrap() {
            element
                .on_end_tag(Box::new(move |end_tag| {
                    let (ref_wrap, anchor) = NativeRefWrap::wrap(end_tag);

                    let rb_end_tag = SelmaHTMLEndTag::new(ref_wrap);

                    let result =
                        rb_handler.funcall::<_, _, Value>(Self::SELMA_ON_END_TAG, (rb_end_tag,));

                    mem::drop(anchor);

                    match result {
                        Ok(_) => Ok(()),
                        Err(err) => Err(err.to_string().into()),
                    }
                }))
                .map_err(|err| {
                    magnus::Error::new(
                        Ruby::get().unwrap().exception_runtime_error(),
                        err.to_string(),
                    )
                })?;
        }

        let (ref_wrap, anchor) = NativeRefWrap::wrap(element);
        let rb_element = SelmaHTMLElement::new(ref_wrap, ancestors);
        let result = rb_handler.funcall::<_, _, Value>(Self::SELMA_HANDLE_ELEMENT, (rb_element,));

        mem::drop(anchor);

        match result {
            Ok(_) => Ok(()),
            Err(err) => Err(magnus::Error::new(
                ruby.exception_runtime_error(),
                format!("{err:?}"),
            )),
        }
    }

    fn process_text_handlers(
        handler: &Handler,
        text_chunk: &mut TextChunk,
    ) -> Result<(), magnus::Error> {
        let ruby = Ruby::get().unwrap();
        let rb_handler = handler.rb_handler.into_value_with(&ruby);

        // prevents missing `handle_text_chunk` function
        let content = text_chunk.as_str();

        // lol-html sometimes returns blank text if
        // last_in_text_node() is true
        if content.is_empty() {
            return Ok(());
        }

        let (ref_wrap, anchor) = NativeRefWrap::wrap(text_chunk);

        let rb_text_chunk = SelmaHTMLTextChunk::new(ref_wrap);
        let result =
            rb_handler.funcall::<_, _, Value>(Self::SELMA_HANDLE_TEXT_CHUNK, (rb_text_chunk,));

        mem::drop(anchor);

        match result {
            Ok(_) => Ok(()),
            Err(err) => Err(magnus::Error::new(
                ruby.exception_runtime_error(),
                format!("{err:?}"),
            )),
        }
    }

    fn get_memory_options(binding: &Ref<Rewriter>) -> MemorySettings {
        let options = &binding.options;
        MemorySettings::new()
            .with_max_allowed_memory_usage(options.max_allowed_memory_usage)
            .with_preallocated_parsing_buffer_size(options.preallocated_parsing_buffer_size)
    }
}

impl RewriterOptions {
    // Mirrors lol_html's `MemorySettings` defaults, which are no longer exposed as
    // public fields as of lol_html 3.0 (settings are now consuming builders).
    const DEFAULT_MAX_ALLOWED_MEMORY_USAGE: usize = usize::MAX;
    const DEFAULT_PREALLOCATED_PARSING_BUFFER_SIZE: usize = 1024;

    pub fn new() -> Self {
        Self {
            max_allowed_memory_usage: Self::DEFAULT_MAX_ALLOWED_MEMORY_USAGE,
            preallocated_parsing_buffer_size: Self::DEFAULT_PREALLOCATED_PARSING_BUFFER_SIZE,
        }
    }
}

pub fn init(m_selma: RModule) -> Result<(), magnus::Error> {
    let ruby = Ruby::get().unwrap();
    let c_rewriter = m_selma
        .define_class("Rewriter", ruby.class_object())
        .expect("cannot define class Selma::Rewriter");

    c_rewriter.define_singleton_method("new", function!(SelmaRewriter::new, -1))?;
    c_rewriter
        .define_method("rewrite", method!(SelmaRewriter::rewrite, 1))
        .expect("cannot define method `rewrite`");

    Ok(())
}
