# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: F821, RUF012
"""Highlight printed TVM script."""

import functools
import os
import shutil
import subprocess
import sys
import warnings
from typing import Any, Union


def cprint(
    printable: Any | str,
    style: str | None = None,
    black_format: bool = False,
) -> None:
    """Print TVMScript string with Pygments highlight and Black auto-formatting.

    Parameters
    ----------
    printable : Union[IRModule, PrimFunc, str]

        The TVMScript to be printed

    style : str, optional

        Pygmentize printing style, auto-detected if None.

    black_format: bool

        If true, use the formatter Black to format the TVMScript

    Notes
    -----

    The style parameter follows the Pygments style names or Style objects. Three
    built-in styles are extended: "light", "dark" and "ansi". By default, "light"
    will be used for notebook environment and terminal style will be "ansi" for
    better style consistency. As an fallback when the optional Pygment library is
    not installed, plain text will be printed with a one-time warning to suggest
    installing the Pygment library. Other Pygment styles can be found in
    https://pygments.org/styles/

    The default pygmentize style can also be set with the environment
    variable "TVM_PYGMENTIZE_STYLE".
    """
    if hasattr(printable, "script") and callable(getattr(printable, "script")):
        printable = printable.script()
    elif not isinstance(printable, str):
        raise TypeError(
            f"Only can print strings or objects with `script` method, but got: {type(printable)}"
        )

    if black_format:
        printable = _format(printable)

    is_in_notebook = "ipykernel" in sys.modules  # in notebook env (support html display).

    style = _get_pygments_style(style, is_in_notebook)

    if style is None:
        print(printable)
        return

    # pylint: disable=import-outside-toplevel
    from pygments import highlight
    from pygments.formatters import HtmlFormatter, Terminal256Formatter
    from pygments.lexers.python import Python3Lexer

    if is_in_notebook:
        from IPython import display  # pylint: disable=import-outside-toplevel

        formatter = HtmlFormatter(style=style)
        formatter.noclasses = True  # inline styles
        html = highlight(printable, Python3Lexer(), formatter)
        display.display(display.HTML(html))
    else:
        print(highlight(printable, Python3Lexer(), Terminal256Formatter(style=style)))


@functools.lru_cache
def _get_formatter(formatter: str | None = None):
    def get_ruff_formatter():
        if shutil.which("ruff") is None:
            return None

        def formatter(code_str):
            proc = subprocess.Popen(
                ["ruff", "format", "--stdin-filename=TVMScript"],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                encoding="utf-8",
            )
            stdout, _stderr = proc.communicate(code_str)
            return stdout

        return formatter

    def get_black_formatter():
        try:
            # pylint: disable=import-outside-toplevel
            import black
        except ImportError:
            return None

        def formatter(code_str):
            return black.format_str(code_str, mode=black.FileMode())

        return formatter

    def get_fallback_formatter():
        def formatter(code_str):
            with warnings.catch_warnings():
                warnings.simplefilter("once", UserWarning)
                ruff_install_cmd = sys.executable + " -m pip install ruff"
                black_install_cmd = (
                    sys.executable + ' -m pip install "black==22.3.0" --upgrade --user'
                )
                warnings.warn(
                    f"Neither the 'ruff' formatter nor the 'black' formatter is available.  "
                    f"To print formatted TVM script, please a formatter.  \n"
                    f"To install ruff: {ruff_install_cmd}\n"
                    f"To install black: {black_install_cmd}",
                    category=UserWarning,
                )
            return code_str

        return formatter

    # formatter = "black"
    if formatter is None:
        options = [get_ruff_formatter, get_black_formatter]
    elif formatter == "ruff":
        options = [get_ruff_formatter]
    elif formatter == "black":
        options = [get_black_formatter]
    else:
        raise ValueError(f"Unknown formatter: {formatter}")

    for option in options:
        func = option()
        if func is not None:
            return func
    return get_fallback_formatter()


