import functools
from typing import Callable

from bs4 import BeautifulSoup
from loguru import logger
import html


class MdAndHTMLConvertor:

    _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

    # BEGIN: Public interface

    @classmethod
    def md_to_html(cls, md_text: str) -> str:
        answer = md_text
        answer = html.escape(answer)
        answer = cls._replace_md_pair_tags(answer)
        answer = cls.reverse_custom_escaping_of_html(answer)
        answer = answer.replace("\n", "<br>")

        return answer

    @classmethod
    def html_to_md(cls, html_text: str, output_without_html_tags: bool = False) -> str:
        answer = html_text
        answer = cls._normalize_html_tags(answer)
        answer = cls.html_to_pwndoc_html(answer)

        if output_without_html_tags:
            answer = answer.replace("<br>", "\n")
            answer = cls._replace_html_pair_tags_with_md_equivalents(answer)
            answer = cls._strip_unsupported_html_tags(answer)
            answer = html.unescape(answer)
        else:
            answer = cls._custom_escape_html(answer)

        return answer

    @classmethod
    def force_unescape_basic_tags(cls, text: str) -> str:
        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 html_to_pwndoc_html(text: str) -> str:
        text = html.unescape(text)

        # PwnDoc is extremely strict in regards to HTML standard.
        text = text.replace("<p>", "").replace("</p>", "<br>")  # HTML doesn't allow <p> in <p>
        text = text.replace('<ul>', '</p><ul>').replace('</ul>', '</ul><p>')  # HTML doesn't allow <ul> inside <p>

        text = f'<p>{text}</p>'  # PwnDoc requires all text to be in some of the allowed tags.

        # Beware of differences in produced HTML trees for invalid documents; https://beautiful-soup-4.readthedocs.io/en/latest/index.html#differences-between-parsers
        bs = BeautifulSoup(text, 'html.parser')
        # bs = BeautifulSoup(text, 'lxml', from_encoding="utf8")  # This parser for general public is faster, but it's not the thing that's actually costing time.

        # PwnDoc doesn't support <a> tags yet.
        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)}]")

        # HTML doesn't allow <br> inside <ul>
        for ul in bs.find_all('ul'):
            for br in ul.find_all('br'):
                br.decompose()  # todo: optimize this - together with the call bellow it takes > 60% of the time

        # this doesn't speed it up
        #for br in bs.find_all('ul br'):
        #    br.decompose()

        # PwnDoc doesn't allow <br> as a top level tag. Replace it with <p><br></p>
        children = bs.find_all("br", recursive=False)
        for br_child in children:
            sup = bs.new_tag('p').contents = bs.new_tag("br")
            br_child.insert_after(sup)  # todo: optimize this - together with the call above it takes > 60% of the time
            br_child.decompose()
            # br_child.replace_with(sup)  # note: this doesn't speed it up

        # Prettify produces HTML, which is valid enough for browsers, but not strictly valid HTML as required by PwnDoc.
        # text = bs.prettify(formatter="minimal")  # minimal does escape <, >, &
        text = bs.prettify(formatter=None)

        # Re-add links, that were extracted before
        if len(links) > 0:
            text += "<p>\n"
            for i, link in enumerate(links):
                text += f"[{i+1}]: {link}<br>\n"
            text += "</p>"

        # Sanitize from htmllaundry does seem to produce clean HTML, but it removes some desired tags (for example <li>).
        # text = htmllaundry.sanitize(text)

        # PwnDoc doesn't support <br/> version of <br> tag.
        text = text.replace("<br/>", "<br>")

        # 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!
        text = html.escape(text)
        text = MdAndHTMLConvertor.force_unescape_basic_tags(text)  # this only unescapes this formats: <tag> </tag>, so no <a href ...>!

        text = text.replace("\n", "")  # Newlines can breakup parsing for beautifying, as well as PwnDoc docx generation.
        return text


# ----- TEST

def _test_helper(function: Callable[[str], str], text_in: str, expected_output: str):
    function_output = function(text_in)
    assert function_output == expected_output, f"Test: {text_in} | Output: {function_output} | Expected output: {expected_output}"


if __name__ == "__main__":
    md_to_html_tests = {
        "**test**": "<strong>test</strong>",
        "*test*": "*test*",
        "&": "&amp;",

        # The following tests are according to spec, but won't work (that same way they don't work in scan2report).
        # "***test***": "***test***",
        # "****": "****",
        # "**test": "**test"
        # "**test**test": "**test**test"
    }

    html_preserving_tests = [
        "<br><strong><FILL_IN_LOREM_IPSUM></strong>"
    ]

    # '\ue022' = MdAndHTMLConvertor._HTML_TO_RESERVED_CHARS["&"]
    _test_helper(MdAndHTMLConvertor.md_to_html, '\ue022', '&')

    to_md_with_custom_escaped_html = functools.partial(MdAndHTMLConvertor.html_to_md, output_without_html_tags=False)
    to_md_without_html = functools.partial(MdAndHTMLConvertor.html_to_md, output_without_html_tags=True)

    _test_helper(to_md_with_custom_escaped_html, '&', '\ue022')
    _test_helper(to_md_without_html, '&', '&')

    # Md to HTML
    for test, expected_result in md_to_html_tests.items():
        _test_helper(MdAndHTMLConvertor.md_to_html, test, expected_result)

    # Html to Md
    for expected_result, test in md_to_html_tests.items():
        _test_helper(to_md_without_html, test, expected_result)

    # Md to HTML to Md
    for test, expected_result in md_to_html_tests.items():
        html_text = MdAndHTMLConvertor.md_to_html(test)
        md_back = to_md_without_html(html_text)
        assert test == md_back, f"Test (MD): {test} | MD -> HTML: {html_text} | MD -> HTML -> MD : {md_back}"

    # HTML to Md to HTML
    for test in html_preserving_tests:
        md_text = to_md_with_custom_escaped_html(test)
        html_back = MdAndHTMLConvertor.md_to_html(md_text)
        assert test == html_back, f"Test (html): {test} | HTML -> MD: {md_text} | HTML -> MD -> HTML : {html_back}"
