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

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

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

#include "mcpu-asm.h"
#include "mcpu-object.h"
#include "mcpu-link.h"
#include "mcpu-exec.h"
#include "mcpu-melf.h"

#define CHECK(COND, MSG) \
  do { if(!(COND)) { fprintf(stderr, "FAIL: %s\n", (MSG)); goto out; } } while(0)

static int
find_symbol( const cpu_asm_object *o, const char *name, size_t *idx )
{
  size_t i;
  for( i = 0; i < o->symbol_count; ++i )
    if( !strcmp( o->symbols[i].name, name ) )
    {
      if( idx ) * idx = i;
      return( 0 );
    }
  return( -1 );
}

static int
write_copy( const char *path, const unsigned char *p, size_t n )
{
  FILE *f = fopen( path, "wb" );
  if( !f ) return( -1 );
  if( n && fwrite( p, 1, n, f ) != n )
  {
    fclose( f );
    return( -1 );
  }
  return( fclose( f ) ? -1 : 0 );
}

static int
expect_load_fail( const char *path, const char *needle )
{
  cpu_exec_loaded ld;
  char err[512];
  char *av[] = {(char *)"prog"};

  if( !cpu_exec_load( path, 1, av, &ld, err, sizeof( err ) ) )
  {
    fprintf( stderr, "malformed executable unexpectedly loaded: %s\n", path );
    return( -1 );
  }
  if( needle && !strstr( err, needle ) )
  {
    fprintf( stderr, "load error mismatch: expected '%s', got '%s'\n", needle, err );
    return( -1 );
  }
  return( 0 );
}

