import logging
import re
import warnings

from bs4 import (
    BeautifulSoup, 
    PageElement, 
    Tag, 
    Comment, 
    Stylesheet, 
    Script, 
    NavigableString, 
    XMLParsedAsHTMLWarning
)
from transformers import BatchEncoding, ProcessorMixin


logger = logging.getLogger(__name__)
warnings.filterwarnings("ignore", category=XMLParsedAsHTMLWarning)


class HTMLLMProcessor(ProcessorMixin):
    attributes = ["tokenizer"]
    tokenizer_class = "AutoTokenizer"

    def __init__(self, tokenizer=None, merge_threshold: int = 2, *args, **kwargs):
        super().__init__(tokenizer, *args, **kwargs)
        self.tokenizer = tokenizer
        self.merge_threshold = merge_threshold

    def _process_tree(self, root: Tag) -> Tag:

        PRESERVE_TAGS = [
            "a", "address", "audio", "br", "button", "canvas", "col", "embed", "figure", "footer", 
            "form", "frame", "header", "hr", "iframe", "img", "input", "label", "menu",
            "meter", "nav", "option", "output", "picture", "progress", "search", "select",
            "td", "th", "tr", "textarea", "video", "i", "svg"
        ]

        def _can_tag_be_removed(tag: Tag) -> bool:
            is_whitelisted = tag.name in PRESERVE_TAGS
            contains_text = tag.get_text(strip=True) != ""
            has_children = any(isinstance(c, Tag) for c in tag.children)
            return not (is_whitelisted or contains_text or has_children)
    
        try:
            # process all children of the tag
            i = len(root.contents) - 1
            while i >= 0:
                child = root.contents[i]
                # A: comments, stylesheets, scripts
                if isinstance(child, (Comment, Stylesheet, Script)):
                    # remove the element
                    child.extract()
                
                # B: tags
                elif isinstance(child, Tag):
                    # first, recursively process the subtree
                    if len(child.contents) > 0:
                        self._process_tree(child)
                    # then, process the tag
                    if _can_tag_be_removed(child):
                        # remove only if all the conditions are met
                        child.extract()
                    # remove all attributes
                    child.attrs = {}

                # C: texts
                elif isinstance(child, NavigableString) and child.get_text(strip=True) == "":
                    # remove empty text or whitespaces 
                    child.extract()

                # D: decrement the index
                i -= 1
        except RecursionError:
            # recursion limit reached -> remove subtree beginning with root
            root.extract()
        finally:
            # remove all attributes
            root.attrs = {}
            return root

    def _merge_tree(self, root: Tag) -> Tag:

        def _should_be_merged(tag: Tag) -> bool:
            num_children = len(tag.contents)
            has_multiple_less_than_threshold = num_children > 1 and num_children < self.merge_threshold
            has_single_tag_only = num_children == 1 and (not isinstance(tag.contents[0], NavigableString))
            return has_multiple_less_than_threshold or has_single_tag_only
            
        try:
            i = 0
            # process all childern of the root
            while i < len(root.contents):
                current_child = root.contents[i]
                if isinstance(current_child, Tag):  # if child is a HTML tag
                    if root.name != "html" and _should_be_merged(current_child):
                        # merge current child children to the root
                        # remove the current child from the tree
                        current_child = current_child.extract()
                        current_child_children = list(current_child.contents)
                        current_child.clear()
                        # insert the extracted current child children to the root
                        for idx, child in enumerate(current_child_children):
                            root.insert(i + idx, child)
                    else:
                        # recursively merge its children
                        self._merge_tree(current_child)
                        i += 1
                else:
                    i += 1
        except RecursionError:
            # recursion limit reached -> remove subtree beginning with root
            root.extract()
        finally:
            return root

    def preprocess_html(self, html: str) -> str | None:
        try:
            root = BeautifulSoup(html, "lxml").html
            if not isinstance(root, PageElement):
                logger.warning("HTML could not be parsed!")
                return None
            root = self._process_tree(root)
            root = self._merge_tree(root)
            html = re.sub(r"\s+", " ", str(root))
            html = re.sub(r">\s*<", "><", str(html))
            return html
        except Exception as e:
            logger.warning(f"HTML preprocesing failed: {e}")
            return None

    def __call__(self, text: str | list[str], *args, **kwargs) -> BatchEncoding:
        if isinstance(text, str):
            text = self.preprocess_html(text)
        elif isinstance(text, list):
            text = [self.preprocess_html(t) for t in text]
        else:
            raise ValueError("HTMLLMProcessor expects the input to be either str or list[str].")
        return self.tokenizer(text, *args, **kwargs)