#!/usr/bin/env python3

################################################################################
# Copyright (c) 2018-2022 Advanced Micro Devices, Inc.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
################################################################################

from __future__ import print_function
import os, sys, re

H_OUT='hsa_prof_str.h'
CPP_OUT='hsa_prof_str.inline.h'
API_TABLES_H = 'hsa_api_trace.h'
API_HEADERS_H = (
  ('CoreApi', 'hsa.h'),
  ('AmdExt', 'hsa_ext_amd.h'),
  ('ImageExt', 'hsa_ext_image.h'),
  ('AmdExt', API_TABLES_H),
)

LICENSE = \
'/* Copyright (c) 2018-2022 Advanced Micro Devices, Inc.\n' + \
'\n' + \
' Permission is hereby granted, free of charge, to any person obtaining a copy\n' + \
' of this software and associated documentation files (the "Software"), to deal\n' + \
' in the Software without restriction, including without limitation the rights\n' + \
' to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n' + \
' copies of the Software, and to permit persons to whom the Software is\n' + \
' furnished to do so, subject to the following conditions:\n' + \
'\n' + \
' The above copyright notice and this permission notice shall be included in\n' + \
' all copies or substantial portions of the Software.\n' + \
'\n' + \
' THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n' + \
' IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n' + \
' FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE\n' + \
' AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n' + \
' LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n' + \
' OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\n' + \
' THE SOFTWARE. */\n'

#############################################################
# Error handler
def fatal(module, msg):
  print (module + ' Error: "' + msg + '"', file = sys.stderr)
  sys.exit(1)

# Get next text block
def NextBlock(pos, record):
  if len(record) == 0: return pos

  space_pattern = re.compile(r'(\s+)')
  word_pattern = re.compile(r'([\w\*]+)')
  if record[pos] != '(':
    m = space_pattern.match(record, pos)
    if not m:
      m = word_pattern.match(record, pos)
    if m:
      return pos + len(m.group(1))
    else:
      fatal('NextBlock', "bad record '" + record + "' pos(" + str(pos) + ")")
  else:
    count = 0
    for index in range(pos, len(record)):
      if record[index] == '(':
        count = count + 1
      elif record[index] == ')':
        count = count - 1
        if count == 0:
          index = index + 1
          break
    if count != 0:
      fatal('NextBlock', "count is not zero (" + str(count) + ")")
    if record[index - 1] != ')':
      fatal('NextBlock', "last char is not ')' '" + record[index - 1] + "'")
    return index

#############################################################
# API table parser class
class API_TableParser:
  def fatal(self, msg):
    fatal('API_TableParser', msg)

  def __init__(self, header, name):
    self.name = name

    if not os.path.isfile(header):
      self.fatal("file '" + header + "' not found")

    self.inp = open(header, 'r')

    self.beg_pattern = re.compile('^\s*struct\s+' + name + 'Table\s*{\s*$')
    self.end_pattern = re.compile('^\s*};\s*$')
    self.array = []
    self.parse()

  # normalizing a line
  def norm_line(self, line):
    return re.sub(r'^\s+', r' ', line[:-1])

  # check for start record
  def is_start(self, record):
    return self.beg_pattern.match(record)

  # check for end record
  def is_end(self, record):
    return self.end_pattern.match(record)

  # check for declaration entry record
  def is_entry(self, record):
    return re.match(r'^\s*decltype\(([^\)]*)\)', record)

  # parse method
  def parse(self):
    active = 0
    for line in self.inp.readlines():
      record = self.norm_line(line)
      if self.is_start(record): active = 1
      if active != 0:
        if self.is_end(record): return
        m = self.is_entry(record)
        if m:
          self.array.append(m.group(1))

