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

#include <errno.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>

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


typedef struct cli_action
{
  mcpu_object_edit_action edit;
  char                    *name;
  char                    *new_name;
  unsigned char           *data;
} cli_action;


typedef struct dump_action
{
  char *name;
  char *path;
} dump_action;


static void
usage( FILE *fp, const char *name )
{
  fprintf( fp,
    "usage: %s [options] INPUT OUTPUT\n\n"
    "Copy and safely edit LibMCPU MELF files and deterministic archives.\n"
    "The output format is always MELF.  Runtime-visible sections and frozen\n"
    "MELF ABI3 execution metadata cannot be removed, renamed or updated.\n\n"
    "Options:\n"
    "  -j, --only-section=NAME        - copy only the named sections\n"
    "  -R, --remove-section=NAME      - remove a safe non-runtime section\n"
    "      --dump-section=NAME=FILE   - write section contents to FILE\n"
    "      --add-section=NAME=FILE    - add non-ALLOC SHT_PROGBITS metadata\n"
    "      --update-section=NAME=FILE - replace safe non-ALLOC section contents\n"
    "      --rename-section=OLD=NEW   - rename a safe non-ALLOC section\n"
    "      --help                     - show this help\n"
    "      --version                  - show program version\n"
    "      --dump-version             - show version number only\n",
    name );

} /* End of usage() */


static int
append_string( char ***items, size_t *count, const char *s )
{
  char **q;
  char *p;

  if( !s || !*s )
    return( -1 );
  q = (char **)realloc( *items, (*count + 1) * sizeof( *q ) );
  if( !q )
    return( -1 );
  p = strdup( s );
  if( !p )
    return( -1 );
  *items = q;
  (*items)[(*count)++] = p;
  return( 0 );

} /* End of append_string() */


static void
free_strings( char **items, size_t count )
{
  size_t i;

  for( i = 0; i < count; ++i )
    free( items[i] );
  free( items );

} /* End of free_strings() */


static int
split_pair( const char *arg, char **left, char **right )
{
  const char *eq = strchr( arg, '=' );
  size_t n;

  if( !eq || eq == arg || !eq[1] )
    return( -1 );
  n = (size_t)( eq - arg );
  *left = (char *)malloc( n + 1 );
  if( !*left )
    return( -1 );
  memcpy( *left, arg, n );
  (*left)[n] = 0;
  *right = strdup( eq + 1 );
  if( !*right )
  {
    free( *left );
    *left = NULL;
    return( -1 );
  }
  return( 0 );

} /* End of split_pair() */


static int
append_edit_action( cli_action **actions, size_t *count,
                    mcpu_object_edit_action_type type,
                    const char *spec, int needs_file,
                    char *err, size_t errsz )
{
  cli_action *q;
  cli_action *a;
  char *left = NULL;
  char *right = NULL;

  if( split_pair( spec, &left, &right ) )
  {
    snprintf( err, errsz, "expected NAME=VALUE: %s", spec );
    return( -1 );
  }
  q = (cli_action *)realloc( *actions, (*count + 1) * sizeof( *q ) );
  if( !q )
  {
    free( left );
    free( right );
    snprintf( err, errsz, "out of memory" );
    return( -1 );
  }
  *actions = q;
  a = &q[*count];
  memset( a, 0, sizeof( *a ) );
  a->name = left;
  a->edit.type = type;
  a->edit.name = a->name;

  if( needs_file )
  {
    if( cpu_melf_readall( right, &a->data, &a->edit.size, err, errsz ) )
    {
      free( right );
      free( a->name );
      memset( a, 0, sizeof( *a ) );
      return( -1 );
    }
    a->edit.data = a->data;
    free( right );
  }
  else
  {
    a->new_name = right;
    a->edit.new_name = a->new_name;
  }
  ++*count;
  return( 0 );

} /* End of append_edit_action() */


static int
append_remove_action( cli_action **actions, size_t *count,
                      const char *name, char *err, size_t errsz )
{
  cli_action *q;
  cli_action *a;

  if( !name || !*name )
  {
    snprintf( err, errsz, "empty section name" );
    return( -1 );
  }
  q = (cli_action *)realloc( *actions, (*count + 1) * sizeof( *q ) );
  if( !q )
  {
    snprintf( err, errsz, "out of memory" );
    return( -1 );
  }
  *actions = q;
  a = &q[*count];
  memset( a, 0, sizeof( *a ) );
  a->name = strdup( name );
  if( !a->name )
  {
    snprintf( err, errsz, "out of memory" );
    return( -1 );
  }
  a->edit.type = MCPU_OBJECT_EDIT_REMOVE;
  a->edit.name = a->name;
  ++*count;
  return( 0 );

} /* End of append_remove_action() */


static void
free_edit_actions( cli_action *actions, size_t count )
{
  size_t i;

  for( i = 0; i < count; ++i )
  {
    free( actions[i].name );
    free( actions[i].new_name );
    free( actions[i].data );
  }
  free( actions );

} /* End of free_edit_actions() */


