# Copyright (c) 2018-2022, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#
__all__ = ["DocumentationBuilderMd"]

import re
from enum import Enum, auto
from pathlib import Path
from typing import Any, Dict, List, NamedTuple, Tuple, Union

import omni.client as client

H1_CHECK = re.compile(r"# .*")
H2_CHECK = re.compile(r"## .*")
H3_CHECK = re.compile(r"### .*")
PARAGRAPH_CHECK = re.compile(r"[a-zA-Z0-9].*")
LIST_CHECK = re.compile(r"^[-+] [a-zA-Z0-9].*$")
CODE_CHECK = re.compile(r"```.*")
IMAGE_CHECK = re.compile('!\\[(?P<alt>[^\\]]*)\\]\\((?P<url>.*?)(?=\\"|\\))(?P<title>\\".*\\")?\\)')
LINK_CHECK = re.compile(r"\[(.*)\]\((.*)\)")
LINK_IN_LIST = re.compile(r"^[-+] \[(.*)\]\((.*)\)")
TABLE_CHECK = re.compile(r"^(?:\s*\|.+)+\|\s*$")


class _BlockType(Enum):
    H1 = auto()
    H2 = auto()
    H3 = auto()
    PARAGRAPH = auto()
    LIST = auto()
    CODE = auto()
    IMAGE = auto()
    LINK = auto()
    TABLE = auto()
    LINK_IN_LIST = auto()


class _Block(NamedTuple):
    block_type: _BlockType
    text: Union[str, Tuple[str]]
    source: str
    metadata: Dict[str, Any]