int
main( void )
{
  static const char main_src[] =
    ".extern func\n"
    ".extern shared\n"
    ".extern scratch\n"
    ".text\n"
    "_start_life:\n"
    "  LDI r64@0x00000, 40\n"
    "  CALL func\n"                 /* CODE64 */
    "  STORE shared, r64@0x00000\n" /* DATA64 */
    "  LOAD r64@0x00018, shared\n"  /* DATA64 */
    "  LEA r64@0x00020, scratch\n"  /* PTR64_CODE */
    "  LDI r32@0x00030, 0\n"
    "  EXIT r32@0x00030\n";

  static const char lib_src[] =
    ".global func, shared, scratch, shared_ptr, func_ptr\n"
    ".data\n"
    ".align 32\n"
    "shared:\n"
    "  .integer 64, 0\n"
    "shared_ptr:\n"
    "  .ptr scratch\n"             /* PTR64_DATA */
    "func_ptr:\n"
    "  .cptr func\n"               /* CPTR64_DATA */
    ".bss\n"
    ".align 64\n"
    "scratch:\n"
    "  .zero 128\n"
    ".text\n"
    "func:\n"
    "  LDI r64@0x00008, 2\n"
    "  IADD r64@0x00000, r64@0x00000, r64@0x00008\n"
    "  RET\n";

  char dtemplate[] = "/tmp/libcpu-melf-XXXXXX";
  char main_o[512], lib_o[512], exe[512], life_exe[512], bad_guard[512], bad_entry[512];
  char err[512];
  cpu_asm_object built[2], readback[2];
  cpu_asm_image im;
  cpu_exec_options opt;
  cpu_exec_loaded ld;
  unsigned char *file = 0;
  size_t file_size = 0;
  size_t ix, i;
  __mpu_uint64_t p;
  int data;
  int rc = 1;

  memset( built, 0, sizeof( built ) );
  memset( readback, 0, sizeof( readback ) );
  memset( &im, 0, sizeof( im ) );

  CHECK( mkdtemp( dtemplate ) != NULL, "mkdtemp" );
  snprintf( main_o, sizeof( main_o ), "%s/main.o", dtemplate );
  snprintf( lib_o, sizeof( lib_o ), "%s/lib.o", dtemplate );
  snprintf( exe, sizeof( exe ), "%s/a.melf", dtemplate );
  snprintf( life_exe, sizeof( life_exe ), "%s/lifecycle.melf", dtemplate );
  snprintf( bad_guard, sizeof( bad_guard ), "%s/bad-guard.melf", dtemplate );
  snprintf( bad_entry, sizeof( bad_entry ), "%s/bad-entry.melf", dtemplate );

  __cpu_init();

  /*
   * LibMPU uses MPU_BYTE_ORDER / MPU_WORD_ORDER as numeric tags (for
   * example 1234), while the *_LITTLE_ENDIAN / *_BIG_ENDIAN macros are
   * boolean selectors.  MELF must key off the boolean selectors rather
   * than compare the tag with 0/1.
   */
#if CPU_DATA_BYTE_ORDER_LITTLE_ENDIAN == 1
  CHECK( cpu_melf_data_encoding() == MELF_DATA_LSB,
        "MELF EI_DATA follows LibMPU little-endian selector");
#elif CPU_DATA_BYTE_ORDER_BIG_ENDIAN == 1
  CHECK( cpu_melf_data_encoding() == MELF_DATA_MSB,
        "MELF EI_DATA follows LibMPU big-endian selector");
#else
# error Unsupported LibMPU byte order for MELF regression
#endif
  {
    MElf_McpuTarget mt;
    cpu_melf_fill_target( &mt );
#if CPU_DATA_WORD_ORDER_LITTLE_ENDIAN == 1
    CHECK( mt.mt_data_word_order == MELF_MCPU_WORD_ORDER_LSB,
          "MELF target metadata follows LibMPU little word-order selector");
#elif CPU_DATA_WORD_ORDER_BIG_ENDIAN == 1
    CHECK( mt.mt_data_word_order == MELF_MCPU_WORD_ORDER_MSB,
          "MELF target metadata follows LibMPU big word-order selector");
#else
# error Unsupported LibMPU word order for MELF regression
#endif
  }

  CHECK( cpu_assemble_object( main_src, &built[0], err, sizeof( err ) ) == 0, err );
  CHECK( cpu_assemble_object( lib_src, &built[1], err, sizeof( err ) ) == 0, err );
  CHECK( built[1].has_bss && built[1].bss_size == 128 && built[1].bss_align == 64,
        "assembler .bss/NOBITS metadata");
  CHECK( built[0].reloc_count == 4 && built[1].reloc_count == 2,
        "all expected relocations emitted");

  /* Assembler-only lifecycle pseudo-label rules. */
  {
    cpu_asm_object q;
    memset( &q, 0, sizeof( q ) );
    CHECK( cpu_assemble_object(
            ".init\n_start:\n  END\n.text\n_start_life:\n  END\n",
            &q, err, sizeof( err )) != 0 && strstr( err, "_start_life" ),
          "_start_life is rejected when .init exists");
    cpu_asm_object_free( &q );
    memset( &q, 0, sizeof( q ) );
    CHECK( cpu_assemble_object(
            ".text\n_start:\n_start_life:\n  END\n",
            &q, err, sizeof( err )) != 0 && strstr( err, "conflicts" ),
          "_start_life conflicts with explicit _start/_life");
    cpu_asm_object_free( &q );
    memset( &q, 0, sizeof( q ) );
    CHECK( cpu_assemble_object(
            ".weak _start\n.text\n_start_life:\n  END\n",
            &q, err, sizeof( err )) != 0 && strstr( err, "cannot be weak" ),
          "lifecycle symbols cannot be weak");
    cpu_asm_object_free( &q );
  }

  /* Full .init/.text/.fini ET_EXEC lifecycle metadata and loader path. */
  {
    static const char life_src[] =
      ".init\n"
      "_start:\n"
      "  END\n"
      ".text\n"
      "_life:\n"
      "  LDI r32@0x00000, 0\n"
      "  EXIT r32@0x00000\n"
      ".fini\n"
      "_fini:\n"
      "  END\n";
    cpu_asm_object q;
    cpu_asm_image qi;
    cpu_exec_loaded ql;
    cpu_exec_options qo = {512, 0, 0, 0};
    char *qv[] = {(char *)"lifecycle.melf"};
    memset( &q, 0, sizeof( q ) );
    memset( &qi, 0, sizeof( qi ) );
    memset( &ql, 0, sizeof( ql ) );
    CHECK( cpu_assemble_object( life_src, &q, err, sizeof( err ) ) == 0, err );
    CHECK( cpu_link_objects( &q, 1, &qi, err, sizeof( err ) ) == 0, err );
    CHECK( qi.has_start && qi.has_life && qi.has_init && qi.has_fini &&
          qi.start_entry == qi.init_entry && qi.start_entry != qi.life_entry,
          "linked INIT/LIFE/FINI lifecycle");
    CHECK( cpu_exec_write( life_exe, &qi, &qo, err, sizeof( err ) ) == 0, err );
    CHECK( cpu_exec_load( life_exe, 1, qv, &ql, err, sizeof( err ) ) == 0, err );
    CHECK( ql.has_init && ql.has_fini && ql.start_entry == ql.init_entry &&
          ql.start_entry != ql.life_entry,
          "ET_EXEC INIT/LIFE/FINI loader metadata");
    __cpu_free_context();
    __cpu_init();
    cpu_asm_image_free( &qi );
    cpu_asm_object_free( &q );
  }

  CHECK( find_symbol( &built[0], "_start", &ix ) == 0,
        "_start_life emits _start");
  CHECK( built[0].symbols[ix].binding == CPU_OBJ_BIND_GLOBAL &&
        built[0].symbols[ix].type == CPU_OBJ_TYPE_FUNC,
        "_start ABI symbol attributes");
  {
    __mpu_uint64_t sv = built[0].symbols[ix].value;
    CHECK( find_symbol( &built[0], "_life", &ix ) == 0,
          "_start_life emits _life");
    CHECK( built[0].symbols[ix].value == sv &&
          built[0].symbols[ix].binding == CPU_OBJ_BIND_GLOBAL &&
          built[0].symbols[ix].type == CPU_OBJ_TYPE_FUNC,
          "_start/_life have same GLOBAL/FUNC address");
  }

  CHECK( cpu_object_write( main_o, &built[0], err, sizeof( err ) ) == 0, err );
  CHECK( cpu_object_write( lib_o, &built[1], err, sizeof( err ) ) == 0, err );
  CHECK( cpu_object_read( main_o, &readback[0], err, sizeof( err ) ) == 0, err );
  CHECK( cpu_object_read( lib_o, &readback[1], err, sizeof( err ) ) == 0, err );
  CHECK( readback[1].has_bss && readback[1].bss_size == 128 &&
        readback[1].bss_align == 64, "ET_REL .bss round trip");
  CHECK( readback[0].reloc_count == 4 && readback[1].reloc_count == 2,
        "ET_REL RELA round trip");

  /* Undefined lifecycle declarations retain STT_FUNC in ET_REL. */
  {
    cpu_asm_object q;
    memset( &q, 0, sizeof( q ) );
    CHECK( cpu_assemble_object(
            ".extern _fini\n.text\n_start_life:\n  END\n",
            &q, err, sizeof(err)) == 0, err);
    CHECK( find_symbol( &q, "_fini", &ix ) == 0 &&
          q.symbols[ix].section == CPU_OBJ_SEC_UNDEF &&
          q.symbols[ix].binding == CPU_OBJ_BIND_GLOBAL &&
          q.symbols[ix].type == CPU_OBJ_TYPE_FUNC,
          "undefined lifecycle symbol is GLOBAL/FUNC");
    cpu_asm_object_free( &q );
  }

  /* WEAK resolution rules fixed by the symbol/RELA ABI. */
  {
    static const char weak_ptr_src[] =
      ".weak missing\n"
      ".data\nwp:\n  .ptr missing\n"
      ".text\n_start_life:\n  END\n";
    cpu_asm_object q;
    cpu_asm_image qi;
    memset( &q, 0, sizeof( q ) );
    memset( &qi, 0, sizeof( qi ) );
    CHECK( cpu_assemble_object( weak_ptr_src, &q, err, sizeof( err ) ) == 0, err );
    CHECK( cpu_link_objects( &q, 1, &qi, err, sizeof( err ) ) == 0, err );
    CHECK( qi.data_size >= CPU_DATA_NULL_GUARD_SIZE + 8 &&
          __cpu_data_get_u64( qi.data + CPU_DATA_NULL_GUARD_SIZE ) == 0,
          "unresolved weak pointer relocation becomes NULL");
    cpu_asm_image_free( &qi );
    cpu_asm_object_free( &q );
  }
  {
    static const char weak_code_src[] =
      ".weak missing\n"
      ".text\n_start_life:\n  CALL missing\n  END\n";
    cpu_asm_object q;
    cpu_asm_image qi;
    memset( &q, 0, sizeof( q ) );
    memset( &qi, 0, sizeof( qi ) );
    CHECK( cpu_assemble_object( weak_code_src, &q, err, sizeof( err ) ) == 0, err );
    CHECK( cpu_link_objects( &q, 1, &qi, err, sizeof( err ) ) != 0 &&
          strstr( err, "unresolved weak" ),
          "unresolved weak CODE64 is a link error");
    cpu_asm_image_free( &qi );
    cpu_asm_object_free( &q );
  }
  {
    static const char caller_src[] =
      ".extern foo\n"
      ".text\n_start_life:\n  CALL foo\n  END\n";
    static const char weak_def_src[] =
      ".weak foo\n.text\nfoo:\n  NOP\n  RET\n";
    static const char strong_def_src[] =
      ".global foo\n.text\nfoo:\n  RET\n";
    cpu_asm_object q[3];
    cpu_asm_image qi;
    __mpu_uint64_t expected_strong, encoded;
    memset( q, 0, sizeof( q ) );
    memset( &qi, 0, sizeof( qi ) );
    CHECK( cpu_assemble_object( caller_src, &q[0], err, sizeof( err ) ) == 0, err );
    CHECK( cpu_assemble_object( weak_def_src, &q[1], err, sizeof( err ) ) == 0, err );
    CHECK( cpu_assemble_object( strong_def_src, &q[2], err, sizeof( err ) ) == 0, err );
    expected_strong = q[0].text_size + q[1].text_size;
    CHECK( cpu_link_objects( q, 3, &qi, err, sizeof( err ) ) == 0, err );
    CHECK( q[0].reloc_count == 1 && q[0].relocs[0].type == CPU_RELOC_CODE64,
          "strong/weak test CALL relocation");
    encoded = __cpu_code_get_u64( qi.code + q[0].relocs[0].offset );
    CHECK( encoded == expected_strong, "strong GLOBAL overrides earlier WEAK definition" );
    cpu_asm_image_free( &qi );
    for( i = 0; i < 3; ++i ) cpu_asm_object_free( &q[i] );
  }

  CHECK( cpu_link_objects( readback, 2, &im, err, sizeof( err ) ) == 0, err );
  CHECK( im.has_start && im.has_life && !im.has_init &&
        im.start_entry == im.life_entry, "linked no-INIT lifecycle");
  CHECK( im.data_size == 56 && im.data_mem_size == 192,
        "DATA filesz/memsz including .bss");
  CHECK( im.bss_present && im.bss_offset == 64 && im.bss_size == 128,
        "final .bss placement");
  for( i = 0; i < CPU_DATA_NULL_GUARD_SIZE; ++i )
    CHECK( im.data[i] == 0, "linker-owned DATA NULL guard" );

  p = __cpu_data_get_u64( im.data + 40 );
  CHECK( __cpu_ptr_segment( p ) == CPU_PTR_SEG_DATA &&
        __cpu_ptr_offset( p ) == 64, "PTR64_DATA relocation to .bss");
  p = __cpu_data_get_u64( im.data + 48 );
  CHECK( __cpu_ptr_segment( p ) == CPU_PTR_SEG_PROGRAM,
        "CPTR64_DATA relocation");

  opt.stack_size = 512;
  opt.have_heap = 1;
  opt.heap_size = 128;
  opt.heap_limit = 1024;
  CHECK( cpu_exec_write( exe, &im, &opt, err, sizeof( err ) ) == 0, err );

  {
    char *av[] = {(char *)"a.melf", (char *)"alpha", (char *)"beta"};
    CHECK( cpu_exec_load( exe, 3, av, &ld, err, sizeof( err ) ) == 0, err );
  }
  CHECK( __cpu_program_size() == im.code_size, "PROGRAM load size" );
  CHECK( __cpu_data_size() == 192, "DATA p_memsz load size" );
  CHECK( __cpu_stack_size() == 512, "STACK size comes from PT_MCPU_STACK" );
  CHECK( __cpu_heap_size() == 128 && __cpu_heap_limit() == 1024,
        "HEAP initial/maximum sizes");
  CHECK( __cpu_sp == ld.startup_base && __cpu_fp == ld.stack_size,
        "initial SP/FP startup state");
  CHECK( __cpu_data_get_u32( __cpu_stack_ptr() + ld.startup_base +
                           MCPU_STARTUP_VERSION_OFF) == MCPU_STARTUP_VERSION_CURRENT,
        "McpuStartup.version");
  CHECK( __cpu_data_get_u64( __cpu_stack_ptr() + ld.startup_base +
                           MCPU_STARTUP_ARGC_OFF) == 3,
        "McpuStartup.argc");
  p = __cpu_data_get_u64( __cpu_stack_ptr() + ld.startup_base +
                         MCPU_STARTUP_ARGV_OFF);
  CHECK( __cpu_ptr_segment( p ) == CPU_PTR_SEG_STACK,
        "McpuStartup.argv segment");
  CHECK( __cpu_data_get_u64( __cpu_stack_ptr() + __cpu_ptr_offset( p ) + 3 * 8 ) == 0,
        "argv[argc] == NULL");
  for( i = 64; i < 192; ++i )
    CHECK( __cpu_data_ptr()[i] == 0, "NOBITS .bss loads as zero" );

  /* Negative loader checks: it validates, never repairs. */
  CHECK( cpu_melf_readall( exe, &file, &file_size, err, sizeof( err ) ) == 0, err );
  data = file[MELF_EI_DATA];
  {
    __mpu_uint64_t phoff = cpu_melf_get64( file + MELF64_EHDR_E_PHOFF_OFF, data );
    __mpu_uint16_t phnum = cpu_melf_get16( file + MELF64_EHDR_E_PHNUM_OFF, data );
    __mpu_uint64_t dataoff = UINT64_MAX;
    for( i = 0; i < phnum; ++i )
    {
      unsigned char *q = file + (size_t)phoff + i * MELF64_PHDR_SIZE;
      if( cpu_melf_get32( q + MELF64_PHDR_P_TYPE_OFF, data ) == MELF_PT_LOAD &&
         MELF_MCPU_PTR64_SEG( cpu_melf_get64( q + MELF64_PHDR_P_VADDR_OFF, data ) ) ==
           MELF_MCPU_SEG_DATA)
      {
        dataoff = cpu_melf_get64( q + MELF64_PHDR_P_OFFSET_OFF, data );
        break;
      }
    }
    CHECK( dataoff != UINT64_MAX && dataoff < file_size, "find DATA PT_LOAD" );
    file[(size_t)dataoff] = 1;
    CHECK( write_copy( bad_guard, file, file_size ) == 0, "write bad guard file" );
    file[(size_t)dataoff] = 0;
  }
  CHECK( expect_load_fail( bad_guard, "DATA NULL guard" ) == 0,
        "loader rejects corrupted DATA guard");

  cpu_melf_put64( file + MELF64_EHDR_E_ENTRY_OFF,
                 MELF_MCPU_PTR64( MELF_MCPU_SEG_PROGRAM, 1 ), data);
  CHECK( write_copy( bad_entry, file, file_size ) == 0, "write bad entry file" );
  CHECK( expect_load_fail( bad_entry, "without INIT" ) == 0,
        "loader rejects inconsistent lifecycle entries");

  rc = 0;
out:
  free( file );
  cpu_asm_image_free( &im );
  cpu_asm_object_free( &built[0] );
  cpu_asm_object_free( &built[1] );
  cpu_asm_object_free( &readback[0] );
  cpu_asm_object_free( &readback[1] );
  __cpu_free_context();
  if( dtemplate[0] )
  {
    unlink( main_o );
    unlink( lib_o );
    unlink( exe );
    unlink( life_exe );
    unlink( bad_guard );
    unlink( bad_entry );
    rmdir( dtemplate );
  }
  if( !rc ) puts( "libcpu MELF ET_REL/ET_EXEC ABI regression: PASS" );
  return( rc );
}