#!/usr/bin/env python3
# Copyright 2020 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.
"""This script finds all Python recipe files that are missing a copyright
notice and adds a copyright, modifying the files in-place."""

import datetime
import os

COPYRIGHT = f"""# Copyright {datetime.datetime.now().year} 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.
"""


def main():
    cwd = os.getcwd()
    recipe_dirs = [os.path.join(cwd, "recipes"), os.path.join(cwd, "recipe_modules")]
    for directory in recipe_dirs:
        for subdir, _, files in os.walk(directory):
            for relpath in files:
                _, ext = os.path.splitext(relpath)
                if ext != ".py":
                    continue
                path = os.path.join(subdir, relpath)
                ensure_copyright(path)


def ensure_copyright(path):
    with open(path, "r") as f:
        orig_lines = f.readlines()

    add_copyright = False
    copyright_line_index = 0
    if orig_lines:
        # Shebangs and Python 2 encoding lines must come on the first line, so
        # make sure to not insert copyright notices ahead of those types of
        # lines.
        if orig_lines[0].startswith(("#!", "# -*- coding")):
            copyright_line_index = 1
        if len(orig_lines) <= 1 or not orig_lines[copyright_line_index].startswith(
            "# Copyright"
        ):
            add_copyright = True
    else:
        add_copyright = True

    if not add_copyright:
        return

    copyright_lines = [l + "\n" for l in COPYRIGHT.splitlines()]

    # Ensure there's a blank line between the copyright statement the next
    # line, if the next line isn't a docstring.
    if copyright_line_index < len(orig_lines):
        next_line = orig_lines[copyright_line_index]
        if next_line.strip() != "" and not next_line.startswith(('"""', "'''")):
            copyright_lines.append("\n")

    with open(path, "w") as f:
        new_lines = orig_lines[:]
        new_lines[copyright_line_index:copyright_line_index] = copyright_lines
        f.write("".join(new_lines))


if __name__ == "__main__":
    main()