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

import shutil
import traceback
from pathlib import Path

import carb
import omni.kit.app
import omni.ui as ui

from .documentation_capture import DocumentationBuilderCapture
from .documentation_md import DocumentationBuilderMd
from .documentation_style import get_style

PRINT_GENERATION_STATUS = True  # Can be made into setting if needed


def _print(message):
    if PRINT_GENERATION_STATUS:
        print(message)
    else:
        carb.log_info(message)


def generate_for_sphinx(extension_id: str, path: str) -> bool:
    """
    Produce the set of files and screenshots readable by Sphinx.

    ### Arguments

        `extension_id: str`
            The extension to generate documentation files
        `path: str`
            The path to produce the files.

    ### Return

        `True` if generation succeded.
    """
    _print(f"Generating for sphinx. Ext: {extension_id}. Path: {path}")

    out_path = Path(path)
    if not out_path.is_dir():
        carb.log_error(f"{path} is not a directory")
        return False

    ext_info = omni.kit.app.get_app().get_extension_manager().get_extension_dict(extension_id)
    if ext_info is None:
        carb.log_error(f"Extension id {extension_id} not found")
        return False

    if not ext_info.get("state", {}).get("enabled", False):
        carb.log_error(f"Extension id {extension_id} is not enabled")
        return False

    ext_path = Path(ext_info.get("path", ""))
    pages = ext_info.get("documentation", {}).get("pages", [])
    if len(pages) == 0:
        carb.log_error(f"Extension id {extension_id} has no pages set in [documentation] section of config.")
        return False

    for page in pages:
        _print(f"Processing page: '{page}'")
        if not page.endswith(".md"):
            _print(f"Not .md, skipping: '{page}'")
            continue

        page_path = Path(ext_path / page)
        md = DocumentationBuilderMd(str(page_path))
        if md.name is None:
            _print(f"md is None, skipping: '{page}'")
            continue

        for code_block in md.code_blocks:
            image_file = out_path / f"{code_block['name']}.png"
            height = int(code_block["height"] or 200)
            with DocumentationBuilderCapture(str(image_file), 800, height):
                stack = ui.ZStack()
                stack.set_style(get_style())
                with stack:
                    # Background
                    ui.Rectangle(name="code_background")

                    # The code
                    with ui.VStack():
                        try:
                            _execute_code(code_block["code"])
                        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
                            )
            _print(f"Generated image: '{image_file}'")

        for image_block in md.image_blocks:
            image_file = Path(image_block["path"])
            image_out = out_path / image_file.relative_to(page_path.parent)
            image_out.parent.mkdir(parents=True, exist_ok=True)
            try:
                shutil.copyfile(image_file, image_out)
            except shutil.SameFileError:
                pass
            _print(f"Copied image: '{image_out}'")

        md_file = out_path / f"{page_path.stem}.md"
        with open(md_file, "w", encoding="utf-8") as f:
            f.write(md.source)
        _print(f"Generated md file: '{md_file}'")
    return True


def _execute_code(code: str):
    # Wrap the code in a class to provide scoping
    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)
    locals_from_execution["Execution"]()