#############################################################
# API declaration parser class
class API_DeclParser:
  def fatal(self, msg):
    fatal('API_DeclParser', msg)

  def __init__(self, header, array, data):
    if not os.path.isfile(header):
      self.fatal("file '" + header + "' not found")

    self.inp = open(header, 'r')

    self.end_pattern = re.compile('\);\s*$')
    self.data = data
    for call in array:
      if call in data:
        self.fatal(call + ' is already found')
      self.parse(call)

  # api record filter
  def api_filter(self, record):
    record = re.sub(r'\sHSA_API\s', r' ', record)
    record = re.sub(r'\sHSA_DEPRECATED\s', r' ', record)
    return record

  # check for start record
  def is_start(self, call, record):
    return re.search('\s' + call + '\s*\(', record)

  # check for API method record
  def is_api(self, call, record):
    record = self.api_filter(record)
    return re.match('\s+\S+\s+' + call + '\s*\(', record)

  # check for end record
  def is_end(self, record):
    return self.end_pattern.search(record)

  # parse method args
  def get_args(self, record):
    struct = {'ret': '', 'args': '', 'astr': {}, 'alst': [], 'tlst': []}
    record = re.sub(r'^\s+', r'', record)
    record = re.sub(r'\s*(\*+)\s*', r'\1 ', record)
    rind = NextBlock(0, record)
    struct['ret'] = record[0:rind]
    pos = record.find('(')
    end = NextBlock(pos, record);
    args = record[pos:end]
    args = re.sub(r'^\(\s*', r'', args)
    args = re.sub(r'\s*\)$', r'', args)
    args = re.sub(r'\s*,\s*', r',', args)
    struct['args'] = re.sub(r',', r', ', args)
    if len(args) == 0: return struct

    pos = 0
    args = args + ','
    while pos < len(args):
      ind1 = NextBlock(pos, args) # type
      ind2 = NextBlock(ind1, args) # space
      if args[ind2] != '(':
        while ind2 < len(args):
          end = NextBlock(ind2, args)
          if args[end] == ',': break
          else: ind2 = end
        name = args[ind2:end]
      else:
        ind3 = NextBlock(ind2, args) # field
        m = re.match(r'\(\s*\*\s*(\S+)\s*\)', args[ind2:ind3])
        if not m:
          self.fatal("bad block3 '" + args + "' : '" + args[ind2:ind3] + "'")
        name = m.group(1)
        end = NextBlock(ind3, args) # the rest
      item = args[pos:end]
      struct['astr'][name] = item
      struct['alst'].append(name)
      struct['tlst'].append(item)
      if args[end] != ',':
        self.fatal("no comma '" + args + "'")
      pos = end + 1

    return struct

  # parse given api
  def parse(self, call):
    record = ''
    active = 0
    found = 0
    api_name = ''
    prev_line = ''

    self.inp.seek(0)
    for line in self.inp.readlines():
      record += ' ' + line[:-1]
      record = re.sub(r'^\s*', r' ', record)

      if active == 0:
        if self.is_start(call, record):
          active = 1
          m = self.is_api(call, record)
          if not m:
            record = ' ' + prev_line + ' ' + record
            m = self.is_api(call, record)
            if not m:
              self.fatal("bad api '" + line + "'")

      if active == 1:
        if self.is_end(record):
          self.data[call] = self.get_args(record)
          active = 0
          found = 0

      if active == 0: record = ''
      prev_line = line

