#!/usr/bin/env python3

# Copyright 2020 Google LLC
#
# Licensed 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
#
#      https://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.
"""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}
{middle_comment} Licensed under the Apache License, Version 2.0 (the "License");
{middle_comment} you may not use this file except in compliance with the License.
{middle_comment} You may obtain a copy of the License at
{middle_comment}
{middle_comment}      https://www.apache.org/licenses/LICENSE-2.0
{middle_comment}
{middle_comment} Unless required by applicable law or agreed to in writing, software
{middle_comment} distributed under the License is distributed on an "AS IS" BASIS,
{middle_comment} WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
{middle_comment} See the License for the specific language governing permissions and
{middle_comment} limitations under the License.{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: Google LLC",
      default="Google LLC")
  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())