#!/usr/bin/env python3

'''
A simple script for expanding macro definitions in source files.
'''

import argparse
import re
import os
import sys

parser = argparse.ArgumentParser(description='Expand RBT macros. The approach used is naive and may fail. Always check the output.')
parser.add_argument(
  'files', metavar='<path>', nargs='+',
  help='The source files in which to expand macros. For each input file, an output file will be created with the chosen extension (see the -e option). Check this file before overwriting the original.'
)
parser.add_argument(
  '-d', '--rbt-dir', metavar='<path>', required=True,
  help='The directory containing the RBT header files.'
)
parser.add_argument(
  '-i', '--include', metavar='<path>', action='append', default=[],
  help='Additional directories in which to search for included files.'
)
parser.add_argument(
  '-e', '--extension', metavar='<extension>', default=".expanded",
  help='The extension to use when saving expanded files. Default: %(default)s'
)
parser.add_argument(
  '-o', '--overwrite', action='store_true',
  help='Overwrite existing output files (with appended extension). Without this option, a numerical suffix will be added after the extension.'
)

def iterate(path, pargs, show=True):
  with open(path, 'r') as f:
    for line in f:
      m = re.search(r'#include\s+"([^"]+)"', line)
      if m:
        isubpath = m.group(1)
        ipath = os.path.join(os.path.dirname(path), isubpath)
        try:
          for s, l in iterate(ipath, pargs, show=False):
            yield s, l
        except FileNotFoundError:
          for idir in pargs.include + [pargs.rbt_dir]:
            ipath = os.path.join(idir, isubpath)
            try:
              for s, l in iterate(ipath, pargs, show=False):
                yield s, l
              break
            except FileNotFoundError:
              pass
          else:
            sys.stderr.write(
              "warning: could not find {} included in {}\n".format(isubpath, path)
            )
      yield show, line


def expand(path, defs, pargs):
  name, ext = os.path.splitext(path)
  name += pargs.extension
  tmp_path = name + ext
  if not pargs.overwrite:
    i = 1
    while os.path.exists(tmp_path):
      tmp_path = "{}.{:d}{}".format(name, i, ext)
      i += 1
  output_path = tmp_path

  local_defs = dict()

  with open(output_path, 'w') as f:
    for show, line in iterate(path, pargs):
      m = re.search(r'^\s*#define\s+(RBT_\S+?_PREFIX_)\s+(\S+)', line)
      if m:
        key = m.group(1)
        value = m.group(2)
        # Macros may be defined in terms of other macros.
        while True:
          try:
            value = local_defs[value]
          except KeyError:
            break
        print("{} -> {}".format(key, value))
        local_defs[key] = value
      if show:
        if 'RBT_' in line:
          copy = line
          for rbt in re.findall(r'RBT_[^\s(]+', copy):
            try:
              prefix, name = defs[rbt]
              expanded_name = local_defs[prefix] + name
              line = line.replace(rbt, expanded_name)
            except KeyError:
              pass
        f.write(line)
  print("expanded {}\n      to {}\n".format(path, output_path))


def main(args=None):
  pargs = parser.parse_args(args)
  defs = dict()
  for root, dnames, fnames in os.walk(pargs.rbt_dir):
    for fname in fnames:
      if fname[-2:] == '.h':
        path = os.path.join(root, fname)
        with open(path, 'r') as f:
           for line in f:
            m = re.search(r'(?:#define|typedef)\s+(\S+)\s+RBT_TOKEN_2_W\(\s*(\S+)\s*,\s*(\S+)\s*\)', line)
            if m:
              defs[m.group(1)] = (m.group(2), m.group(3))

  for path in pargs.files:
    expand(path, defs, pargs)

if __name__ == '__main__':
  try:
    main()
  except (KeyboardInterrupt, BrokenPipeError):
    pass