## Copyright 2022 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Helpers that build CMake rules.

Each function takes a list of parameters and returns a string ready to be
included in a CMakeLists.txt file. Builder functions handle optional arguments,
lists, formatting, etc.

For example:

build_iree_fetch_artifact(
    target_name="abcd",
    source_url="https://example.com/abcd.tflite",
    output="./abcd.tflite",
    unpack=True)

Outputs:

iree_fetch_artifact(
  NAME
    "abcd"
  SOURCE_URL
    "https://example.com/abcd.tflite"
  OUTPUT
    "./abcd.tflite"
  UNPACK
)
"""

from typing import Dict, List, Optional, Sequence

INDENT_SPACES = " " * 2


def _get_string_list(values: Sequence[str], quote: bool = True) -> List[str]:
    if quote:
        return [f'"{value}"' for value in values]
    return list(values)


def _get_block_body(body: List[str]) -> List[str]:
    return [INDENT_SPACES + line for line in body]


def _get_string_arg_block(
    keyword: str, value: Optional[str], quote: bool = True
) -> List[str]:
    if value is None:
        return []
    if quote:
        value = f'"{value}"'
    return [f"{keyword} {value}"]


def _get_string_list_arg_block(
    keyword: str, values: Sequence[str], quote: bool = True
) -> List[str]:
    if len(values) == 0:
        return []
    body = _get_string_list(values, quote)
    return [keyword] + _get_block_body(body)


def _get_option_arg_block(keyword: str, value: Optional[bool]) -> List[str]:
    if value is True:
        return [keyword]
    return []


def _build_call_rule(
    rule_name: str, parameter_blocks: Sequence[List[str]]
) -> List[str]:
    output = [f"{rule_name}("]
    for block in parameter_blocks:
        if len(block) == 0:
            continue
        output.extend(_get_block_body(block))
    output.append(")")
    return output


def _convert_block_to_string(block: List[str]) -> str:
    # Hack to append the terminating newline and only copies the list instead of
    # the whole string.
    return "\n".join(block + [""])


def build_iree_bytecode_module(
    target_name: str,
    src: str,
    module_name: str,
    flags: List[str] = [],
    compile_tool_target: Optional[str] = None,
    c_identifier: Optional[str] = None,
    static_lib_path: Optional[str] = None,
    deps: List[str] = [],
    friendly_name: Optional[str] = None,
    testonly: bool = False,
    public: bool = True,
) -> str:
    name_block = _get_string_arg_block("NAME", target_name)
    src_block = _get_string_arg_block("SRC", src)
    module_name_block = _get_string_arg_block("MODULE_FILE_NAME", module_name)
    c_identifier_block = _get_string_arg_block("C_IDENTIFIER", c_identifier)
    static_lib_block = _get_string_arg_block("STATIC_LIB_PATH", static_lib_path)
    compile_tool_target_block = _get_string_arg_block(
        "COMPILE_TOOL", compile_tool_target
    )
    flags_block = _get_string_list_arg_block("FLAGS", flags)
    deps_block = _get_string_list_arg_block("DEPS", deps)
    friendly_name_block = _get_string_arg_block("FRIENDLY_NAME", friendly_name)
    testonly_block = _get_option_arg_block("TESTONLY", testonly)
    public_block = _get_option_arg_block("PUBLIC", public)
    return _convert_block_to_string(
        _build_call_rule(
            rule_name="iree_bytecode_module",
            parameter_blocks=[
                name_block,
                src_block,
                module_name_block,
                c_identifier_block,
                compile_tool_target_block,
                static_lib_block,
                flags_block,
                friendly_name_block,
                deps_block,
                testonly_block,
                public_block,
            ],
        )
    )


def build_iree_fetch_artifact(
    target_name: str, source_url: str, output: str, unpack: bool
) -> str:
    name_block = _get_string_arg_block("NAME", target_name)
    source_url_block = _get_string_arg_block("SOURCE_URL", source_url)
    output_block = _get_string_arg_block("OUTPUT", output)
    unpack_block = _get_option_arg_block("UNPACK", unpack)
    return _convert_block_to_string(
        _build_call_rule(
            rule_name="iree_fetch_artifact",
            parameter_blocks=[name_block, source_url_block, output_block, unpack_block],
        )
    )


def build_iree_import_tf_model(
    target_path: str, source: str, import_flags: List[str], output_mlir_file: str
) -> str:
    target_name_block = _get_string_arg_block("TARGET_NAME", target_path)
    source_block = _get_string_arg_block("SOURCE", source)
    import_flags_block = _get_string_list_arg_block("IMPORT_FLAGS", import_flags)
    output_mlir_file_block = _get_string_arg_block("OUTPUT_MLIR_FILE", output_mlir_file)
    return _convert_block_to_string(
        _build_call_rule(
            rule_name="iree_import_tf_model",
            parameter_blocks=[
                target_name_block,
                source_block,
                import_flags_block,
                output_mlir_file_block,
            ],
        )
    )


def build_iree_import_tflite_model(
    target_path: str, source: str, import_flags: List[str], output_mlir_file: str
) -> str:
    target_name_block = _get_string_arg_block("TARGET_NAME", target_path)
    source_block = _get_string_arg_block("SOURCE", source)
    import_flags_block = _get_string_list_arg_block("IMPORT_FLAGS", import_flags)
    output_mlir_file_block = _get_string_arg_block("OUTPUT_MLIR_FILE", output_mlir_file)
    return _convert_block_to_string(
        _build_call_rule(
            rule_name="iree_import_tflite_model",
            parameter_blocks=[
                target_name_block,
                source_block,
                import_flags_block,
                output_mlir_file_block,
            ],
        )
    )


def build_iree_benchmark_suite_module_test(
    target_name: str,
    driver: str,
    expected_output: str,
    platform_module_map: Dict[str, str],
    runner_args: Sequence[str],
    timeout_secs: Optional[int] = None,
    labels: Sequence[str] = [],
    xfail_platforms: Sequence[str] = [],
) -> str:
    name_block = _get_string_arg_block("NAME", target_name)
    driver_block = _get_string_arg_block("DRIVER", driver)
    expected_output_block = _get_string_arg_block("EXPECTED_OUTPUT", expected_output)
    modules_block = _get_string_list_arg_block(
        "MODULES",
        [f"{platform}={path}" for platform, path in platform_module_map.items()],
    )
    timeout_block = _get_string_arg_block(
        "TIMEOUT", str(timeout_secs) if timeout_secs is not None else None
    )
    runner_args_block = _get_string_list_arg_block("RUNNER_ARGS", runner_args)
    labels_block = _get_string_list_arg_block("LABELS", labels)
    xfail_platforms_block = _get_string_list_arg_block(
        "XFAIL_PLATFORMS", xfail_platforms
    )
    return _convert_block_to_string(
        _build_call_rule(
            rule_name="iree_benchmark_suite_module_test",
            parameter_blocks=[
                name_block,
                driver_block,
                expected_output_block,
                timeout_block,
                modules_block,
                runner_args_block,
                labels_block,
                xfail_platforms_block,
            ],
        )
    )


def build_add_dependencies(target: str, deps: List[str]) -> str:
    if len(deps) == 0:
        raise ValueError("Target dependencies can't be empty.")
    deps_list = _get_string_list(deps, quote=False)
    return _convert_block_to_string(
        [f"add_dependencies({target}"] + _get_block_body(deps_list) + [")"]
    )


def build_set(variable_name: str, value: str) -> str:
    return _convert_block_to_string(
        [f"set({variable_name}"] + _get_block_body([value]) + [")"]
    )