static int
append_dump_action( dump_action **actions, size_t *count,
                    const char *spec, char *err, size_t errsz )
{
  dump_action *q;
  char *name = NULL;
  char *path = NULL;

  if( split_pair( spec, &name, &path ) )
  {
    snprintf( err, errsz, "expected NAME=FILE: %s", spec );
    return( -1 );
  }
  q = (dump_action *)realloc( *actions, (*count + 1) * sizeof( *q ) );
  if( !q )
  {
    free( name );
    free( path );
    snprintf( err, errsz, "out of memory" );
    return( -1 );
  }
  *actions = q;
  q[*count].name = name;
  q[*count].path = path;
  ++*count;
  return( 0 );

} /* End of append_dump_action() */


static void
free_dump_actions( dump_action *actions, size_t count )
{
  size_t i;

  for( i = 0; i < count; ++i )
  {
    free( actions[i].name );
    free( actions[i].path );
  }
  free( actions );

} /* End of free_dump_actions() */


static int
write_output( const char *input, const char *output,
              const unsigned char *data, size_t size,
              char *err, size_t errsz )
{
  struct stat st;

  if( stat( input, &st ) )
  {
    snprintf( err, errsz, "%s: %s", input, strerror( errno ) );
    return( -1 );
  }
  if( cpu_melf_writeall( output, data, size, err, errsz ) )
    return( -1 );
  if( chmod( output, st.st_mode & 07777 ) )
  {
    snprintf( err, errsz, "%s: %s", output, strerror( errno ) );
    unlink( output );
    return( -1 );
  }
  return( 0 );

} /* End of write_output() */


static int
perform_dumps( const unsigned char *data, size_t size,
               const dump_action *actions, size_t count,
               char *err, size_t errsz )
{
  size_t i;

  for( i = 0; i < count; ++i )
  {
    unsigned char *p = NULL;
    size_t n = 0;

    if( mcpu_object_dump_section( data, size, actions[i].name,
                                  &p, &n, err, errsz ) )
      return( -1 );
    if( cpu_melf_writeall( actions[i].path, p, n, err, errsz ) )
    {
      free( p );
      return( -1 );
    }
    free( p );
  }
  return( 0 );

} /* End of perform_dumps() */


static int
copy_archive( const char *input, const char *output,
              const mcpu_object_edit_options *options,
              char *err, size_t errsz )
{
  mcpu_archive archive;
  struct stat  st;
  size_t       i;
  int          rc = -1;

  mcpu_archive_init( &archive );
  if( stat( input, &st ) )
  {
    snprintf( err, errsz, "%s: %s", input, strerror( errno ) );
    return( -1 );
  }
  if( mcpu_archive_read( input, &archive, err, errsz ) )
    return( -1 );

  for( i = 0; i < archive.member_count; ++i )
  {
    unsigned char *p = NULL;
    size_t n = 0;
    int changed = 0;

    if( mcpu_object_edit_data( archive.members[i].data,
                               archive.members[i].size,
                               options, &p, &n, &changed, err, errsz ) )
      goto out;
    free( archive.members[i].data );
    archive.members[i].data = p;
    archive.members[i].size = n;
    (void)changed;
  }

  if( mcpu_archive_write( output, &archive, archive.had_symbol_index,
                          err, errsz ) )
    goto out;
  if( chmod( output, st.st_mode & 07777 ) )
  {
    snprintf( err, errsz, "%s: %s", output, strerror( errno ) );
    unlink( output );
    goto out;
  }
  rc = 0;

out:
  mcpu_archive_free( &archive );
  return( rc );

} /* End of copy_archive() */


