#!/usr/bin/env python3
# Copyright 2026 The Emscripten Authors.  All rights reserved.
# Emscripten is available under two separate licenses, the MIT license and the
# University of Illinois/NCSA Open Source License.  Both these licenses can be
# found in the LICENSE file.

"""Extract native signature information for use in pointer conversion wrappers.

It generates a file (or prints) a signature mapping for native functions.
"""

import argparse
import os
import sys

__scriptdir__ = os.path.dirname(os.path.abspath(__file__))
__rootdir__ = os.path.dirname(os.path.dirname(__scriptdir__))
sys.path.insert(0, __rootdir__)

from tools import shared, utils, webassembly

# Overrides to convert 'p' to 'P' (allow undefined) in signatures.
# Maps symbol name to the list of 0-based argument indices that should be 'P'.
signature_overrides = {
  'sbrk': [0],      # sbrk(increment)
  'main': [1, 2],   # main(argc, argv, envp)
}


def apply_overrides(sym, sig):
  if sym in signature_overrides:
    indices = signature_overrides[sym]
    sig_chars = list(sig)
    for idx in indices:
      sig_idx = idx + 1
      if sig_idx < len(sig_chars) and sig_chars[sig_idx] == 'p':
        sig_chars[sig_idx] = 'P'
    return ''.join(sig_chars)
  return sig


# WebIDL functions might not be linked in unless we use webidl binder.
# We can just define them in our dummy file to get their signatures,
# since they are just wrappers with predictable signatures.
# Same for other symbols that are hard to link.
extra_declarations = '''
#include <stddef.h>
#include <stdint.h>
extern "C" {
  void* webidl_malloc(size_t len) { return nullptr; }
  void webidl_free(void* ptr) {}
  __attribute__((weak)) uintptr_t __get_tp() { return 0; }
  __attribute__((weak)) void __set_stack_limits(void* low, void* high) {}
}
'''


def extract_sigs(wasm_file):
  sig_info = {}
  with webassembly.Module(wasm_file) as mod:
    exports = mod.get_exports()
    for exp in exports:
      # Ignore C++ mangled symbols (starting with '_Z')
      if exp.name.startswith('_Z'):
        continue
      # Ignore `dynCall_` symbols generated by binaryen.
      if exp.name.startswith('dynCall_'):
        continue
      if exp.kind == webassembly.ExternType.FUNC:
        sig_info[exp.name] = mod.get_function_type(exp.index)
  return sig_info


def valuetype_to_chr(t, t64):
  if t == webassembly.Type.I32 and t64 == webassembly.Type.I64:
    return 'p'
  return '_'


def functype_to_str(t, t64):
  assert len(t.returns) == len(t64.returns)
  assert len(t.params) == len(t64.params)
  if t.returns:
    assert len(t.returns) == 1
    rtn = valuetype_to_chr(t.returns[0], t64.returns[0])
  else:
    rtn = '_'
  for p, p64 in zip(t.params, t64.params, strict=True):
    rtn += valuetype_to_chr(p, p64)
  return rtn


def run_build(config_flags, explicit_exports, main_args, suffix):
  utils.safe_ensure_dirs(utils.path_from_root('out'))
  cpp_file = utils.path_from_root('out/gen_native_sig_info.cpp')
  # We write the dummy file with extra declarations to ensure they are defined
  # and exported (since MAIN_MODULE=1 will export them if they are defined in the main module).
  utils.write_file(cpp_file, f'int main({main_args}) {{ return 0; }}\n' + extra_declarations)

  wasm32 = utils.path_from_root(f'out/test_{suffix}32')
  wasm64 = utils.path_from_root(f'out/test_{suffix}64')

  flags = config_flags.copy()

  if explicit_exports:
    exports_str = ','.join(explicit_exports)
    flags.append(f'-sEXPORTED_FUNCTIONS=_main,{exports_str}')
  else:
    flags += ['-Wl,--whole-archive', '-Wl,--export-all']

  # Wasm32 build
  cmd = [shared.EMXX, cpp_file, '-o', wasm32 + '.js', *flags]
  print(f"Building wasm32 for {suffix}...")
  shared.check_call(cmd)

  # Wasm64 build
  cmd_64 = [shared.EMXX, cpp_file, '-o', wasm64 + '.js', '-m64', *flags]
  print(f"Building wasm64 for {suffix}...")
  shared.check_call(cmd_64)

  return wasm32 + '.wasm', wasm64 + '.wasm'


def main(args):
  parser = argparse.ArgumentParser()
  default_output = utils.path_from_root('tools/native_sigs.py')
  parser.add_argument('-o', '--output', default=default_output, help='Python output file (default: tools/native_sigs.py, use "-" for stdout)')
  args = parser.parse_args()

  # Define configurations to cover different sets of symbols.
  # We use MAIN_MODULE=1 to export everything that gets linked.
  configs = {
    'base': (['-pthread', '-sASYNCIFY', '-sSTB_IMAGE', '--bind', '-sSTACK_OVERFLOW_CHECK=1'], [], 'int argc, char* argv[], char* envp[]'),
    'wasmfs': (['-sWASMFS', '-sSTB_IMAGE', '--bind', '-sSTACK_OVERFLOW_CHECK=1'], [], ''),
    'wasm_workers': (['-sWASM_WORKERS', '-sSTACK_OVERFLOW_CHECK=1'], ['__emscripten_wasm_worker_initialize'], ''),
    'main_module': (['-sMAIN_MODULE', '-pthread', '-Wno-experimental'], [], ''),
  }

  merged_sigs = {}

  for name, (flags, exports, main_args) in configs.items():
    wasm32, wasm64 = run_build(flags, exports, main_args, name)
    sigs32 = extract_sigs(wasm32)
    sigs64 = extract_sigs(wasm64)

    for sym, sig32 in sigs32.items():
      sig64 = sigs64[sym]
      type_str = functype_to_str(sig32, sig64)
      if sym in merged_sigs:
        if merged_sigs[sym] != type_str:
          print(f'symbol type mismatch for {sym}: existing={merged_sigs[sym]} vs new={type_str}')
          continue
      merged_sigs[sym] = type_str

  # Filter results to only include signatures containing pointers ('p')
  result_sigs = {}
  for sym, sig in merged_sigs.items():
    if 'p' in sig.lower():
      result_sigs[sym] = apply_overrides(sym, sig)

  # Duplicate main signature to __main_argc_argv
  assert '__main_argc_argv' not in result_sigs
  result_sigs['__main_argc_argv'] = result_sigs['main']

  content = '# Auto-generated by tools/maint/gen_native_sig_info.py. DO NOT EDIT.\n'
  content += 'native_sigs = {\n'
  for sym, sig in sorted(result_sigs.items()):
    content += f"  '{sym}': '{sig}',\n"
  content += '}\n'

  if args.output == '-':
    print(content, end='')
  else:
    utils.write_file(args.output, content)


if __name__ == '__main__':
  sys.exit(main(sys.argv[1:]))
