#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <inttypes.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <libmcpu.h>
#include <mcpu-elf.h>

#include "mcpu-archive.h"
#include "mcpu-melf.h"
#include "mcpu-version.h"


typedef struct mcpu_nm_shdr
{
  __mpu_uint32_t  type;
  __mpu_uint64_t  flags;
  __mpu_uint64_t  offset;
  __mpu_uint64_t  size;
  __mpu_uint32_t  link;
  __mpu_uint64_t  entsize;
} mcpu_nm_shdr;


typedef struct mcpu_nm_file
{
  const unsigned char *data;
  size_t               size;
  int                  encoding;
  __mpu_uint16_t       shnum;
  mcpu_nm_shdr        *shdr;
} mcpu_nm_file;


typedef struct mcpu_nm_symbol
{
  char                *name;
  __mpu_uint64_t       value;
  __mpu_uint64_t       size;
  __mpu_uint16_t       shndx;
  unsigned int         bind;
  unsigned int         type;
  char                 code;
  size_t               order;
} mcpu_nm_symbol;


typedef struct mcpu_nm_symbols
{
  mcpu_nm_symbol *symbol;
  size_t          count;
} mcpu_nm_symbols;


typedef struct mcpu_nm_options
{
  int all_symbols;
  int extern_only;
  int undefined_only;
  int numeric_sort;
  int reverse_sort;
  int no_sort;
  int dynamic;
  int print_file_name;
} mcpu_nm_options;


static const mcpu_nm_options *sort_options;


static void
usage( FILE *fp, const char *name )
{
  fprintf( fp,
    "usage: %s [options] MELF-file(s)\n\n"
    "List symbols from LibMCPU MELF objects, executables, shared objects and archives.\n\n"
    "Options:\n"
    "  -a, --debug-syms          - display all named symbols, including FILE/SECTION\n"
    "  -g, --extern-only         - display only external GLOBAL/WEAK symbols\n"
    "  -u, --undefined-only      - display only undefined symbols\n"
    "  -n, --numeric-sort        - sort symbols numerically by value\n"
    "  -r, --reverse-sort        - reverse the selected sort order\n"
    "  -p, --no-sort             - preserve symbol-table order\n"
    "  -D, --dynamic             - display .dynsym instead of .symtab\n"
    "  -A, --print-file-name     - prefix each symbol with its input file name\n"
    "      --help                - show this help\n"
    "      --version             - show program version\n"
    "      --dump-version        - show version number only\n",
    name );

} /* End of usage() */


static void
nm_file_free( mcpu_nm_file *file )
{
  if( !file )
    return;
  free( file->shdr );
  memset( file, 0, sizeof( *file ) );

} /* End of nm_file_free() */


static void
nm_symbols_free( mcpu_nm_symbols *symbols )
{
  size_t i;

  if( !symbols )
    return;
  for( i = 0; i < symbols->count; ++i )
    free( symbols->symbol[i].name );
  free( symbols->symbol );
  memset( symbols, 0, sizeof( *symbols ) );

} /* End of nm_symbols_free() */


static void
get_shdr( const unsigned char *p, int data, mcpu_nm_shdr *sh )
{
  memset( sh, 0, sizeof( *sh ) );
  sh->type = cpu_melf_get32( p + MELF64_SHDR_SH_TYPE_OFF, data );
  sh->flags = cpu_melf_get64( p + MELF64_SHDR_SH_FLAGS_OFF, data );
  sh->offset = cpu_melf_get64( p + MELF64_SHDR_SH_OFFSET_OFF, data );
  sh->size = cpu_melf_get64( p + MELF64_SHDR_SH_SIZE_OFF, data );
  sh->link = cpu_melf_get32( p + MELF64_SHDR_SH_LINK_OFF, data );
  sh->entsize = cpu_melf_get64( p + MELF64_SHDR_SH_ENTSIZE_OFF, data );

} /* End of get_shdr() */


