#include <stdio.h>
#include <string.h>
#include <libmcpu.h>
#include "mcpu-asm.h"
#include "mcpu-link.h"

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


static int
expect_link_fail( const char *a, const char *b, const char *needle )
{
  cpu_asm_object o[2];
  cpu_asm_image im;
  char err[512];
  int ok = 0;
  memset( o, 0, sizeof( o ) );
  memset( &im, 0, sizeof( im ) );
  if( cpu_assemble_object( a, &o[0], err, sizeof( err ) ) )
  {
    fprintf( stderr, "negative object A did not assemble: %s\n", err );
    goto out;
  }
  if( cpu_assemble_object( b, &o[1], err, sizeof( err ) ) )
  {
    fprintf( stderr, "negative object B did not assemble: %s\n", err );
    goto out;
  }
  if( !cpu_link_objects( o, 2, &im, err, sizeof( err ) ) )
  {
    fprintf( stderr, "negative link unexpectedly succeeded\n" );
    cpu_asm_image_free( &im );
    goto out;
  }
  if( needle && !strstr( err, needle ) )
  {
    fprintf( stderr, "negative link error mismatch: %s\n", err );
    goto out;
  }
  ok = 1;
out:
  cpu_asm_object_free( &o[0] );
  cpu_asm_object_free( &o[1] );
  return( ok ? 0 : -1 );
}

