#!/usr/bin/env python3

# Copyright 2020 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
"""Prepends a license header to files that don't already have one.

By default, only operates on known filetypes but behavior can be overridden with
flags. Ignores files already containing a license as determined by the presence
of a block that looks like "Copyright SOME_YEAR"
"""

import argparse
import datetime
import os
import re
import sys

COPYRIGHT_PATTERN = re.compile(r"Copyright\s+\d+")

LICENSE_HEADER_FORMATTER = """{shebang}{start_comment} Copyright {year} {holder}
{middle_comment} Licensed under the Apache License v2.0 with LLVM Exceptions.
{middle_comment} See https://llvm.org/LICENSE.txt for license information.
{middle_comment} SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception{end_comment}

"""


class CommentSyntax(object):
    def __init__(self, start_comment, middle_comment=None, end_comment=""):
        self.start_comment = start_comment
        self.middle_comment = middle_comment if middle_comment else start_comment
        self.end_comment = end_comment


def comment_arg_parser(v):
    """Can be used to parse a comment syntax triple."""
    if v is None:
        return None
    if not isinstance(v, str):
        raise argparse.ArgumentTypeError("String expected")
    return CommentSyntax(*v.split(","))


def create_multikey(d):
    # pylint: disable=g-complex-comprehension
    return {k: v for keys, v in d.items() for k in keys}


filename_to_comment = create_multikey(
    {
        ("BUILD", "CMakeLists.txt"): CommentSyntax("#"),
    }
)

ext_to_comment = create_multikey(
    {
        (".bzl", ".cfg", ".cmake", ".overlay", ".py", ".sh", ".yml"): CommentSyntax(
            "#"
        ),
        (".cc", ".cpp", ".comp", ".fbs", ".h", ".hpp", ".inc", ".td"): CommentSyntax(
            "//"
        ),
        (".def",): CommentSyntax(";;"),
    }
)


def get_comment_syntax(args):
    """Deterime the comment syntax to use."""
    if args.comment:
        return args.comment
    basename = os.path.basename(args.filename)
    from_filename = filename_to_comment.get(basename)
    if from_filename:
        return from_filename
    _, ext = os.path.splitext(args.filename)
    return ext_to_comment.get(ext, args.default_comment)


def parse_arguments():
    """Parses command line arguments."""
    current_year = datetime.date.today().year
    parser = argparse.ArgumentParser()
    input_group = parser.add_mutually_exclusive_group()
    input_group.add_argument(
        "infile",
        nargs="?",
        type=argparse.FileType("r", encoding="UTF-8"),
        help="Input file to format. Default: stdin",
        default=sys.stdin,
    )
    parser.add_argument(
        "--filename",
        "--assume-filename",
        type=str,
        default=None,
        help=(
            "Filename to use for determining comment syntax. Default: actual name"
            "of input file."
        ),
    )
    parser.add_argument(
        "--year",
        "-y",
        help="Year to add copyright. Default: the current year ({})".format(
            current_year
        ),
        default=current_year,
    )
    parser.add_argument(
        "--holder",
        help="Copyright holder. Default: The IREE Authors",
        default="The IREE Authors",
    )
    parser.add_argument(
        "--quiet",
        help=(
            "Don't raise a runtime error on encountering an unhandled filetype."
            "Useful for running across many files at once. Default: False"
        ),
        action="store_true",
        default=False,
    )
    output_group = parser.add_mutually_exclusive_group()
    output_group.add_argument(
        "-o",
        "--outfile",
        "--output",
        help="File to send output. Default: stdout",
        type=argparse.FileType("w", encoding="UTF-8"),
        default=sys.stdout,
    )
    output_group.add_argument(
        "--in_place",
        "-i",
        action="store_true",
        help="Run formatting in place. Default: False",
        default=False,
    )
    comment_group = parser.add_mutually_exclusive_group()
    comment_group.add_argument(
        "--comment",
        "-c",
        type=comment_arg_parser,
        help="Override comment syntax.",
        default=None,
    )
    comment_group.add_argument(
        "--default_comment",
        type=comment_arg_parser,
        help="Fallback comment syntax if filename is unknown. Default: None",
        default=None,
    )
    args = parser.parse_args()

    if args.in_place and args.infile == sys.stdin:
        raise parser.error("Cannot format stdin in place")

    if not args.filename and args.infile != sys.stdin:
        args.filename = args.infile.name

    return args


def main(args):
    first_line = args.infile.readline()
    already_has_license = False
    shebang = ""
    content_lines = []
    if first_line.startswith("#!"):
        shebang = first_line
    else:
        content_lines = [first_line]
    content_lines.extend(args.infile.readlines())
    for line in content_lines:
        if COPYRIGHT_PATTERN.search(line):
            already_has_license = True
            break
    if already_has_license:
        header = shebang
    else:
        comment_syntax = get_comment_syntax(args)
        if not comment_syntax:
            if args.quiet:
                header = shebang
            else:
                raise ValueError(
                    "Could not determine comment syntax for " + args.filename
                )
        else:
            header = LICENSE_HEADER_FORMATTER.format(
                # Add a blank line between shebang and license.
                shebang=(shebang + "\n" if shebang else ""),
                start_comment=comment_syntax.start_comment,
                middle_comment=comment_syntax.middle_comment,
                # Add a blank line before the end comment.
                end_comment=(
                    "\n" + comment_syntax.end_comment
                    if comment_syntax.end_comment
                    else ""
                ),
                year=args.year,
                holder=args.holder,
            )

    # Have to open for write after we're done reading.
    if args.in_place:
        args.outfile = open(args.filename, "w", encoding="UTF-8")
    args.outfile.write(header)
    args.outfile.writelines(content_lines)


if __name__ == "__main__":
    main(parse_arguments())