static int
nm_file_open_data( const unsigned char *data, size_t size, mcpu_nm_file *file,
                   char *err, size_t errsz )
{
  __mpu_uint64_t shoff;
  __mpu_uint16_t shentsize;
  size_t         i;

  memset( file, 0, sizeof( *file ) );
  file->data = data;
  file->size = size;

  if( size < MELF64_EHDR_SIZE )
  {
    snprintf( err, errsz, "file is too small for a MELF64 header" );
    return( -1 );
  }
  if( data[MELF_EI_MAG0] != MELFMAG0 || data[MELF_EI_MAG1] != MELFMAG1 ||
      data[MELF_EI_MAG2] != MELFMAG2 || data[MELF_EI_MAG3] != MELFMAG3 )
  {
    snprintf( err, errsz, "not a MELF file" );
    return( -1 );
  }
  if( data[MELF_EI_CLASS] != MELF_CLASS_64 )
  {
    snprintf( err, errsz, "unsupported MELF class %u",
              (unsigned int)data[MELF_EI_CLASS] );
    return( -1 );
  }

  file->encoding = data[MELF_EI_DATA];
  if( file->encoding != MELF_DATA_LSB && file->encoding != MELF_DATA_MSB )
  {
    snprintf( err, errsz, "invalid MELF data encoding %d", file->encoding );
    return( -1 );
  }
  if( cpu_melf_get16( data + MELF64_EHDR_E_EHSIZE_OFF, file->encoding ) !=
      MELF64_EHDR_SIZE )
  {
    snprintf( err, errsz, "unsupported MELF64 header size" );
    return( -1 );
  }

  shoff = cpu_melf_get64( data + MELF64_EHDR_E_SHOFF_OFF, file->encoding );
  shentsize = cpu_melf_get16( data + MELF64_EHDR_E_SHENTSIZE_OFF,
                              file->encoding );
  file->shnum = cpu_melf_get16( data + MELF64_EHDR_E_SHNUM_OFF,
                                file->encoding );

  if( !file->shnum )
    return( 0 );
  if( shentsize != MELF64_SHDR_SIZE ||
      !cpu_melf_span( size, shoff,
                      (__mpu_uint64_t)file->shnum * MELF64_SHDR_SIZE ) )
  {
    snprintf( err, errsz, "malformed section-header table" );
    return( -1 );
  }

  file->shdr = (mcpu_nm_shdr *)calloc( file->shnum, sizeof( *file->shdr ) );
  if( !file->shdr )
  {
    snprintf( err, errsz, "out of memory" );
    return( -1 );
  }

  for( i = 0; i < file->shnum; ++i )
    get_shdr( data + (size_t)shoff + i * MELF64_SHDR_SIZE,
              file->encoding, &file->shdr[i] );

  return( 0 );

} /* End of nm_file_open_data() */


static const char *
string_at( const unsigned char *table, size_t size, __mpu_uint64_t offset )
{
  const unsigned char *p;

  if( !table || offset >= size )
    return( NULL );
  p = table + (size_t)offset;
  if( !memchr( p, '\0', size - (size_t)offset ) )
    return( NULL );
  return( (const char *)p );

} /* End of string_at() */


static int
symbol_code( const mcpu_nm_file *file, unsigned int bind,
             __mpu_uint16_t shndx, char *code )
{
  char c;

  if( bind == MELF_STB_WEAK )
  {
    *code = shndx == MELF_SHN_UNDEF ? 'w' : 'W';
    return( 0 );
  }
  if( shndx == MELF_SHN_UNDEF )
  {
    *code = 'U';
    return( 0 );
  }
  if( shndx == MELF_SHN_ABS )
    c = 'A';
  else if( shndx == MELF_SHN_COMMON )
    c = 'C';
  else if( shndx >= file->shnum )
    c = '?';
  else if( file->shdr[shndx].flags & MELF_SHF_EXECINSTR )
    c = 'T';
  else if( file->shdr[shndx].type == MELF_SHT_NOBITS &&
           (file->shdr[shndx].flags & MELF_SHF_WRITE) )
    c = 'B';
  else if( file->shdr[shndx].flags & MELF_SHF_WRITE )
    c = 'D';
  else if( file->shdr[shndx].flags & MELF_SHF_ALLOC )
    c = 'R';
  else
    c = 'N';

  if( bind == MELF_STB_LOCAL && c >= 'A' && c <= 'Z' )
    c = (char)(c - 'A' + 'a');
  *code = c;
  return( c == '?' ? -1 : 0 );

} /* End of symbol_code() */


