#!/usr/bin/env python3
# Copyright 2023 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import argparse
import ast
import dataclasses
import json
import os


@dataclasses.dataclass
class Finding:
    message: str
    line: int
    col: int = 0
    end_line: int = 0
    end_col: int = 0
    replacements: list = ()
    level: str = "error"


def main():
    parser = argparse.ArgumentParser(description="Enforce the recipe style guide")
    parser.add_argument("filename", help="Name of the file to check.")
    args = parser.parse_args()

    with open(args.filename) as f:
        tree = ast.parse(f.read(), filename=args.filename)

    findings = []

    # PROPERTIES declarations may appear in recipe files or in recipe module
    # __init__.py files.
    findings.extend(check_properties_type(args.filename, tree))
    if is_recipe_file(args.filename):
        findings.extend(check_recipe_function_naming(tree))
        findings.extend(check_recipe_function_ordering(tree))

    print(
        json.dumps(
            [{"filepath": args.filename, **dataclasses.asdict(f)} for f in findings],
            indent=2,
        )
    )


def is_recipe_file(path):
    if not path.endswith(".py"):
        return False

    path_parts = path.split(os.sep)
    if any(p.endswith(".resources") for p in path_parts):
        return False
    if path_parts[0] == "recipes":
        return True
    if path_parts[0] == "recipe_modules" and path_parts[-2] in ("tests", "examples"):
        return True
    return False


def check_recipe_function_naming(tree: ast.AST):
    for node in tree.body:
        if isinstance(node, ast.FunctionDef) and node.name.startswith("_"):
            # node.col_offset is the column of the "def" keyword, but it's nicer
            # if we only highlight the function name instead of the whole "def"
            # line. Note that this is fragile, as it assumes only one space
            # between the "def" keyword and the function name, which may not be
            # the case if the code has not been auto-formatted.
            col = node.col_offset + len("def ") + 1
            yield Finding(
                line=node.lineno,
                end_line=node.lineno,
                col=col,
                end_col=col + len(node.name),
                message=f"{node.name} should not start with an underscore",
                replacements=[node.name.lstrip("_")],
            )


def check_recipe_function_ordering(tree: ast.AST):
    function_defs = [node for node in tree.body if isinstance(node, ast.FunctionDef)]
    for index, node in enumerate(function_defs):
        if node.name == "RunSteps" and index != 0:
            yield Finding(
                line=node.lineno,
                end_line=node.lineno,
                message=(
                    f"{node.name} should be the first function in a recipe. "
                    "Helper functions should be between RunSteps and GenTests."
                ),
            )
        if node.name == "GenTests" and index != len(function_defs) - 1:
            yield Finding(
                line=node.lineno,
                end_line=node.lineno,
                message=(
                    f"{node.name} should be the last function in a recipe. "
                    "Helper functions should be between RunSteps and GenTests."
                ),
            )


# TODO(fxbug.dev/88250): Trim down this allowlist. *Do not* add new files to
# this allowlist.
DICT_PROPERTIES_ALLOWLIST = [
    "recipes/contrib/dart_toolchain.py",
    "recipes/contrib/goma_client.py",
    "recipes/contrib/goma_gcp_configurator.py",
    "recipes/contrib/goma_gcp_deployer.py",
    "recipes/contrib/goma_toolchain.py",
    "recipes/contrib/goma_windows_images.py",
    "recipes/contrib/gomatools_builder.py",
    "recipes/contrib/gomatools_images_roller.py",
]


def check_properties_type(filename: str, tree: ast.AST):
    path_parts = filename.split(os.sep)
    # Recipe module test recipes are allowed to use dict-style properties for
    # simplicity.
    if path_parts[0] == "recipe_modules" and path_parts[-1] != "__init__.py":
        return
    for node in tree.body:
        if not isinstance(node, ast.Assign) or node.targets[0].id != "PROPERTIES":
            continue
        if isinstance(node.value, ast.Dict):
            if filename not in DICT_PROPERTIES_ALLOWLIST:
                yield Finding(
                    message="use proto properties instead of dict properties",
                    line=node.lineno,
                    end_line=node.end_lineno,
                    col=node.col_offset + 1,
                    end_col=node.end_col_offset + 1,
                )
        elif filename in DICT_PROPERTIES_ALLOWLIST:
            yield Finding(
                message=f"remove {filename} from DICT_PROPERTIES_ALLOWLIST",
                line=node.lineno,
                end_line=node.end_lineno,
                col=node.col_offset + 1,
                end_col=node.end_col_offset + 1,
            )


if __name__ == "__main__":
    main()