# 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__ = ["DocumentationBuilderPage"]

import re
import subprocess
import sys
import traceback
import weakref
import webbrowser
from functools import partial
from typing import Any, Dict, List, Optional

import omni.ui as ui
from PIL import Image

from .documentation_md import LINK_CHECK, DocumentationBuilderMd, _Block, _BlockType


class DocumentationBuilderPage:
    """
    Widget that draws .md page with the cache.
    """

    def __init__(self, md: DocumentationBuilderMd, scroll_to: Optional[str] = None):
        """
        ### Arguments

        `md : DocumentationBuilderMd`
            MD cache object

        `scroll_to : Optional[str]`
            The name of the section to scroll to
        """
        self.__md = md
        self.__scroll_to = scroll_to
        self.__images: List[ui.Image] = []
        self.__code_stacks: List[ui.ZStack] = []
        self.__code_locals: Dict[str, Any] = {}
        self._copy_context_menu = None
        self.__bookmarks: Dict[str, ui.Widget] = {}
        self._dont_scale_images = False

        ui.Frame(build_fn=self.__build)

    def destroy(self):
        self.__bookmarks: Dict[str, ui.Widget] = {}

        if self._copy_context_menu:
            self._copy_context_menu.destroy()
            self._copy_context_menu = None

        self.__code_locals = {}

        for stack in self.__code_stacks:
            stack.destroy()
        self.__code_stacks = []

        for image in self.__images:
            image.destroy()
        self.__images = []

    @property
    def name(self):
        """The name of the page"""
        return self.__md.name

    def scroll_to(self, where: str):
        """Scroll to the given section"""
        widget = self.__bookmarks.get(where, None)
        if widget:
            widget.scroll_here()

    def _execute_code(self, code: str):
        indented = "\n".join(["    " + line for line in code.splitlines()])
        source = f"class Execution:\n  def __init__(self):\n{indented}"
        locals_from_execution = {}
        # flake8: noqa: PLW0122
        exec(source, None, locals_from_execution)
        self.__code_locals[code] = locals_from_execution["Execution"]()

    def __build(self):
        self.__bookmarks: Dict[str, ui.Widget] = {}

        with ui.HStack():
            ui.Spacer(width=50)
            with ui.VStack(height=0):
                ui.Spacer(height=50)
                for block in self.__md.blocks:
                    # Build sections one by one

                    # Space between sections
                    space = ui.Spacer(height=10)
                    self.__bookmarks[block.text] = space
                    if block.text == self.__scroll_to:
                        # Scroll to this section if necessary
                        space.scroll_here()

                    if block.block_type == _BlockType.H1:
                        self._build_h1(block)
                    elif block.block_type == _BlockType.H2:
                        self._build_h2(block)
                    elif block.block_type == _BlockType.H3:
                        self._build_h3(block)
                    elif block.block_type == _BlockType.PARAGRAPH:
                        self._build_paragraph(block)
                    elif block.block_type == _BlockType.LIST:
                        self._build_list(block)
                    elif block.block_type == _BlockType.TABLE:
                        self._build_table(block)
                    elif block.block_type == _BlockType.LINK:
                        self._build_link(block)
                    elif block.block_type == _BlockType.CODE:
                        self._build_code(block)
                    elif block.block_type == _BlockType.IMAGE:
                        self._build_image(block)
                    elif block.block_type == _BlockType.LINK_IN_LIST:
                        self._build_link_in_list(block)

                ui.Spacer(height=50)
            ui.Spacer(width=50)

    def _build_h1(self, block: _Block):
        with ui.ZStack():
            ui.Rectangle(name="h1")
            ui.Label(block.text, name="h1", style_type_name_override="Text")

    def _build_h2(self, block: _Block):
        with ui.VStack():
            ui.Line(name="h2")
            ui.Label(block.text, name="h2", style_type_name_override="Text")
            ui.Line(name="h2")

    def _build_h3(self, block: _Block):
        with ui.VStack(height=0):
            ui.Line(name="h3")
            ui.Label(block.text, name="h3", style_type_name_override="Text")
            ui.Line(name="h3")

    def _build_paragraph(self, block: _Block):
        ui.Label(block.text, name="paragraph", word_wrap=True, skip_draw_when_clipped=True)

    def _build_list(self, block: _Block):
        with ui.VStack(spacing=-5):
            for list_entry in block.text:
                list_entry = list_entry.lstrip("+")
                list_entry = list_entry.lstrip("-")

                with ui.HStack(spacing=0):
                    ui.Circle(
                        width=10, size_policy=ui.CircleSizePolicy.FIXED, alignment=ui.Alignment.LEFT_CENTER, radius=3
                    )
                    ui.Label(
                        list_entry,
                        alignment=ui.Alignment.LEFT_TOP,
                        name="paragraph",
                        word_wrap=True,
                        skip_draw_when_clipped=True,
                    )
                    # ui.Spacer(width=5)

    def _build_link_in_list(self, block: _Block):

        with ui.VStack(spacing=-5):
            for list_entry in block.text:
                list_entry = list_entry.lstrip("+")
                list_entry = list_entry.lstrip("-")
                list_entry = list_entry.lstrip()

                with ui.HStack(spacing=0):
                    ui.Circle(
                        width=10, size_policy=ui.CircleSizePolicy.FIXED, alignment=ui.Alignment.LEFT_CENTER, radius=3
                    )
                    reg = re.match(LINK_CHECK, list_entry)
                    if reg and len(reg.groups()) == 2:
                        self._build_link_impl(reg.group(1), reg.group(2))

    def _build_link(self, block: _Block):
        text = block.text
        reg = re.match(LINK_CHECK, text)
        label = ""
        url = ""
        if reg and len(reg.groups()) == 2:
            label = reg.group(1)
            url = reg.group(2)
            self._build_link_impl(label, url)

    def _build_link_impl(self, label, url):
        def _on_clicked(a, b, c, d, url):
            try:
                webbrowser.open(url)
            except OSError:
                print("Please open a browser on: " + url)

        with ui.HStack():
            spacer = ui.Spacer(width=5)
            label = ui.Label(
                label,
                style_type_name_override="url_link",
                name="paragraph",
                word_wrap=True,
                skip_draw_when_clipped=True,
                tooltip=url,
            )

            label.set_mouse_pressed_fn(lambda a, b, c, d, u=url: _on_clicked(a, b, c, d, u))

    def _build_table(self, block: _Block):
        cells = []
        num_columns = 0
        for line in block.text:
            cell_text_list = line.split("|")
            cell_text_list = [c.strip() for c in cell_text_list]
            cell_text_list = [c for c in cell_text_list if c]
            num_columns = max(num_columns, len(cell_text_list))
            cells.append(cell_text_list)

        # Assume 2nd row is separators "|------|---|"
        del cells[1]

        frame_pixel_width = 800
        grid_column_width = frame_pixel_width / num_columns
        frame_pixel_height = len(cells) * 20

        with ui.ScrollingFrame(
            width=frame_pixel_width,
            height=frame_pixel_height,
            horizontal_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_OFF,
            vertical_scrollbar_policy=ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_OFF,
        ):
            with ui.VGrid(column_width=grid_column_width, row_height=20):
                for i in range(len(cells)):
                    for j in range(num_columns):
                        with ui.ZStack():
                            bg_color = 0xFFCCCCCC if i == 0 else 0xFFFFFFF
                            ui.Rectangle(
                                style={
                                    "border_color": 0xFF000000,
                                    "background_color": bg_color,
                                    "border_width": 1,
                                    "margin": 0,
                                }
                            )
                            cell_val = cells[i][j]
                            if len(cell_val) > 66:
                                cell_val = cell_val[:65] + "..."
                            ui.Label(f"{cell_val}", style_type_name_override="table_label")

    def _build_code(self, block: _Block):
        def label_size_changed(label: ui.Label, short_frame: ui.Frame, overlay: ui.Frame):
            max_height = ui.Pixel(200)
            if label.computed_height > max_height.value:
                short_frame.height = max_height
                short_frame.vertical_clipping = True

                with overlay:
                    with ui.VStack():
                        ui.Spacer()
                        with ui.ZStack():
                            ui.Rectangle(style_type_name_override="TextOverlay")
                            ui.InvisibleButton(clicked_fn=partial(remove_overlay, label, short_frame, overlay))

        def remove_overlay(label: ui.Label, short_frame: ui.Frame, overlay: ui.Frame):
            short_frame.height = ui.Pixel(label.computed_height)
            with overlay:
                ui.Spacer()

        code_type = block.metadata.get("code_type", None)
        is_execute_manual = code_type == "execute-manual"
        is_execute = code_type == "execute"
        if is_execute:
            with ui.Frame(horizontal_clipping=True):
                with ui.ZStack():
                    # Background
                    ui.Rectangle(name="code_background")

                    # The code
                    with ui.VStack():
                        try:
                            self._execute_code(block.text)
                        except Exception as e:
                            tb = traceback.format_exception(Exception, e, e.__traceback__)
                            ui.Label(
                                "".join(tb), name="code", style_type_name_override="Text", skip_draw_when_clipped=True
                            )

            ui.Spacer(height=10)

        # Remove double-comments
        clean = DocumentationBuilderMd.get_clean_code(block)

        with ui.Frame(horizontal_clipping=True):
            stack = ui.ZStack(height=0)
            with stack:
                ui.Rectangle(name="code")
                short_frame = ui.Frame()
                with short_frame:
                    label = ui.Label(clean, name="code", style_type_name_override="Text", skip_draw_when_clipped=True)

                overlay = ui.Frame()
                # Enable clipping/extending the label after we know its size
                label.set_computed_content_size_changed_fn(partial(label_size_changed, label, short_frame, overlay))

            self.__code_stacks.append(stack)

        if is_execute_manual:
            ui.Spacer(height=4)
            with ui.HStack():
                # ui.Spacer(height=10)
                def on_click(text=block.text):
                    try:
                        self._execute_code(block.text)
                    except Exception as e:
                        tb = traceback.format_exception(Exception, e, e.__traceback__)
                        ui.Label("".join(tb), name="code", style_type_name_override="Text", skip_draw_when_clipped=True)

                button = ui.Button(
                    "RUN CODE", style_type_name_override="RunButton", width=50, height=20, clicked_fn=on_click
                )
                ui.Spacer()

        if sys.platform != "linux":
            stack.set_mouse_pressed_fn(lambda x, y, b, m, text=clean: b == 1 and self._show_copy_menu(x, y, b, m, text))

    def _build_image(self, block: _Block):
        source_url = block.text

        def change_resolution_legacy(image_weak, size):
            """
            Change the widget size proportionally to the image
            We set a max height of 400 for anything.
            """
            MAX_HEIGHT = 400

            image = image_weak()

            if not image:
                return

            width, height = size

            image_height = min(image.computed_content_width * height / width, MAX_HEIGHT)
            image.height = ui.Pixel(image_height)

        def change_resolution(image_weak, size):
            """
            This leaves the image scale untouched
            """
            image = image_weak()

            if not image:
                return

            width, height = size
            image.height = ui.Pixel(height)
            image.width = ui.Pixel(width)

        im = Image.open(source_url)

        image = ui.Image(source_url, height=200)

        # Even though the legacy behaviour is probably wrong there are extensions out there with images that are
        # too big, we don't want them to display differently than before
        if self._dont_scale_images:
            image.set_computed_content_size_changed_fn(partial(change_resolution, weakref.ref(image), im.size))
        else:
            image.set_computed_content_size_changed_fn(partial(change_resolution_legacy, weakref.ref(image), im.size))

        self.__images.append(image)

    def _show_copy_menu(self, x, y, button, modifier, text):
        """The context menu to copy the text"""
        # Display context menu only if the right button is pressed
        def copy_text(text_value):
            # we need to find a cross plartform way currently Window Only
            subprocess.run(["clip.exe"], input=text_value.strip().encode("utf-8"), check=True)

        # Reset the previous context popup
        self._copy_context_menu = ui.Menu("Copy Context Menu")
        with self._copy_context_menu:
            ui.MenuItem("Copy Code", triggered_fn=lambda t=text: copy_text(t))

        # Show it
        self._copy_context_menu.show()