static int
symbols_append( mcpu_nm_symbols *symbols, const char *name,
                __mpu_uint64_t value, __mpu_uint64_t size,
                __mpu_uint16_t shndx, unsigned int bind, unsigned int type,
                char code, size_t order, char *err, size_t errsz )
{
  mcpu_nm_symbol *q;
  char           *copy;

  q = (mcpu_nm_symbol *)realloc( symbols->symbol,
                  (symbols->count + 1) * sizeof( *q ) );
  if( !q )
  {
    snprintf( err, errsz, "out of memory" );
    return( -1 );
  }
  copy = strdup( name );
  if( !copy )
  {
    snprintf( err, errsz, "out of memory" );
    return( -1 );
  }

  symbols->symbol = q;
  q = &symbols->symbol[symbols->count++];
  memset( q, 0, sizeof( *q ) );
  q->name = copy;
  q->value = value;
  q->size = size;
  q->shndx = shndx;
  q->bind = bind;
  q->type = type;
  q->code = code;
  q->order = order;
  return( 0 );

} /* End of symbols_append() */


static int
collect_table( const mcpu_nm_file *file, size_t table_index,
               const mcpu_nm_options *options, mcpu_nm_symbols *symbols,
               size_t *order, char *err, size_t errsz )
{
  const mcpu_nm_shdr  *sh = &file->shdr[table_index];
  const mcpu_nm_shdr  *str;
  const unsigned char *strtab;
  size_t               strsize;
  size_t               count;
  size_t               i;

  if( sh->entsize != MELF64_SYM_SIZE || sh->size % MELF64_SYM_SIZE ||
      sh->size > SIZE_MAX || !cpu_melf_span( file->size, sh->offset, sh->size ) )
  {
    snprintf( err, errsz, "malformed symbol table" );
    return( -1 );
  }
  if( sh->link >= file->shnum )
  {
    snprintf( err, errsz, "symbol table has invalid string-table link" );
    return( -1 );
  }
  str = &file->shdr[sh->link];
  if( str->type != MELF_SHT_STRTAB || str->size > SIZE_MAX ||
      !cpu_melf_span( file->size, str->offset, str->size ) )
  {
    snprintf( err, errsz, "malformed symbol string table" );
    return( -1 );
  }

  strtab = file->data + (size_t)str->offset;
  strsize = (size_t)str->size;
  count = (size_t)(sh->size / MELF64_SYM_SIZE);

  for( i = 0; i < count; ++i )
  {
    const unsigned char *p;
    __mpu_uint32_t       name_off;
    unsigned int         bind;
    unsigned int         type;
    __mpu_uint16_t       shndx;
    __mpu_uint64_t       value;
    __mpu_uint64_t       ssize;
    const char          *name;
    char                 code;

    p = file->data + (size_t)sh->offset + i * MELF64_SYM_SIZE;
    name_off = cpu_melf_get32( p + MELF64_SYM_ST_NAME_OFF, file->encoding );
    bind = MELF64_ST_BIND( p[MELF64_SYM_ST_INFO_OFF] );
    type = MELF64_ST_TYPE( p[MELF64_SYM_ST_INFO_OFF] );
    shndx = cpu_melf_get16( p + MELF64_SYM_ST_SHNDX_OFF, file->encoding );
    value = cpu_melf_get64( p + MELF64_SYM_ST_VALUE_OFF, file->encoding );
    ssize = cpu_melf_get64( p + MELF64_SYM_ST_SIZE_OFF, file->encoding );
    name = string_at( strtab, strsize, name_off );

    if( !name )
    {
      snprintf( err, errsz, "symbol has invalid name offset" );
      return( -1 );
    }
    if( !*name )
      continue;
    if( !options->all_symbols &&
        (type == MELF_STT_FILE || type == MELF_STT_SECTION) )
      continue;
    if( options->extern_only && bind == MELF_STB_LOCAL )
      continue;
    if( options->undefined_only && shndx != MELF_SHN_UNDEF )
      continue;
    if( symbol_code( file, bind, shndx, &code ) )
    {
      snprintf( err, errsz, "symbol '%s' has invalid section index", name );
      return( -1 );
    }

    if( symbols_append( symbols, name, value, ssize, shndx, bind, type,
                        code, (*order)++, err, errsz ) )
      return( -1 );
  }
  return( 0 );

} /* End of collect_table() */


