#!/usr/bin/env python3
# Copyright 2016 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.
"""Script to check C and C++ file header guards.

This script accepts a list of file or directory arguments. If a given
path is a file, it runs the checker on it. If the path is a directory,
it runs the checker on all files in that directory.

In addition, this script checks for potential header guard
collisions. This is useful since we munge / to _, and so
    lib/abc/xyz/xyz.h
and
    lib/abc_xyz/xyz.h
both want to use LIB_ABC_XYZ_XYZ_H_ as a header guard.

"""

import argparse
import collections
import fileinput
import os.path
import re
import string
import sys

SOURCE_ROOT = os.path.dirname(  # $root
    os.path.dirname(  # tools
        os.path.dirname(  # style
            os.path.realpath(os.path.abspath(__file__)))))

SYSROOT_PREFIXES = [
    'ZIRCON_SYSTEM_PUBLIC',
    'ZIRCON_THIRD_PARTY_ULIB_MUSL_INCLUDE',
]
sysroot_prefix = re.compile('^(' + '|'.join(SYSROOT_PREFIXES) + ')_')

PUBLIC_PREFIXES = [
    'ZIRCON_SYSTEM_ULIB_.*_INCLUDE', 'GARNET_PUBLIC', 'PERIDOT_PUBLIC',
    'TOPAZ_PUBLIC', 'SDK'
]
public_prefix = re.compile('^(' + '|'.join(PUBLIC_PREFIXES) + ')_')

all_header_guards = collections.defaultdict(list)

pragma_once = re.compile('^#pragma once$')
disallowed_header_characters = re.compile('[^a-zA-Z0-9_]')


def adjust_for_location(header_guard):
  """Remove internal location prefix from public headers if applicable."""
  # Remove public prefixes
  header_guard = public_prefix.sub('', header_guard, 1)
  # Replace sysroot prefixes
  header_guard = sysroot_prefix.sub('SYSROOT_', header_guard, 1)

  return header_guard


def header_guard_from_path(path):
  """Compute the header guard from the path"""
  assert (path.startswith(SOURCE_ROOT))
  relative_path = path[len(SOURCE_ROOT):].strip('/')
  upper_path = relative_path.upper()
  header_guard = re.sub(disallowed_header_characters, '_', upper_path) + '_'
  header_guard = adjust_for_location(header_guard)
  return 'COBALT_' + header_guard


def check_file(path, fix_guards=False):
  """Check whether the file has a correct header guard.

    A header guard can either be a #pragma once, or else a matching set of
        #ifndef PATH_TO_FILE_
        #define PATH_TO_FILE_
        ...
        #endif  // PATH_TO_FILE_
    preprocessor directives, where both '.' and '/' in the path are
    mapped to '_', and a trailing '_' is appended.

    In either the #pragma once case or the header guard case, it is
    assumed that there is no trailing or leading whitespace.

    """

  # Only check .h files
  if path[-2:] != '.h':
    return True

  header_guard = header_guard_from_path(path)
  all_header_guards[header_guard].append(path)

  ifndef = re.compile('^#ifndef %s$' % header_guard)
  define = re.compile('^#define %s$' % header_guard)
  endif = re.compile('^#endif +// *%s$' % header_guard)

  found_pragma_once = False
  found_ifndef = False
  found_define = False
  found_endif = False

  with open(path, 'r') as f:
    for line in f.readlines():
      match = pragma_once.match(line)
      if match:
        if found_pragma_once:
          print('%s contains multiple #pragma once' % path)
          return False
        found_pragma_once = True

      match = ifndef.match(line)
      if match:
        if found_ifndef:
          print('%s contains multiple ifndef header guards' % path)
          return False
        found_ifndef = True

      match = define.match(line)
      if match:
        if found_define:
          print('%s contains multiple define header guards' % path)
          return False
        found_define = True

      match = endif.match(line)
      if match:
        if found_endif:
          print('%s contains multiple endif header guards' % path)
          return False
        found_endif = True

  if found_pragma_once:
    if found_ifndef or found_define or found_endif:
      print('%s contains both #pragma once and header guards' % path)
      return False
    if not fix_guards:
      return True

  if found_ifndef and found_define and found_endif:
    return True

  if not found_ifndef:
    print('%s did not contain ifndef part of its header guard' % path)
  elif not found_define:
    print('%s did not contain define part of its header guard' % path)
  elif not found_endif:
    print('%s did not contain endif part of its header guard' % path)
  elif fix_guards:
    if found_pragma_once:
      print('%s contained #pragma once instead of a header guard' % path)
    else:
      print('%s did not contain a header guard or the header guard did '
            'not match the file path' % path)
  else:
    print('%s contained neither a proper header guard nor #pragma once' % path)

  header_guards_fixed = False
  if fix_guards:
    header_guards_fixed = fix_header_guard(path, header_guard)

  if not header_guards_fixed:
    print('Expected header guard value is %s' % header_guard)

  return False