def _format(code_str: str, formatter: str | None = None) -> str:
    """Format a code string using Black.

    Parameters
    ----------
    code_str: str

        The string containing Python/TVMScript code to format

    formatter: Optional[str]

        The formatter to use.  Can specify `ruff`, `black`, or
        auto-select by passing `None`.

    Returns
    -------
    formatted: str

        The formatted Python/TVMScript code

    """
    return _get_formatter(formatter)(code_str)


def _get_pygments_style(
    style: str | None, is_in_notebook: bool
) -> Union["pygments.style.Style", str] | None:
    """Select a pygments style to use

    Parameters
    ----------
    style: str

        The style specifier to use.  If None, auto-select a style.

    is_in_notebook: bool

        Whether python is currently running in a jupyter notebook.
        Used for automatic selection.

    Returns
    -------
    style: Optional[Union['pygments.style.Style',str]]

        If pygments is installed, the style object or string, suitable
        for use as the "style" argument to pygments formatters.  If
        pygments is not installed, returns None.

    """
    try:
        # pylint: disable=import-outside-toplevel
        import pygments
        from packaging import version
        from pygments.style import Style
        from pygments.token import Comment, Keyword, Name, Number, Operator, String

        if version.parse(pygments.__version__) < version.parse("2.4.0"):
            raise ImportError("Required Pygments version >= 2.4.0 but got " + pygments.__version__)
    except ImportError as err:
        if err.name == "packaging":
            name = "packaging"
        elif err.name == "pygments":
            name = "Pygments>=2.4.0"
        else:
            raise ValueError(f'Package "{err.name}" should not be used')

        with warnings.catch_warnings():
            warnings.simplefilter("once", UserWarning)
            install_cmd = sys.executable + f' -m pip install "{name}" --upgrade --user'
            warnings.warn(
                str(err)
                + "\n"
                + f"To print highlighted TVM script, please install {name}:\n"
                + install_cmd,
                category=UserWarning,
            )
        return None

    class JupyterLight(Style):
        """A Jupyter-Notebook-like Pygments style configuration (aka. "light")"""

        background_color = ""
        styles = {
            Keyword: "bold #008000",
            Keyword.Type: "nobold #008000",
            Name.Function: "#0000FF",
            Name.Class: "bold #0000FF",
            Name.Decorator: "#AA22FF",
            String: "#BA2121",
            Number: "#008000",
            Operator: "bold #AA22FF",
            Operator.Word: "bold #008000",
            Comment: "italic #007979",
        }

    class VSCDark(Style):
        """A VSCode-Dark-like Pygments style configuration (aka. "dark")"""

        background_color = ""
        styles = {
            Keyword: "bold #c586c0",
            Keyword.Type: "#82aaff",
            Keyword.Namespace: "#4ec9b0",
            Name.Class: "bold #569cd6",
            Name.Function: "bold #dcdcaa",
            Name.Decorator: "italic #fe4ef3",
            String: "#ce9178",
            Number: "#b5cea8",
            Operator: "#bbbbbb",
            Operator.Word: "#569cd6",
            Comment: "italic #6a9956",
        }

    class AnsiTerminalDefault(Style):
        """The default style for terminal display with ANSI colors (aka. "ansi")"""

        background_color = ""
        styles = {
            Keyword: "bold ansigreen",
            Keyword.Type: "nobold ansigreen",
            Name.Class: "bold ansiblue",
            Name.Function: "bold ansiblue",
            Name.Decorator: "italic ansibrightmagenta",
            String: "ansiyellow",
            Number: "ansibrightgreen",
            Operator: "bold ansimagenta",
            Operator.Word: "bold ansigreen",
            Comment: "italic ansibrightblack",
        }

    if style == "light":
        return JupyterLight
    elif style == "dark":
        return VSCDark
    elif style == "ansi":
        return AnsiTerminalDefault

    if style is not None:
        return style

    style_from_environment = os.environ.get("TVM_PYGMENTIZE_STYLE", "").strip()
    if style_from_environment:
        return style_from_environment

    if is_in_notebook:
        return JupyterLight

    return AnsiTerminalDefault