static int
collect_symbols( const unsigned char *data, size_t size,
                 const mcpu_nm_options *options, mcpu_nm_symbols *symbols,
                 char *err, size_t errsz )
{
  mcpu_nm_file file;
  size_t       i;
  size_t       order = 0;
  int          want;

  memset( symbols, 0, sizeof( *symbols ) );
  if( nm_file_open_data( data, size, &file, err, errsz ) )
    return( -1 );

  want = options->dynamic ? MELF_SHT_DYNSYM : MELF_SHT_SYMTAB;
  for( i = 0; i < file.shnum; ++i )
  {
    if( file.shdr[i].type != (unsigned int)want )
      continue;
    if( collect_table( &file, i, options, symbols, &order, err, errsz ) )
    {
      nm_file_free( &file );
      nm_symbols_free( symbols );
      return( -1 );
    }
  }

  nm_file_free( &file );
  return( 0 );

} /* End of collect_symbols() */


static int
compare_symbols( const void *ap, const void *bp )
{
  const mcpu_nm_symbol *a = (const mcpu_nm_symbol *)ap;
  const mcpu_nm_symbol *b = (const mcpu_nm_symbol *)bp;
  int c;

  if( sort_options->numeric_sort )
  {
    if( a->value < b->value )
      c = -1;
    else if( a->value > b->value )
      c = 1;
    else
    {
      c = strcmp( a->name, b->name );
      if( !c )
        c = a->order < b->order ? -1 : a->order > b->order ? 1 : 0;
    }
  }
  else
  {
    c = strcmp( a->name, b->name );
    if( !c )
      c = a->order < b->order ? -1 : a->order > b->order ? 1 : 0;
  }

  return( sort_options->reverse_sort ? -c : c );

} /* End of compare_symbols() */


static void
print_one_symbol( const mcpu_nm_symbol *symbol, const char *prefix,
                  const mcpu_nm_options *options )
{
  if( options->print_file_name && prefix )
    printf( "%s: ", prefix );

  if( symbol->shndx == MELF_SHN_UNDEF )
    printf( "%16s %c %s\n", "", symbol->code, symbol->name );
  else
    printf( "%016" PRIx64 " %c %s\n",
            (uint64_t)symbol->value, symbol->code, symbol->name );

} /* End of print_one_symbol() */


static int
print_symbols_data( const unsigned char *data, size_t size,
                    const char *header, const char *prefix,
                    const mcpu_nm_options *options, int show_header,
                    char *err, size_t errsz )
{
  mcpu_nm_symbols symbols;
  size_t          i;

  if( collect_symbols( data, size, options, &symbols, err, errsz ) )
    return( -1 );

  if( !options->no_sort && symbols.count > 1 )
  {
    sort_options = options;
    qsort( symbols.symbol, symbols.count, sizeof( *symbols.symbol ),
           compare_symbols );
  }

  if( show_header && !options->print_file_name )
    printf( "%s:\n", header );
  for( i = 0; i < symbols.count; ++i )
    print_one_symbol( &symbols.symbol[i], prefix, options );

  nm_symbols_free( &symbols );
  return( 0 );

} /* End of print_symbols_data() */


static int
is_archive_file( const unsigned char *data, size_t size )
{
  static const unsigned char magic[] = "!<arch>\n";

  return( size >= sizeof( magic ) - 1 &&
          !memcmp( data, magic, sizeof( magic ) - 1 ) );

} /* End of is_archive_file() */


static int
process_archive( const char *path, const mcpu_nm_options *options,
                 char *err, size_t errsz )
{
  mcpu_archive archive;
  size_t       i;

  if( mcpu_archive_read( path, &archive, err, errsz ) )
    return( -1 );

  for( i = 0; i < archive.member_count; ++i )
  {
    const mcpu_archive_member *member = &archive.members[i];
    char                       label[1024];

    if( snprintf( label, sizeof( label ), "%s:%s", path, member->name ) >=
        (int)sizeof( label ) )
    {
      snprintf( err, errsz, "archive/member name is too long" );
      mcpu_archive_free( &archive );
      return( -1 );
    }

    if( print_symbols_data( member->data, member->size, label, label,
                            options, 1, err, errsz ) )
    {
      mcpu_archive_free( &archive );
      return( -1 );
    }
    if( !options->print_file_name && i + 1 < archive.member_count )
      printf( "\n" );
  }

  mcpu_archive_free( &archive );
  return( 0 );

} /* End of process_archive() */


