import functools
from typing import Callable, List, Tuple

from bs4 import BeautifulSoup
from loguru import logger
import html
from helpers.html_soup import wrap_bare_words, pwndoc_extract_tags_from_forbidden_parents, PwnDocFormatterInstance

BS_HTML_PARSER = 'html.parser'
# BS_HTML_PARSER = 'lxml'  # lxml is MUCH faster, but produces HTML that PwnDoc doesn't accept and there isn't an easy way how to fix it

# PwnDoc HTML example is in tests/pwndoc_valid_template.html


class MdAndHTMLConvertor:

    @classmethod
    def custom_html_to_pwndoc_html(cls, text: str, force_import_of_html_for_basic_tags: bool = False) -> str:
        # The following function does few additional/unnecessary steps, but that shouldn't change the result.
        return cls.custom_md_to_pwndoc_html(text, convert_md=False, force_import_of_html_for_basic_tags=force_import_of_html_for_basic_tags)

    @staticmethod
    def custom_md_to_pwndoc_html(text: str, convert_md: bool = True, force_import_of_html_for_basic_tags: bool = False) -> str:
        answer = text

        # html escape (make sure the special encoding chars survive)
        answer = html.escape(answer)

        if force_import_of_html_for_basic_tags:
            # Burp uses html for both document structure and PoC examples. PoCs are not escaped, so there is no way to distinguish them.
            # The only workaround I see is to escape everything, and unescape only supported tags.
            # Note: This can slightly change the text!
            # todo: Reasonable heuristic seems to be: answer.startswith("&lt;p&gt;") and fid.startswith("burp_")
            answer = MdAndHTMLConvertorLowLevel.force_unescape_basic_tags(answer)  # this only unescapes this formats: <tag> </tag>, so no <a href ...>!

        if convert_md:  # md to HTML
            answer = MdAndHTMLConvertorLowLevel.md_to_html(answer)

        answer = MdAndHTMLConvertorLowLevel.reverse_custom_escaping_of_html(answer)  # custom unescape
        answer = MdAndHTMLConvertorLowLevel.html_to_pwndoc_html(answer)  # pwndoc HTML cleanup

        return answer

    @staticmethod
    def html_to_custom_pwndoc_html(text: str) -> str:
        answer = text
        answer = MdAndHTMLConvertorLowLevel.html_to_pwndoc_html(answer)  # make sure the pwndoc syntax is clean
        answer = MdAndHTMLConvertorLowLevel.custom_escape_html(answer)  # custom escape
        return answer

    @classmethod
    def pwndoc_html_to_md_native(cls, text: str) -> str:
        # note: This export is NOT reversible without loosing data.
        answer = cls.html_to_custom_pwndoc_html(text)
        answer = MdAndHTMLConvertorLowLevel.reverse_custom_escaping_of_html(answer)

        answer = answer.replace("<br>", "\n")
        answer = MdAndHTMLConvertorLowLevel.replace_html_pair_tags_with_md_equivalents(answer)
        answer = MdAndHTMLConvertorLowLevel.strip_unsupported_html_tags(answer)
        answer = html.unescape(answer)
        answer = answer.strip()  # todo: this also have to work for HTML output

        return answer

    @staticmethod
    def unescape_custom_html(text: str) -> str:
        return MdAndHTMLConvertorLowLevel.reverse_custom_escaping_of_html(text)

    @staticmethod
    def html_to_pwndoc_html(text: str) -> str:
        return MdAndHTMLConvertorLowLevel.html_to_pwndoc_html(text)