class DocumentationBuilderMd:
    """
    The cache for .md file.

    It contains parsed data, catalog, links, etc...
    """

    def __init__(self, filename: str):
        """
        ### Arguments

        `filename : str`
            Path to .md file
        """

        self.__filename = filename
        self.__blocks: List[_Block] = []

        result, _, content = client.read_file(f"{filename}")

        if result != client.Result.OK:
            content = ""
        else:
            content = memoryview(content).tobytes().decode("utf-8")

        self._process_lines(content.splitlines())

    def destroy(self):
        pass

    @property
    def blocks(self):
        return self.__blocks

    @property
    def catalog(self):
        """
        Return the document outline in the following format:
        {"name": "h1", "children": []}
        """
        result = []
        blocks = [b for b in self.__blocks if b.block_type in [_BlockType.H1, _BlockType.H2, _BlockType.H3]]
        if blocks:
            self._get_catalog_recursive(blocks, result, blocks[0].block_type.value)
        return result

    @property
    def source(self) -> str:
        content = []
        code_id = 0
        page_name = self.name

        for block in self.__blocks:
            if block.block_type == _BlockType.CODE:
                content.append(self._source_code(block, page_name, code_id))
                code_id += 1
            else:
                content.append(block.source)

        return "\n\n".join(content) + "\n"

    @property
    def code_blocks(self) -> List[Dict[str, str]]:
        page_name = self.name
        blocks = [b for b in self.__blocks if b.block_type == _BlockType.CODE]
        return [
            {"name": f"{page_name}_{i}", "code": block.text, "height": block.metadata["height"]}
            for i, block in enumerate(blocks)
            if block.metadata.get("code_type", None) == "execute"
        ]

    @property
    def image_blocks(self) -> List[Dict[str, str]]:
        return [
            {"path": block.text, "alt": block.metadata["alt"], "title": block.metadata["title"]}
            for block in self.__blocks
            if block.block_type == _BlockType.IMAGE
        ]

    @property
    def name(self):
        first_header = next(
            (b for b in self.__blocks if b.block_type in [_BlockType.H1, _BlockType.H2, _BlockType.H3]), None
        )
        return first_header.text if first_header else None

    def _get_catalog_recursive(self, blocks: List[_Block], children: List[Dict], level: int):
        while blocks:
            top = blocks[0]
            top_level = top.block_type.value
            if top_level == level:
                children.append({"name": top.text, "children": []})
                blocks.pop(0)
            elif top_level > level:
                self._get_catalog_recursive(blocks, children[-1]["children"], top_level)
            else:  # top_level < level
                return

    def _add_block(self, block_type, text, source, **kwargs):
        self.__blocks.append(_Block(block_type, text, source, kwargs))

    def _process_lines(self, content):

        # Remove empty lines
        while True:
            while content and not content[0].strip():
                content.pop(0)

            if not content:
                break

            if H1_CHECK.match(content[0]):
                self._process_h1(content)
            elif H2_CHECK.match(content[0]):
                self._process_h2(content)
            elif H3_CHECK.match(content[0]):
                self._process_h3(content)
            elif CODE_CHECK.match(content[0]):
                self._process_code(content)
            elif IMAGE_CHECK.match(content[0]):
                self._process_image(content)
            elif LINK_CHECK.match(content[0]):
                self._process_link(content)
            elif LINK_IN_LIST.match(content[0]):
                self._process_link_in_list(content)
            elif LIST_CHECK.match(content[0]):
                self._process_list(content)
            elif TABLE_CHECK.match(content[0]):
                self._process_table(content)
            else:
                self._process_paragraph(content)

    def _process_h1(self, content):
        text = content.pop(0)
        self._add_block(_BlockType.H1, text[2:], text)

    def _process_h2(self, content):
        text = content.pop(0)
        self._add_block(_BlockType.H2, text[3:], text)

    def _process_h3(self, content):
        text = content.pop(0)
        self._add_block(_BlockType.H3, text[4:], text)

    def _process_list(self, content: List[str]):
        text_lines = []
        while True:
            current_line = content.pop(0)
            text_lines.append(current_line.strip())

            if not content or not LIST_CHECK.match(content[0]):
                # The next line is not a paragraph.
                break
        self._add_block(_BlockType.LIST, tuple(text_lines), "\n".join(text_lines))

    def _process_table(self, content: List[str]):
        text_lines = []
        while True:
            current_line = content.pop(0)
            text_lines.append(current_line.strip())

            if not content or not TABLE_CHECK.match(content[0]):
                # The next line is not a paragraph.
                break
        self._add_block(_BlockType.TABLE, tuple(text_lines), "\n".join(text_lines))

    def _process_link(self, content):
        text = content.pop(0)
        self._add_block(_BlockType.LINK, text, text)

    def _process_link_in_list(self, content: List[str]):
        text_lines = []
        while True:
            current_line = content.pop(0)
            text_lines.append(current_line.strip())

            if not content or not LINK_CHECK.match(content[0]):
                # The next line is not a link in a list.
                break
        self._add_block(_BlockType.LINK_IN_LIST, tuple(text_lines), "\n".join(text_lines))

    def _process_paragraph(self, content: List[str]):
        text_lines = []
        while True:
            current_line = content.pop(0)
            text_lines.append(current_line.strip())

            if current_line.endswith("  "):
                # To create a line break, end a line with two or more spaces.
                break

            if not content or not PARAGRAPH_CHECK.match(content[0]):
                # The next line is not a paragraph.
                break

        self._add_block(_BlockType.PARAGRAPH, " ".join(text_lines), "\n".join(text_lines))

    def _process_code(self, content):
        source = []

        text = content.pop(0)
        source.append(text)

        code_type = text[3:].strip().split()

        # Extract height
        if code_type and len(code_type) > 1:
            height = code_type[1]
        else:
            height = None

        if code_type:
            code_type = code_type[0]
        else:
            code_type = ""

        code_lines = []

        while content and not CODE_CHECK.match(content[0]):
            text = content.pop(0)
            source.append(text)
            code_lines.append(text)

        source.append(content.pop(0))

        self._add_block(_BlockType.CODE, "\n".join(code_lines), "\n".join(source), code_type=code_type, height=height)

    def _process_image(self, content: List[str]):
        text = content.pop(0)
        parts = IMAGE_CHECK.match(text.strip()).groupdict()
        path = Path(self.__filename).parent.joinpath(parts["url"])
        self._add_block(_BlockType.IMAGE, f"{path}", text, alt=parts["alt"], title=parts["title"])

    @staticmethod
    def get_clean_code(block):
        # Remove double-comments
        clean = []
        skip = False
        for line in block.text.splitlines():
            if line.lstrip().startswith("##"):
                skip = not skip
                continue

            if skip:
                continue

            if not clean and not line.strip():
                continue

            clean.append(line)

        return "\n".join(clean)

    def _source_code(self, block, page_name, code_id):
        result = []
        code_type = block.metadata.get("code_type", "")
        if code_type == "execute":
            image_name = f"{page_name}_{code_id}.png"
            result.append(f"![Code Result]({image_name})\n")
            code_type = "python"
        elif code_type == "execute-manual":
            code_type = "python"
        result.append(f"```{code_type}")
        result.append(self.get_clean_code(block))
        result.append("```")
        return "\n".join(result)