static int
process_path( const char *path, const mcpu_nm_options *options,
              int show_header, char *err, size_t errsz )
{
  unsigned char *data = NULL;
  size_t         size = 0;
  int            rc;

  if( cpu_melf_readall( path, &data, &size, err, errsz ) )
    return( -1 );

  if( is_archive_file( data, size ) )
  {
    free( data );
    return( process_archive( path, options, err, errsz ) );
  }

  rc = print_symbols_data( data, size, path, path, options,
                           show_header, err, errsz );
  free( data );
  return( rc );

} /* End of process_path() */


static int
parse_short_options( const char *arg, mcpu_nm_options *options )
{
  const char *p = arg + 1;

  while( *p )
  {
    switch( *p )
    {
      case 'a': options->all_symbols = 1; break;
      case 'g': options->extern_only = 1; break;
      case 'u': options->undefined_only = 1; break;
      case 'n': options->numeric_sort = 1; options->no_sort = 0; break;
      case 'r': options->reverse_sort = 1; break;
      case 'p': options->no_sort = 1; options->numeric_sort = 0; break;
      case 'D': options->dynamic = 1; break;
      case 'A': options->print_file_name = 1; break;
      default: return( -1 );
    }
    ++p;
  }
  return( 0 );

} /* End of parse_short_options() */


int
main( int argc, char **argv )
{
  mcpu_nm_options options;
  char          **files = NULL;
  size_t          file_count = 0;
  size_t          file_capacity = 0;
  int             i;
  int             rc = 0;

  memset( &options, 0, sizeof( options ) );

  for( i = 1; i < argc; ++i )
  {
    const char *arg = argv[i];

    if( !strcmp( arg, "--help" ) )
    {
      usage( stdout, argv[0] );
      free( files );
      return( 0 );
    }
    else if( !strcmp( arg, "--version" ) ||
             !strcmp( arg, "--dump-version" ) )
    {
      mcpu_version_action version_action;
      int rc;

      mcpu_version_parse_option( arg, &version_action );
      rc = mcpu_version_print( argv[0], version_action );
      free( files );
      return( rc == 0 ? 0 : 1 );
    }
    else if( !strcmp( arg, "--debug-syms" ) )
      options.all_symbols = 1;
    else if( !strcmp( arg, "--extern-only" ) )
      options.extern_only = 1;
    else if( !strcmp( arg, "--undefined-only" ) )
      options.undefined_only = 1;
    else if( !strcmp( arg, "--numeric-sort" ) )
    {
      options.numeric_sort = 1;
      options.no_sort = 0;
    }
    else if( !strcmp( arg, "--reverse-sort" ) )
      options.reverse_sort = 1;
    else if( !strcmp( arg, "--no-sort" ) )
    {
      options.no_sort = 1;
      options.numeric_sort = 0;
    }
    else if( !strcmp( arg, "--dynamic" ) )
      options.dynamic = 1;
    else if( !strcmp( arg, "--print-file-name" ) )
      options.print_file_name = 1;
    else if( arg[0] == '-' && arg[1] != '\0' )
    {
      if( parse_short_options( arg, &options ) )
      {
        fprintf( stderr, "mcpu-nm: unknown option '%s'\n", arg );
        usage( stderr, argv[0] );
        free( files );
        return( 1 );
      }
    }
    else
    {
      if( file_count == file_capacity )
      {
        size_t  next = file_capacity ? file_capacity * 2U : 4U;
        char  **q = (char **)realloc( files, next * sizeof( *q ) );

        if( !q )
        {
          fprintf( stderr, "mcpu-nm: out of memory\n" );
          free( files );
          return( 1 );
        }
        files = q;
        file_capacity = next;
      }
      files[file_count++] = argv[i];
    }
  }

  if( !file_count )
  {
    usage( stderr, argv[0] );
    free( files );
    return( 1 );
  }

  for( i = 0; i < (int)file_count; ++i )
  {
    char err[512];

    if( i && !options.print_file_name )
      printf( "\n" );
    if( process_path( files[i], &options,
                      file_count > 1, err, sizeof( err ) ) )
    {
      fprintf( stderr, "mcpu-nm: %s: %s\n", files[i], err );
      rc = 1;
    }
  }

  free( files );
  return( rc );

} /* End of main() */