int
main( int argc, char **argv )
{
  const char    *prog = argv[0];
  const char    *input = NULL;
  const char    *output = NULL;
  char         **only = NULL;
  size_t         only_count = 0;
  cli_action    *actions = NULL;
  size_t         action_count = 0;
  dump_action   *dumps = NULL;
  size_t         dump_count = 0;
  unsigned char *data = NULL;
  size_t         size = 0;
  mcpu_object_edit_options options;
  mcpu_object_edit_action *edit_actions = NULL;
  char           err[512];
  int            i;
  int            rc = EXIT_FAILURE;

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

    if( !strcmp( arg, "--help" ) )
    {
      usage( stdout, prog );
      rc = EXIT_SUCCESS;
      goto out;
    }
    if( !strcmp( arg, "--version" ) ||
        !strcmp( arg, "--dump-version" ) )
    {
      mcpu_version_action version_action;

      mcpu_version_parse_option( arg, &version_action );
      rc = mcpu_version_print(prog, version_action) == 0 ?
           EXIT_SUCCESS : EXIT_FAILURE;
      goto out;
    }
    if( !strcmp( arg, "-j" ) )
    {
      if( ++i >= argc || append_string( &only, &only_count, argv[i] ) )
      {
        snprintf( err, sizeof( err ), "option -j requires a section name" );
        goto bad_usage;
      }
      continue;
    }
    if( !strncmp( arg, "-j", 2 ) && arg[2] )
    {
      if( append_string( &only, &only_count, arg + 2 ) )
      {
        snprintf( err, sizeof( err ), "out of memory or empty section name" );
        goto bad_usage;
      }
      continue;
    }
    if( !strncmp( arg, "--only-section=", 15 ) )
    {
      if( append_string( &only, &only_count, arg + 15 ) )
      {
        snprintf( err, sizeof( err ), "--only-section requires a section name" );
        goto bad_usage;
      }
      continue;
    }
    if( !strcmp( arg, "-R" ) )
    {
      if( ++i >= argc || append_remove_action( &actions, &action_count,
                                                argv[i], err, sizeof( err ) ) )
        goto bad_usage;
      continue;
    }
    if( !strncmp( arg, "-R", 2 ) && arg[2] )
    {
      if( append_remove_action( &actions, &action_count, arg + 2,
                                err, sizeof( err ) ) )
        goto bad_usage;
      continue;
    }
    if( !strncmp( arg, "--remove-section=", 17 ) )
    {
      if( append_remove_action( &actions, &action_count, arg + 17,
                                err, sizeof( err ) ) )
        goto bad_usage;
      continue;
    }
    if( !strncmp( arg, "--dump-section=", 15 ) )
    {
      if( append_dump_action( &dumps, &dump_count, arg + 15,
                              err, sizeof( err ) ) )
        goto bad_usage;
      continue;
    }
    if( !strncmp( arg, "--add-section=", 14 ) )
    {
      if( append_edit_action( &actions, &action_count, MCPU_OBJECT_EDIT_ADD,
                              arg + 14, 1, err, sizeof( err ) ) )
        goto bad_usage;
      continue;
    }
    if( !strncmp( arg, "--update-section=", 17 ) )
    {
      if( append_edit_action( &actions, &action_count, MCPU_OBJECT_EDIT_UPDATE,
                              arg + 17, 1, err, sizeof( err ) ) )
        goto bad_usage;
      continue;
    }
    if( !strncmp( arg, "--rename-section=", 17 ) )
    {
      if( append_edit_action( &actions, &action_count, MCPU_OBJECT_EDIT_RENAME,
                              arg + 17, 0, err, sizeof( err ) ) )
        goto bad_usage;
      continue;
    }
    if( arg[0] == '-' && arg[1] )
    {
      snprintf( err, sizeof( err ), "unknown option: %s", arg );
      goto bad_usage;
    }
    if( !input )
      input = arg;
    else if( !output )
      output = arg;
    else
    {
      snprintf( err, sizeof( err ), "exactly INPUT and OUTPUT are required" );
      goto bad_usage;
    }
  }

  if( !input || !output )
  {
    snprintf( err, sizeof( err ), "both INPUT and OUTPUT are required" );
    goto bad_usage;
  }
  if( !strcmp( input, output ) )
  {
    snprintf( err, sizeof( err ), "INPUT and OUTPUT must be different files" );
    goto fail;
  }

  if( action_count )
  {
    size_t k;

    edit_actions = (mcpu_object_edit_action *)calloc( action_count,
                                                       sizeof( *edit_actions ) );
    if( !edit_actions )
    {
      snprintf( err, sizeof( err ), "out of memory" );
      goto fail;
    }
    for( k = 0; k < action_count; ++k )
      edit_actions[k] = actions[k].edit;
  }
  memset( &options, 0, sizeof( options ) );
  options.only_sections = (const char *const *)only;
  options.only_count = only_count;
  options.actions = edit_actions;
  options.action_count = action_count;

  if( cpu_melf_readall( input, &data, &size, err, sizeof( err ) ) )
    goto fail;

  if( size >= 8 && !memcmp( data, "!<arch>\n", 8 ) )
  {
    if( dump_count )
    {
      snprintf( err, sizeof( err ),
                "--dump-section is ambiguous for archive input" );
      goto fail;
    }
    if( !only_count && !action_count )
    {
      if( write_output( input, output, data, size, err, sizeof( err ) ) )
        goto fail;
    }
    else if( copy_archive( input, output, &options, err, sizeof( err ) ) )
      goto fail;
  }
  else
  {
    unsigned char *result = NULL;
    size_t result_size = 0;
    int changed = 0;

    if( perform_dumps( data, size, dumps, dump_count, err, sizeof( err ) ) )
      goto fail;
    if( mcpu_object_edit_data( data, size, &options,
                               &result, &result_size, &changed,
                               err, sizeof( err ) ) )
      goto fail;
    if( write_output( input, output, result, result_size,
                      err, sizeof( err ) ) )
    {
      free( result );
      goto fail;
    }
    free( result );
    (void)changed;
  }

  rc = EXIT_SUCCESS;
  goto out;

bad_usage:
  fprintf( stderr, "%s: %s\n", prog, err );
  usage( stderr, prog );
  goto out;

fail:
  fprintf( stderr, "%s: %s\n", prog, err );

out:
  free( data );
  free( edit_actions );
  free_dump_actions( dumps, dump_count );
  free_edit_actions( actions, action_count );
  free_strings( only, only_count );
  return( rc );

} /* End of main() */