int
main( void )
{
  const char *main_src=
    ".global _start, _life\n"
    ".extern add_two\n"
    ".extern shared\n"
    ".text\n"
    "_start_life:\n"
    " LDI r64@0x00000, 40\n"
    " CALL add_two\n"
    " STORE shared, r64@0x00000\n"
    " LOAD r64@0x00018, shared\n"
    " LEA r64@0x00020, shared+0\n"
    " LDI r32@0x00030, 0\n"
    " EXIT r32@0x00030\n";
  const char *lib_src=
    ".global add_two\n"
    ".global shared\n"
    ".global shared_ptr\n"
    ".global fn_ptr\n"
    ".data\n"
    ".align 32\n"
    "shared:\n"
    " .integer 64, 0\n"
    "shared_ptr:\n"
    " .ptr shared\n"
    "fn_ptr:\n"
    " .cptr add_two\n"
    ".text\n"
    "add_two:\n"
    " LDI r64@0x00008, 2\n"
    " IADD r64@0x00000, r64@0x00000, r64@0x00008\n"
    " RET\n";
  cpu_asm_object o[2];
  cpu_asm_image im;
  char err[512];
  size_t i, ix;
  __mpu_uint64_t p;
  __mpu_init();
  memset( o, 0, sizeof( o ) );
  memset( &im, 0, sizeof( im ) );
  if( cpu_assemble_object( main_src, &o[0], err, sizeof( err ) ) )
  {
    fprintf( stderr, "main object: %s\n", err );
    return( 1 );
  }
  if( cpu_assemble_object( lib_src, &o[1], err, sizeof( err ) ) )
  {
    fprintf( stderr, "lib object: %s\n", err );
    return( 1 );
  }
  if( o[0].data_size != 0 )
  {
    fprintf( stderr, "relocatable DATA unexpectedly contains NULL guard\n" );
    return( 1 );
  }
  if( o[1].data_align != 32 )
  {
    fprintf( stderr, "DATA alignment not preserved: %llu\n", (unsigned long long)o[1].data_align );
    return( 1 );
  }
  if( sym( &o[0], "add_two", &ix ) || o[0].symbols[ix].section != CPU_OBJ_SEC_UNDEF )
  {
    fprintf( stderr, "extern symbol missing\n" );
    return( 1 );
  }
  if( sym( &o[1], "shared", &ix ) || o[1].symbols[ix].binding != CPU_OBJ_BIND_GLOBAL || o[1].symbols[ix].value != 0 )
  {
    fprintf( stderr, "defined global shared bad\n" );
    return( 1 );
  }
  if( o[0].reloc_count != 4 || o[1].reloc_count != 2 )
  {
    fprintf( stderr, "unexpected relocation counts %zu %zu\n", o[0].reloc_count, o[1].reloc_count );
    return( 1 );
  }
  if( cpu_link_objects( o, 2, &im, err, sizeof( err ) ) )
  {
    fprintf( stderr, "link: %s\n", err );
    return( 1 );
  }
  if( !im.has_life || im.life_entry != 0 )
  {
    fprintf( stderr, "bad linked _life entry\n" );
    return( 1 );
  }
  if( im.data_size != 56 )
  {
    fprintf( stderr, "linked DATA size/alignment/guard mismatch: %zu\n", im.data_size );
    return( 1 );
  }
  for( i = 0; i < CPU_DATA_NULL_GUARD_SIZE; i++ ) if( im.data[i] )
  {
    fprintf( stderr, "DATA guard is not zero\n" );
    return( 1 );
  }
  /* shared is aligned at final DATA offset 32; shared_ptr immediately follows. */
  p = __cpu_data_get_u64( im.data + 40 );
  if( __cpu_ptr_segment( p ) != CPU_PTR_SEG_DATA || __cpu_ptr_offset( p ) != 32 )
  {
    fprintf( stderr, ".ptr relocation bad: 0x%llx\n", (unsigned long long)p );
    return( 1 );
  }
  p = __cpu_data_get_u64( im.data + 48 );
  if( __cpu_ptr_segment( p ) != CPU_PTR_SEG_PROGRAM )
  {
    fprintf( stderr, ".cptr segment bad\n" );
    return( 1 );
  }
  cpu_asm_image_free( &im );
  cpu_asm_object_free( &o[0] );
  cpu_asm_object_free( &o[1] );

  /* _end is an executable linker-defined HEAP pointer at HEAP:16. */
  {
    cpu_asm_object x =
    {
      0
    };
    cpu_asm_image xi;
    const char *src=
      ".global _start, _life\n"
      ".extern _end\n"
      ".data\n"
      "end_ptr:\n"
      " .ptr _end\n"
      ".text\n"
      "_start_life:\n"
      " LEA r64@0x00000, _end\n"
      " LDI r32@0x00010, 0\n"
      " EXIT r32@0x00010\n";
    memset( &xi, 0, sizeof( xi ) );
    if( cpu_assemble_object( src, &x, err, sizeof( err ) ) )
    {
      fprintf( stderr, "_end object: %s\n", err );
      return( 1 );
    }
    if( cpu_link_objects( &x, 1, &xi, err, sizeof( err ) ) )
    {
      fprintf( stderr, "_end link: %s\n", err );
      return( 1 );
    }
    if( xi.data_size < 24 )
    {
      fprintf( stderr, "_end DATA image too small\n" );
      return( 1 );
    }
    p = __cpu_data_get_u64( xi.data + CPU_DATA_NULL_GUARD_SIZE );
    if( __cpu_ptr_segment( p ) != CPU_PTR_SEG_HEAP || __cpu_ptr_offset( p ) != CPU_HEAP_BEGIN )
    {
      fprintf( stderr, "_end .ptr relocation bad: 0x%llx\n", (unsigned long long)p );
      return( 1 );
    }
    p = __cpu_code_get_u64( xi.code + 5 );
    if( __cpu_ptr_segment( p ) != CPU_PTR_SEG_HEAP || __cpu_ptr_offset( p ) != CPU_HEAP_BEGIN )
    {
      fprintf( stderr, "_end LEA relocation bad: 0x%llx\n", (unsigned long long)p );
      return( 1 );
    }
    cpu_asm_image_free( &xi );

    memset( &xi, 0, sizeof( xi ) );
    if( cpu_link_objects_mode( &x, 1, &xi, CPU_LINK_DYNAMIC_EXEC, err, sizeof( err ) ) )
    {
      fprintf( stderr, "_end dynamic link: %s\n", err );
      return( 1 );
    }
    if( xi.data_size < 24 )
    {
      fprintf( stderr, "_end dynamic DATA image too small\n" );
      return( 1 );
    }
    p = __cpu_data_get_u64( xi.data + CPU_DATA_NULL_GUARD_SIZE );
    if( __cpu_ptr_segment( p ) != CPU_PTR_SEG_HEAP || __cpu_ptr_offset( p ) != CPU_HEAP_BEGIN )
    {
      fprintf( stderr, "_end dynamic .ptr relocation bad: 0x%llx\n", (unsigned long long)p );
      return( 1 );
    }
    p = __cpu_code_get_u64( xi.code + 5 );
    if( __cpu_ptr_segment( p ) != CPU_PTR_SEG_HEAP || __cpu_ptr_offset( p ) != CPU_HEAP_BEGIN )
    {
      fprintf( stderr, "_end dynamic LEA relocation bad: 0x%llx\n", (unsigned long long)p );
      return( 1 );
    }
    if( xi.dyn_reloc_count != 0 )
    {
      fprintf( stderr, "_end must not remain a dynamic relocation\n" );
      return( 1 );
    }
    cpu_asm_image_free( &xi );
    cpu_asm_object_free( &x );
  }

  /* A hexadecimal absolute expression must not mistake the `x` in 0x20 for
     a relocatable symbol. */
  {
    cpu_asm_object x =
    {
      0
    };
    const char *src=
      ".global _start, _life\n"
      ".text\n"
      "_start_life:\n"
      " LEA r64@0x00000, 0x20\n"
      " LDI r32@0x00010, 0\n"
      " EXIT r32@0x00010\n";
    if( cpu_assemble_object( src, &x, err, sizeof( err ) ) )
    {
      fprintf( stderr, "hex absolute expression: %s\n", err );
      return( 1 );
    }
    if( x.code_size < 13 )
    {
      fprintf( stderr, "hex absolute expression produced short code\n" );
      return( 1 );
    }
    p = __cpu_code_get_u64( x.code + 5 );
    if( __cpu_ptr_segment( p ) != CPU_PTR_SEG_DATA || __cpu_ptr_offset( p ) != 0x20 )
    {
      fprintf( stderr, "hex absolute LEA bad: 0x%llx\n", (unsigned long long)p );
      return( 1 );
    }
    cpu_asm_object_free( &x );
  }

  /* .equ and relocatable global/extern symbols occupy different namespaces
     and may not reuse one name. */
  {
    cpu_asm_object x =
    {
      0
    };
    const char *src=
      ".equ N, 4*512\n"
      ".global N\n"
      ".global _start, _life\n"
      ".text\n_start_life:\n LDI r32@0x00000, 0\n EXIT r32@0x00000\n";
    if( !cpu_assemble_object( src, &x, err, sizeof( err ) ) )
    {
      fprintf( stderr, ".equ/global conflict unexpectedly accepted\n" );
      cpu_asm_object_free( &x );
      return( 1 );
    }
    if( !strstr( err, "conflicts with .equ" ) )
    {
      fprintf( stderr, ".equ/global conflict error mismatch: %s\n", err );
      return( 1 );
    }
  }

  if( expect_link_fail(
      ".global _start, _life\n.extern missing\n.text\n_start_life:\n CALL missing\n LDI r32@0x10,0\n EXIT r32@0x10\n",
      ".data\nlocal:\n .integer 64,0\n",
      "undefined symbol: missing")) return( 1 );

  if( expect_link_fail(
      ".global _start, _life, foo\n.data\nfoo:\n .integer 64,1\n.text\n_start_life:\n LDI r32@0x10,0\n EXIT r32@0x10\n",
      ".global foo\n.data\nfoo:\n .integer 64,2\n",
      "multiple definition of global symbol: foo")) return( 1 );

  if( expect_link_fail(
      ".global _start, _life\n.extern data_sym\n.data\np:\n .cptr data_sym\n.text\n_start_life:\n LDI r32@0x10,0\n EXIT r32@0x10\n",
      ".global data_sym\n.data\ndata_sym:\n .integer 64,0\n",
      ".cptr target is not PROGRAM: data_sym")) return( 1 );

  __mpu_free_context();
  return( 0 );
}