def fix_header_guard(path, header_guard):
  """Attempt to fix the header guard in the given file."""
  ifndef = re.compile('^#ifndef [^\s]+_H_$')
  define = re.compile('^#define [^\s]+_H_$')
  endif = re.compile('^#endif +// *[^\s]+_H_$')
  fixed_ifndef = False
  fixed_define = False
  fixed_endif = False
  fixed_pragma_once = False

  for line in fileinput.input(path, inplace=1):
    (new_line, changes) = re.subn(ifndef, '#ifndef %s' % header_guard, line)
    if changes:
      fixed_ifndef = True
      sys.stdout.write(new_line)
      continue
    (new_line, changes) = re.subn(define, '#define %s' % header_guard, line)
    if changes:
      fixed_define = True
      sys.stdout.write(new_line)
      continue
    (new_line, changes) = re.subn(endif, '#endif  // %s' % header_guard, line)
    if changes:
      fixed_endif = True
      sys.stdout.write(new_line)
      continue
    if pragma_once.match(line):
      fixed_pragma_once = True
      sys.stdout.write('#ifndef %s\n' % header_guard)
      sys.stdout.write('#define %s\n' % header_guard)
      continue
    sys.stdout.write(line)

  if fixed_pragma_once:
    with open(path, 'a') as file:
      file.write('\n')
      file.write('#endif  // %s\n' % header_guard)

  if (fixed_ifndef and fixed_define and fixed_endif) or fixed_pragma_once:
    print('Fixed!')
    return True

  print('Not fixed...')
  return False


def check_dir(p, fix_guards=False):
  """Walk recursively over a directory checking .h files"""

  def prune(d):
    if d[0] == '.' or d == 'third_party':
      return True
    return False

  bad_files = 0
  for root, dirs, paths in os.walk(p):
    # Prune dot directories like .git
    [dirs.remove(d) for d in list(dirs) if prune(d)]
    for path in paths:
      if not check_file(os.path.join(root, path), fix_guards=fix_guards):
        bad_files += 1

  return bad_files


def check_collisions():
  collisions = 0
  for header_guard, paths in iter(all_header_guards.items()):
    if len(paths) == 1:
      continue
    collisions += 1
    print('Multiple files could use %s as a header guard:' % header_guard)
    for path in paths:
      print('    %s' % path)

  return collisions


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--fix', help='Correct wrong header guards', action='store_true')
  (arg_results, other_args) = parser.parse_known_args()
  fix_guards = arg_results.fix
  errors = 0
  for p in other_args:
    p = os.path.realpath(os.path.abspath(p))
    if os.path.isdir(p):
      errors += check_dir(p, fix_guards=fix_guards)
    else:
      if not check_file(p, fix_guards=fix_guards):
        errors += 1
  errors += check_collisions()

  if not fix_guards:
    return errors
  else:
    return 0


if __name__ == '__main__':
  sys.exit(main())