class MdAndHTMLConvertorLowLevel:

    HTML_TO_RESERVED_CHARS = {
        '<': '\ue020',
        '>': '\ue021',
        '&': '\ue022',
        '\'': '\ue023',
        '"': '\ue024',
    }

    _HTML_TAG_NORMALIZATION = {
        "<b>": "<strong>",
        "</b>": "</strong>",
        "<i>": "<em>",
        "</i>": "</em>",
    }

    _MD_TO_HTML_PAIR_TAGS = {
        "**": ("<strong>", "</strong>"),
        "__": ("<em>", "</em>"),
    }

    _HTML_TAGS_TO_STRIP_FOR_RAW_MD = {
        's': '',  # ~~strike-through~~
        'u': '',  # underlined
        'code': '',
        'p': '\n',
        'ul': '\n',
        'li': '\n-',  # todo: this works for opening tag, not so much for the closing tag
    }
    # missing: div and it's style and classes
    # missing: span and it's style and classes
    # note: <br> and (un)escaping is handled separately

    @staticmethod
    def _get_basic_html_tags():
        simple_tags = ['b', 'i', 'br', 'p', 'ul', 'li', 'strong', 's', 'u', 'p', 'code', 'strong', 'em', 'div', 'span']
        answer = []
        for tag in simple_tags:
            answer.append(f'<{tag}>')
            answer.append(f'</{tag}>')
        return answer

    # BEGIN: MD -> HTML

    @staticmethod
    def __is_md_separator(char: str) -> bool:
        # warning: This is not up to spec, but slightly more involved than scan2report version.
        # correct spec would be: https://spec.commonmark.org/0.30/#emphasis-and-strong-emphasis
        return char.isspace() or char.isnumeric() or char in '''!"#$%&'()+, -./:;<=>?@[]^`{|}~'''

    @classmethod
    def __is_valid_md_state_change(cls, current_state: bool, char_before: str, char_after: str):
        # warning: This is not up to spec.
        # if (not current_state and cls.__is_md_separator(char_before) and not cls.__is_md_separator(char_after)) or \
        #    (current_state and not cls.__is_md_separator(char_before) and cls.__is_md_separator(char_after)):
        return True

    @classmethod
    def __replace_simple_tag_with_pair_tag(cls, text: str, original_tag: str, opening_tag: str, closing_tag: str) -> str:
        answer = ""
        remaining_string = text
        current_state = False

        while True:
            pre, sep, post = remaining_string.partition(original_tag)
            answer += pre
            if len(sep) == 0:
                break

            char_before = (" " + pre)[-1]
            char_after = (post + " ")[0]

            if cls.__is_valid_md_state_change(current_state, char_before, char_after):
                current_state = not current_state
                answer += opening_tag if current_state else closing_tag
            else:
                answer += sep

            remaining_string = post

        if current_state:
            logger.warning(f"Conversion found possible unmatched tag {original_tag}")

        return answer

    @classmethod
    def replace_md_pair_tags(cls, text: str) -> str:
        for md_tag, html_tags in cls._MD_TO_HTML_PAIR_TAGS.items():
            text = cls.__replace_simple_tag_with_pair_tag(text, md_tag, html_tags[0], html_tags[1])
        return text

    @classmethod
    def reverse_custom_escaping_of_html(cls, text: str) -> str:
        for html_char, custom_escaped_char in cls.HTML_TO_RESERVED_CHARS.items():
            text = text.replace(custom_escaped_char, html_char)

            unicode_repr = ascii(custom_escaped_char).strip("'")
            text = text.replace(unicode_repr, html_char)  # Handles unicode repr (e.g. \ue020) which happens for example when outputting to JSON.
        return text

    # BEGIN: HTML -> MD

    @classmethod
    def _normalize_html_tags(cls, text: str) -> str:
        """Normalize (partially) the HTML tags (i.e. <b> -> <strong>)"""
        for original_tag, new_tag in cls._HTML_TAG_NORMALIZATION.items():
            text = text.replace(original_tag, new_tag)
        return text

    @classmethod
    def replace_html_pair_tags_with_md_equivalents(cls, text: str) -> str:
        """ Convert HTML formatting to MD (i.e. <strong> -> **) """
        for md_tag, html_tags in cls._MD_TO_HTML_PAIR_TAGS.items():
            for html_tag in html_tags:
                text = text.replace(html_tag, md_tag)
        return text

    @classmethod
    def strip_unsupported_html_tags(cls, text: str) -> str:
        """ Strip (some) unsupported HTML tags (i.e. <p> and </p> to <br>) """
        for html_tag, md_tag in cls._HTML_TAGS_TO_STRIP_FOR_RAW_MD.items():
            text = text.replace(f"<{html_tag}>", md_tag)
            text = text.replace(f"</{html_tag}>", md_tag)
        return text

    @classmethod
    def custom_escape_html(cls, text: str) -> str:
        """ Escapes HTML special chars so that it's possible to distinguish them from any HTML chars added later in the process by scan2report. """
        for html_char, custom_escaped_char in cls.HTML_TO_RESERVED_CHARS.items():
            text = text.replace(html_char, custom_escaped_char)
        return text

    @classmethod
    def md_to_html(cls, md_text: str) -> str:
        answer = cls.replace_md_pair_tags(md_text)
        answer = answer.replace("\n", "<br>")
        return answer

    @classmethod
    def force_unescape_basic_tags(cls, text: str) -> str:
        logger.debug("force_unescape_basic_tags is deprecated.")  # todo
        simple_tags = cls._get_basic_html_tags()

        for full_tag_html in simple_tags:
            full_tag_escaped = html.escape(full_tag_html)
            text = text.replace(full_tag_escaped, full_tag_html)
            # note: this only covers first level encoded, i.e. &lt;p&gt; but not &amp;lt;p&amp;gt;
            #  it also does not cover custom encoding, however neither of those should occur
        return text

    @staticmethod
    def extract_links(bs: BeautifulSoup) -> List[str]:
        links = []
        for a in bs.find_all('a'):
            content = a.string

            url = a.get('href')
            if url is None:
                a.replace_with(content)
            else:
                links.append(url)
                a.replace_with(f"{content} [{len(links)}]")

        return links

    @staticmethod
    def insert_links(text: str, links: List[str]) -> str:
        if len(links) > 0:
            text += "<p>\n"
            for i, link in enumerate(links):
                text += f"[{i+1}]: {link}<br>\n"
            text += "</p>"
        return text

    @staticmethod
    def remove_br_from_ul(bs: BeautifulSoup):
        # todo: Check performance of this for loop. Previously this and similar loops were HUGE performance slow down.
        for br in bs.find_all('ul br'):
            br.decompose()

    @staticmethod
    def convert_top_level_br_to_empty_p_tag(bs: BeautifulSoup):
        # I can't use wrap_bare_words, because I don't want <p><br></p> but just <p></p>
        children = bs.find_all("br", recursive=False)
        for br_child in children:
            br_child.replace_with(bs.new_tag("p"))

    @staticmethod
    def _remove_html_body_wrapper(text: str) -> str:
        """ lxml parser is MUCH faster, but adds HTML ballast around the input"""
        # todo: python 3.9 has remove prefix and suffix
        prefix = "<html><body>"
        if text.startswith(prefix):
            text = text[len(prefix):]
        suffix = "</body></html>"
        if text.endswith(suffix):
            text = text[:-len(suffix)]
        return text

    @classmethod
    @functools.lru_cache(maxsize=1024)
    def html_to_pwndoc_html(cls, text: str) -> str:
        bs = BeautifulSoup(text, BS_HTML_PARSER)  # Beware of differences in produced HTML trees for invalid documents; https://beautiful-soup-4.readthedocs.io/en/latest/index.html#differences-between-parsers
        wrap_bare_words(bs)

        links: List[str] = cls.extract_links(bs)  # PwnDoc doesn't support <a> tags yet.

        cls.remove_br_from_ul(bs)  # HTML doesn't allow <br> inside <ul>
        cls.convert_top_level_br_to_empty_p_tag(bs)  # PwnDoc doesn't allow <br> as a top level tag. Replace it with <p><br></p>
        text = pwndoc_extract_tags_from_forbidden_parents(bs)  # Some elements can be inside other elements (e.g. <img> inside <p>). Extract them to higher levels.

        bs = BeautifulSoup(text, BS_HTML_PARSER)
        wrap_bare_words(bs)

        text = PwnDocFormatterInstance.to_text(bs)

        text = cls.insert_links(text, links)  # Re-add links, that were extracted before

        text = text.replace("<br/>", "<br>")  # PwnDoc doesn't support <br/> version of <br> tag.
        text = text.replace("\n", "")  # Newlines can breakup parsing for beautifying, as well as PwnDoc docx generation.

        if BS_HTML_PARSER == 'lxml':
            text = cls._remove_html_body_wrapper(text)

        return text

# todo: add lru cache on str -> str steps?