#############################################################
# API description parser class
class API_DescrParser:
  def fatal(self, msg):
    fatal('API_DescrParser', msg)

  def __init__(self, out_h_file, hsa_dir, api_table_h, api_headers, license):
    out_macro = re.sub(r'[\/\.]', r'_', out_h_file.upper()) + '_'

    self.h_content = ''
    self.cpp_content = ''
    self.api_names = []
    self.api_calls = {}
    self.api_rettypes = set()
    self.api_id = {}

    api_data = {}
    api_list = []
    ns_calls = []

    for i in range(0, len(api_headers)):
      (name, header) = api_headers[i]

      if i < len(api_headers) - 1:
        api = API_TableParser(hsa_dir + api_table_h, name)
        api_list = api.array
        self.api_names.append(name)
        self.api_calls[name] = api_list
      else:
        api_list = ns_calls
        ns_calls = []

      for call in api_list:
        if call in api_data:
          self.fatal("call '"  + call + "' is already found")

      API_DeclParser(hsa_dir + header, api_list, api_data)

      for call in api_list:
        if not call in api_data:
          # Not-supported functions
          ns_calls.append(call)
        else:
          # API ID map
          self.api_id[call] = 'HSA_API_ID_' + call
          # Return types
          self.api_rettypes.add(api_data[call]['ret'])

    self.api_rettypes.discard('void')
    self.api_data = api_data
    self.ns_calls = ns_calls

    self.h_content += "/* Generated by " + os.path.basename(__file__) + " */\n" + license + "\n\n"

    self.h_content += "/* HSA API tracing primitives\n"
    for (name, header) in api_headers:
      self.h_content += " '" + name + "', header '" + header + "', " + str(len(self.api_calls[name])) + ' funcs\n'
    for call in self.ns_calls:
      self.h_content += ' ' + call + ' was not parsed\n'
    self.h_content += " */\n"
    self.h_content += '\n'
    self.h_content += '#ifndef ' + out_macro + '\n'
    self.h_content += '#define ' + out_macro + '\n'

    self.h_content += self.add_section('API ID enumeration', '  ', self.gen_id_enum)

    self.h_content += '/* Declarations of APIs intended for use only by tools. */\n'
    self.h_content += 'typedef void (*hsa_amd_queue_intercept_packet_writer)(const void*, uint64_t);\n'
    self.h_content += 'typedef void (*hsa_amd_queue_intercept_handler)(const void*, uint64_t, uint64_t, void*,\n'
    self.h_content += '                                                hsa_amd_queue_intercept_packet_writer);\n'
    self.h_content += 'typedef void (*hsa_amd_runtime_queue_notifier)(const hsa_queue_t*, hsa_agent_t, void*);\n'

    self.h_content += self.add_section('API arg structure', '    ', self.gen_arg_struct)
    self.h_content += self.add_section('API output stream', '    ', self.gen_out_stream)
    self.h_content += '#endif /* ' + out_macro + ' */\n'

    self.cpp_content += "/* Generated by " + os.path.basename(__file__) + " */\n" + license + "\n\n"

    self.cpp_content += '#include <hsa/hsa_api_trace.h>\n'
    self.cpp_content += '#include <atomic>\n'
    self.cpp_content += 'namespace roctracer::hsa_support::detail {\n'

    self.cpp_content += 'static CoreApiTable CoreApi_saved_before_cb;\n'
    self.cpp_content += 'static AmdExtTable AmdExt_saved_before_cb;\n'
    self.cpp_content += 'static ImageExtTable ImageExt_saved_before_cb;\n\n'

    self.cpp_content += self.add_section('API callback functions', '', self.gen_callbacks)
    self.cpp_content += self.add_section('API intercepting code', '', self.gen_intercept)
    self.cpp_content += self.add_section('API get_name function', '    ', self.gen_get_name)
    self.cpp_content += self.add_section('API get_code function', '  ', self.gen_get_code)
    self.cpp_content += '\n};\n'

  # add code section
  def add_section(self, title, gap, fun):
    content = ''
    n = 0
    content +=  '\n/* section: ' + title + ' */\n\n'
    content += fun(-1, '-', '-', {})
    for index in range(len(self.api_names)):
      last = (index == len(self.api_names) - 1)
      name = self.api_names[index]
      if n != 0:
        if gap == '': content += fun(n, name, '-', {})
        content += '\n'
      content += gap + '/* block: ' + name + ' API */\n'
      for call in self.api_calls[name]:
        content += fun(n, name, call, self.api_data[call])
        n += 1
    content += fun(n, '-', '-', {})
    return content

  # generate API ID enumeration
  def gen_id_enum(self, n, name, call, data):
    content = ''
    if n == -1:
      content += 'enum hsa_api_id_t {\n'
      return content
    if call != '-':
      content += '  ' + self.api_id[call] + ' = ' + str(n) + ',\n'
    else:
      content += '\n'
      content += '  HSA_API_ID_DISPATCH = ' + str(n) + ',\n'
      content += '  HSA_API_ID_NUMBER = ' + str(n + 1) + ',\n'
      content += '};\n'
    return content

  # generate API args structure
  def gen_arg_struct(self, n, name, call, struct):
    content = ''
    if n == -1:
      content += 'typedef struct hsa_api_data_s {\n'
      content += '  uint64_t correlation_id;\n'
      content += '  uint32_t phase;\n'
      content += '  union {\n'
      for ret_type in self.api_rettypes:
        content += '    ' + ret_type + ' ' + ret_type + '_retval;\n'
      content += '  };\n'
      content += '  union {\n'
      return content
    if call != '-':
      content +=   '    struct {\n'
      for (var, item) in struct['astr'].items():
        content += '      ' + item + ';\n'
        if call == "hsa_amd_memory_async_copy_rect" and item == "const hsa_dim3_t* range":
          content += '      hsa_dim3_t range__val;\n'
      content +=   '    } ' + call + ';\n'
    else:
      content += '  } args;\n'
      content += '  uint64_t *phase_data;\n'
      content += '} hsa_api_data_t;\n'
    return content

  # generate API callbacks
  def gen_callbacks(self, n, name, call, struct):
    content = ''
    if n == -1:
      content += '/* section: Static declarations */\n'
      content += '\n'
    if call != '-':
      call_id = self.api_id[call];
      ret_type = struct['ret']
      content += 'static ' + ret_type + ' ' + call + '_callback(' + struct['args'] + ') {\n'

      content += '  hsa_trace_data_t trace_data;\n'
      content += '  bool enabled{false};\n'
      content += '\n'
      content += '  if (auto function = report_activity.load(std::memory_order_relaxed); function &&\n'
      content += '      (enabled =\n'
      content += '           function(ACTIVITY_DOMAIN_HSA_API, ' + call_id + ', &trace_data) == 0)) {\n'
      content += '    if (trace_data.phase_enter != nullptr) {\n'

      for var in struct['alst']:
        item = struct['astr'][var];
        if re.search(r'char\* ', item):
          # FIXME: we should not strdup the char* arguments here, as the callback will not outlive the scope of this function. Instead, we
          # should generate a helper function to capture the content of the arguments similar to hipApiArgsInit for HIP. We also need a
          # helper to free the memory that is allocated to capture the content.
          content += '      trace_data.api_data.args.' + call + '.' + var + ' = ' + '(' + var + ' != NULL) ? strdup(' + var + ')' + ' : NULL;\n'
        else:
          content += '      trace_data.api_data.args.' + call + '.' + var + ' = ' + var + ';\n'
          if call == 'hsa_amd_memory_async_copy_rect' and var == 'range':
            content += '      trace_data.api_data.args.' + call + '.' + var + '__val = ' + '*(' + var + ');\n'

      content += '      trace_data.phase_enter(' + call_id + ', &trace_data);\n'
      content += '    }\n'
      content += '  }\n'
      content += '\n'

      if ret_type != 'void':
        content +=  '  trace_data.api_data.' + ret_type + '_retval = '
      content += '  ' + name + '_saved_before_cb.' + call + '_fn(' + ', '.join(struct['alst']) + ');\n'

      content += '\n'
      content += '  if (enabled && trace_data.phase_exit != nullptr)\n'
      content += '    trace_data.phase_exit(' + call_id + ', &trace_data);\n'

      if ret_type != 'void':
        content += '  return trace_data.api_data.' + ret_type + '_retval;\n'
      content += '}\n'

    return content

  # generate API intercepting code
  def gen_intercept(self, n, name, call, struct):
    content = ''
    if n > 0 and call == '-':
      content += '};\n'
    if n == 0 or (call == '-' and name != '-'):
      content += 'static void Install' + name + 'Wrappers(' + name + 'Table* table) {\n'
      content += '  ' + name + '_saved_before_cb = *table;\n'
    if call != '-':
      if call != 'hsa_shut_down':
        content += '  table->' + call + '_fn = ' + call + '_callback;\n'
      else:
        content += '  { void* p = (void*)' + call + '_callback; (void)p; }\n'
    return content

  # generate API name function
  def gen_get_name(self, n, name, call, struct):
    content = ''
    if n == -1:
      content += 'static const char* GetApiName(uint32_t id) {\n'
      content += '  switch (id) {\n'
      return content
    if call != '-':
      content += '    case ' + self.api_id[call] + ': return "' + call + '";\n'
    else:
      content += '  }\n'
      content += '  return "unknown";\n'
      content += '}\n'
    return content

  # generate API code function
  def gen_get_code(self, n, name, call, struct):
    content = ''
    if n == -1:
      content += 'static uint32_t GetApiCode(const char* str) {\n'
      return content
    if call != '-':
      content += '  if (strcmp("' + call + '", str) == 0) return ' + self.api_id[call] + ';\n'
    else:
      content += '  return HSA_API_ID_NUMBER;\n'
      content += '}\n'
    return content

  # generate stream operator
  def gen_out_stream(self, n, name, call, struct):
    content = ''
    if n == -1:
      content += '#ifdef __cplusplus\n'
      content += '#include "hsa_ostream_ops.h"\n'
      content += 'typedef std::pair<uint32_t, hsa_api_data_t> hsa_api_data_pair_t;\n'
      content += 'inline std::ostream& operator<< (std::ostream& out, const hsa_api_data_pair_t& data_pair) {\n'
      content += '  const uint32_t cid = data_pair.first;\n'
      content += '  const hsa_api_data_t& api_data = data_pair.second;\n'
      content += '  switch(cid) {\n'
      return content
    if call != '-':
      content += '    case ' + self.api_id[call] + ': {\n'
      content += '      out << "' + call + '(";\n'
      arg_list = struct['alst']
      if len(arg_list) != 0:
        for ind in range(len(arg_list)):
          arg_var = arg_list[ind]
          arg_val = 'api_data.args.' + call + '.' + arg_var
          if re.search(r'char\* ', struct['astr'][arg_var]):
            content += '      out << "0x" << std::hex << (uint64_t)' + arg_val
          else:
            content += '      out << ' + arg_val
            if call == "hsa_amd_memory_async_copy_rect" and arg_var == "range":
              content += ' << ", ";\n'
              content += '      out << ' + arg_val + '__val'
          '''
          arg_item = struct['tlst'][ind]
          if re.search(r'\(\* ', arg_item): arg_pref = ''
          elif re.search(r'void\* ', arg_item): arg_pref = ''
          elif re.search(r'\*\* ', arg_item): arg_pref = '**'
          elif re.search(r'\* ', arg_item): arg_pref = '*'
          else: arg_pref = ''
          if arg_pref != '':
            content += '      if (' + arg_val + ') out << ' + arg_pref + '(' + arg_val + '); else out << ' + arg_val
          else:
            content += '      out << ' + arg_val
          '''
          if ind < len(arg_list) - 1: content += ' << ", ";\n'
          else: content += ';\n'
      if struct['ret'] != 'void':
        content += '      out << ") = " << api_data.' + struct['ret'] + '_retval;\n'
      else:
        content += '      out << ") = void";\n'
      content += '      break;\n'
      content += '    }\n'
    else:
      content += '    default:\n'
      content += '      out << "ERROR: unknown API";\n'
      content += '      abort();\n'
      content += '  }\n'
      content += '  return out;\n'
      content += '}\n'
      content += '#endif\n'
    return content

#############################################################
# main
# Usage
if len(sys.argv) != 3:
  print ("Usage:", sys.argv[0], " <OUT prefix> <HSA runtime include path>", file=sys.stderr)
  sys.exit(1)
else:
  PREFIX = sys.argv[1] + '/'
  HSA_DIR = sys.argv[2] + '/'

descr = API_DescrParser(H_OUT, HSA_DIR, API_TABLES_H, API_HEADERS_H, LICENSE)

out_file = PREFIX + H_OUT
print ('Generating "' + out_file + '"')
f = open(out_file, 'w')
f.write(descr.h_content[:-1])
f.close()

out_file = PREFIX + CPP_OUT
print ('Generating "' + out_file + '"')
f = open(out_file, 'w')
f.write(descr.cpp_content[:-1])
f.close()
#############################################################