#!/usr/bin/env python3
# Copyright (C) 2022 The Android Open Source Project
#
# 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
#
#      http://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.

import os
import sys
import logging
import re

ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

AMALGAMATION_MAP = {
    'python/tools/record_android_trace.py': 'tools/record_android_trace',
    'python/tools/tracebox.py': 'tools/tracebox',
    'python/tools/traceconv.py': 'tools/traceconv',
    'python/tools/trace_processor.py': 'tools/trace_processor',
    'python/tools/cpu_profile.py': 'tools/cpu_profile',
    'python/tools/heap_profile.py': 'tools/heap_profile',
}


def amalgamate_file(fname, stack=None, done=None, in_progress=None):
  stack = [] if stack is None else stack
  done = set() if done is None else done
  in_progress = set() if in_progress is None else in_progress
  if fname in in_progress:
    cycle = ' > '.join(stack + [fname])
    logging.fatal('Cycle detected in %s', cycle)
    sys.exit(1)
  if fname in done:
    return []
  logging.debug('Processing %s', fname)
  done.add(fname)
  in_progress.add(fname)
  with open(fname, encoding='utf-8') as f:
    lines = f.readlines()
  outlines = []
  for line in lines:
    if line.startswith('from perfetto') or line.startswith('import perfetto'):
      if not re.match('from perfetto[.][.\w]+\s+import\s+[*]$', line):
        logging.fatal('Error in %s on line \"%s\"', fname, line.rstrip())
        logging.fatal('Only "from perfetto.foo import *" is supported in '
                      'sources that are used in //tools and get amalgamated')
        sys.exit(1)
      pkg = line.split()[1]
      fpath = os.path.join('python', pkg.replace('.', os.sep) + '.py')
      outlines.append('\n# ----- Amalgamator: begin of %s\n' % fpath)
      outlines += amalgamate_file(fpath, stack + [fname], done, in_progress)
      outlines.append('\n# ----- Amalgamator: end of %s\n' % fpath)
    elif '__file__' in line and not 'amalgamator:nocheck' in line:
      logging.fatal('__file__ is not allowed in sources that get amalgamated.'
                    'In %s on line \"%s\"', fname, line.rstrip())
      sys.exit(1)

    else:
      outlines.append(line)
  in_progress.remove(fname)
  logging.debug('%s: %d lines', fname, len(outlines))
  return outlines


def amalgamate(src, dst, check_only=False):
  lines = amalgamate_file(src)
  banner = '''
# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
# DO NOT EDIT. Auto-generated by %s
# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
'''
  lines.insert(lines.index('\n'), banner % os.path.relpath(__file__, ROOT_DIR))
  new_content = ''.join(lines)

  if check_only:
    if not os.path.exists(dst):
      return False
    with open(dst, encoding='utf-8') as f:
      return f.read() == new_content

  logging.info('Amalgamating %s -> %s', src, dst)
  with open(dst + '.tmp', 'w', encoding='utf-8') as f:
    f.write(new_content)
  os.chmod(dst + '.tmp', 0o755)
  os.rename(dst + '.tmp', dst)
  return True


def main():
  check_only = '--check-only' in sys.argv
  logging.basicConfig(
    format='%(levelname)-8s: %(message)s',
    level=logging.DEBUG if '-v' in sys.argv else logging.INFO)
  os.chdir(ROOT_DIR)  # Make the execution cwd-independent.
  success = True
  for src, dst in AMALGAMATION_MAP.items():
    success = success and amalgamate(src, dst, check_only)
  return 0 if success else 1


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