#include <ctype.h>
#include <errno.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <libmcpu.h>
#include <libmpuio.h>
#include "mcpu-asm.h"
#include "mcpu-object.h"
#include "mcpu-version.h"

#define CPU_AS_LINE_MAX 2147483647ULL

typedef struct source_marker
{
  size_t physical_line;
  uint64_t logical_line;
  char *filename;
} source_marker;

typedef struct source_map
{
  source_marker *marker;
  size_t count;
  size_t capacity;
} source_map;

static char *
readall( const char *name )
{
  FILE *f;
  long z;
  char *p;

  f = fopen( name, "rb" );
  if( !f )
    return( NULL );
  if( fseek( f, 0, SEEK_END ) != 0 )
  {
    fclose( f );
    return( NULL );
  }
  z = ftell( f );
  if( z < 0 )
  {
    fclose( f );
    return( NULL );
  }
  rewind( f );
  p = (char *)malloc( (size_t)z + 1U );
  if( !p )
  {
    fclose( f );
    return( NULL );
  }
  if( fread( p, 1, (size_t)z, f ) != (size_t)z )
  {
    free( p );
    fclose( f );
    return( NULL );
  }
  p[z] = 0;
  fclose( f );
  return( p );
}

static void
usage( FILE *fp, const char *name )
{
  fprintf( fp,
    "usage: %s [options] -o output.o input.s\n\n"
    "Assemble UTF-8 LibMCPU source into a MELF ET_REL relocatable object.\n\n"
    "Options:\n"
    "  -o <file>                  - write output to <file>\n"
    "  -v, --verbose              - print output-object summary\n"
    "  -h, --help                 - show this help\n"
    "      --version              - show program version\n"
    "      --dump-version         - show version number only\n\n"
    "Exactly one input source is required.  The input file and -o option may\n"
    "appear in either order.\n",
    name);
}

static void
source_map_free( source_map *map )
{
  size_t i;

  if( !map )
    return;
  for( i = 0; i < map->count; ++i )
    free( map->marker[i].filename );
  free( map->marker );
  memset( map, 0, sizeof( *map ) );
}

static int
source_map_add( source_map *map, size_t physical_line, uint64_t logical_line,
               const char *filename)
{
  source_marker *q;
  size_t capacity;

  if( map->count == map->capacity )
  {
    capacity = map->capacity ? map->capacity * 2U : 8U;
    q = ( source_marker * ) realloc( map->marker, capacity *sizeof( *q ) );
    if( !q )
      return( -1 );
    map->marker = q;
    map->capacity = capacity;
  }

  map->marker[map->count].physical_line = physical_line;
  map->marker[map->count].logical_line = logical_line;
  map->marker[map->count].filename = strdup( filename ? filename : "<input>" );
  if( !map->marker[map->count].filename )
    return( -1 );
  ++map->count;
  return( 0 );
}

static const source_marker *
source_map_lookup( const source_map *map, size_t physical_line )
{
  size_t i;
  const source_marker *found = NULL;

  for( i = 0; i < map->count; ++i )
  {
    if( map->marker[i].physical_line > physical_line )
      break;
    found = &map->marker[i];
  }
  return( found );
}

static uint64_t
source_map_logical_line( const source_marker *m, size_t physical_line )
{
  uint64_t delta;

  if( !m || physical_line < m->physical_line )
    return (uint64_t)physical_line;
  delta = (uint64_t)(physical_line - m->physical_line);
  if( UINT64_MAX - m->logical_line < delta )
    return( UINT64_MAX );
  return( m->logical_line + delta );
}

static int
horizontal_space( int c )
{
  return( c == ' ' || c == '\t' || c == '\v' || c == '\f' || c == '\r' );
}

/*
 * Classify source-control lines left by a preprocessor.
 *
 * Return values:
 *   0  ordinary source line
 *   1  other '#' line: ignore its contents, but count the line normally
 *   2  #pragma / #lang: ignore its contents and do not advance logical line
 *   3  preprocessor line marker: # decimal ["filename"] [flags...]
 *  -1  malformed line marker
 *
 * A line marker applies to the physical source line immediately following it.
 * Its decimal number is always base 10 and is restricted to 0..2147483647.
 * If the filename is omitted, the previous logical filename is retained.  Any
 * text following the closing filename quote is deliberately ignored: those are
 * preprocessor flags for later consumers and have no assembler semantics.
 *
 * The historical #line spelling is intentionally not recognized here.  Like
 * every other unrecognized '#' line, it is discarded without affecting
 * assembly syntax and counts as one ordinary logical source line.
 */
static int
parse_source_control( const char *line, uint64_t *next_line, char **new_filename,
                      char *why, size_t whysz )
{
  const char *p = line;
  const char *digits;
  const char *start;
  uint64_t value = 0;
  size_t n;
  char *filename = NULL;

  while( horizontal_space( (unsigned char)*p ) )
    ++p;
  if( *p != '#' )
    return( 0 );
  ++p;
  while( horizontal_space( (unsigned char)*p ) )
    ++p;

  if( !strncmp( p, "pragma", 6 ) &&
      (!p[6] || horizontal_space( (unsigned char)p[6] )) )
    return( 2 );
  if( !strncmp( p, "lang", 4 ) &&
      (!p[4] || horizontal_space( (unsigned char)p[4] )) )
    return( 2 );

  if( !isdigit( (unsigned char)*p ) )
    return( 1 );

  digits = p;
  while( isdigit( (unsigned char)*p ) )
    ++p;

  if( *p && !horizontal_space( (unsigned char)*p ) )
    return( 1 );

  while( horizontal_space( (unsigned char)*p ) )
    ++p;
  if( *p && *p != '"' )
    return( 1 );

  p = digits;
  do
  {
    unsigned int digit = (unsigned int)(*p - '0');

    if( value > ( CPU_AS_LINE_MAX - digit ) / 10ULL )
    {
      snprintf( why, whysz,
                "preprocessor line marker value is outside 0..2147483647" );
      return( -1 );
    }
    value = value * 10ULL + digit;
    ++p;
  }
  while( isdigit( (unsigned char)*p ) );

  while( horizontal_space( (unsigned char)*p ) )
    ++p;
  if( *p == '"' )
  {
    ++p;
    start = p;
    while( *p && *p != '"' )
      ++p;
    if( *p != '"' )
    {
      snprintf( why, whysz,
                "unterminated filename in preprocessor line marker" );
      return( -1 );
    }
    n = (size_t)(p - start);
    filename = (char *)malloc( n + 1U );
    if( !filename )
    {
      snprintf( why, whysz, "out of memory" );
      return( -1 );
    }
    memcpy( filename, start, n );
    filename[n] = 0;
  }

  *next_line = value;
  *new_filename = filename;
  return( 3 );
}

static void
format_location_error( char *err, size_t errsz, const char *filename,
                      uint64_t line, const char *message)
{
  snprintf( err, errsz, "%s:%llu: error: %s",
           filename ? filename : "<input>",
           (unsigned long long)line,
           message ? message : "invalid source");
}

/*
 * cpu_asm's historical line scanner uses strtok_r(..., "\n", ...), which
 * skips zero-length fields between adjacent newlines.  That is harmless for
 * assembly semantics, but it means its internal diagnostic line number does
 * not advance across a physically empty source line.
 *
 * Preprocessor line markers, however, define logical source coordinates, and
 * #pragma/#lang lines may be excluded from the logical counter.  The physical
 * line counter used by cpu_asm therefore has to stay synchronized with the
 * source map maintained by mcpu-as, including across empty lines.
 *
 * Give each physically empty line a single horizontal-space character before
 * passing the source to cpu_asm.  trim() still treats the line as empty, so no
 * assembly semantics change, while strtok_r now returns one token for every
 * physical line and its diagnostic counter remains synchronized with the
 * source map maintained by mcpu-as.
 */
static char *
copy_preserving_empty_lines( const char *source )
{
  size_t len, extra = 0, i;
  int at_line_start = 1;
  char *copy, *q;

  if( !source )
    return( NULL );

  len = strlen( source );
  for( i = 0; i < len; ++i )
  {
    if( source[i] == '\n' )
    {
      if( at_line_start )
        ++extra;
      at_line_start = 1;
    }
    else
      at_line_start = 0;
  }

  if( extra >= SIZE_MAX - len )
    return( NULL );
  copy = (char *)malloc( len + extra + 1U );
  if( !copy )
    return( NULL );

  q = copy;
  at_line_start = 1;
  for( i = 0; i < len; ++i )
  {
    if( source[i] == '\n' )
    {
      if( at_line_start )
        *q++ = ' ';
      *q++ = '\n';
      at_line_start = 1;
    }
    else
    {
      *q++ = source[i];
      at_line_start = 0;
    }
  }
  *q = 0;
  return( copy );
}

/*
 * Remove preprocessor-control lines from the token stream without changing
 * physical line numbering and build a physical->logical source map for
 * diagnostics.
 */
static int
prepare_source( const char *source, const char *input_name, char **prepared,
               source_map *map, char *err, size_t errsz)
{
  char *copy;
  char *line;
  size_t physical_line = 1;
  const char *logical_file;

  if( !source || !prepared || !map )
    return( -1 );

  memset( map, 0, sizeof( *map ) );
  copy = copy_preserving_empty_lines( source );
  if( !copy )
  {
    snprintf( err, errsz, "out of memory" );
    return( -1 );
  }

  logical_file = input_name ? input_name : "<input>";
  if( source_map_add( map, 1U, 1U, logical_file ) )
  {
    free( copy );
    source_map_free( map );
    snprintf( err, errsz, "out of memory" );
    return( -1 );
  }

  line = copy;
  while( line && *line )
  {
    char *end = strchr( line, '\n' );
    char saved = 0;
    uint64_t requested_line = 0;
    char *requested_file = NULL;
    char why[160];
    int parsed;
    const source_marker *m;

    if( end )
    {
      saved = *end;
      *end = 0;
    }

    why[0] = 0;
    parsed = parse_source_control( line, &requested_line, &requested_file,
                                   why, sizeof( why ) );
    if( parsed < 0 )
    {
      m = source_map_lookup( map, physical_line );
      format_location_error( err, errsz,
                            m ? m->filename : logical_file,
                            source_map_logical_line( m, physical_line ), why);
      free( requested_file );
      if( end )
        *end = saved;
      free( copy );
      source_map_free( map );
      return( -1 );
    }

    if( parsed > 0 )
    {
      size_t len = strlen( line );

      memset( line, ' ', len );
      if( parsed == 3 )
      {
        const source_marker *current = source_map_lookup( map, physical_line );
        const char *next_file = requested_file ? requested_file :
                                (current ? current->filename : logical_file);

        if( source_map_add( map, physical_line + 1U, requested_line, next_file ) )
        {
          free( requested_file );
          if( end )
            *end = saved;
          free( copy );
          source_map_free( map );
          snprintf( err, errsz, "out of memory" );
          return( -1 );
        }
      }
      else if( parsed == 2 )
      {
        const source_marker *current = source_map_lookup( map, physical_line );
        uint64_t current_line = source_map_logical_line( current, physical_line );
        const char *current_file = current ? current->filename : logical_file;

        if( source_map_add( map, physical_line + 1U, current_line, current_file ) )
        {
          free( requested_file );
          if( end )
            *end = saved;
          free( copy );
          source_map_free( map );
          snprintf( err, errsz, "out of memory" );
          return( -1 );
        }
      }
    }

    free( requested_file );
    if( end )
    {
      *end = saved;
      line = end + 1;
      ++physical_line;
    }
    else
      break;
  }

  *prepared = copy;
  return( 0 );
}

static void
remap_assembler_error( char *err, size_t errsz, const source_map *map,
                      const char *input_name)
{
  char old[512];
  const char *p;
  char *end;
  unsigned long long physical;
  const source_marker *m;
  uint64_t logical;

  if( !err || !*err )
    return;
  if( strncmp( err, "line ", 5 ) != 0 )
    return;

  errno = 0;
  p = err + 5;
  physical = strtoull( p, &end, 10 );
  if( errno || end == p || *end != ':' || physical == 0ULL ||
     physical > (unsigned long long)SIZE_MAX)
    return;

  snprintf( old, sizeof( old ), "%s", end + 1 );
  p = old;
  while( *p == ' ' || *p == '\t' )
    ++p;

  m = source_map_lookup( map, (size_t)physical );
  logical = source_map_logical_line( m, (size_t)physical );
  format_location_error( err, errsz,
                        m ? m->filename : (input_name ? input_name : "<input>"),
                        logical, p);
}

int
main( int ac, char **av )
{
  const char *input = NULL;
  const char *output = NULL;
  char *src = NULL;
  char *prepared = NULL;
  char err[512];
  cpu_asm_object object;
  source_map map;
  int i;
  int rc = 1;
  int verbose = 0;

  for( i = 1; i < ac; ++i )
  {
    mcpu_version_action version_action;

    if( mcpu_version_parse_option(av[i], &version_action) )
      return( mcpu_version_print(av[0], version_action) == 0 ? 0 : 1 );

    if( !strcmp( av[i], "-h" ) || !strcmp( av[i], "--help" ) )
    {
      usage( stdout, av[0] );
      return( 0 );
    }
    else if( !strcmp( av[i], "-o" ) )
    {
      if(++i >= ac || output )
      {
        usage( stderr, av[0] );
        return( 2 );
      }
      output = av[i];
    }
    else if( !strcmp( av[i], "-v" ) || !strcmp( av[i], "--verbose" ) )
      verbose = 1;
    else if( av[i][0] == '-' )
    {
      fprintf( stderr, "mcpu-as: unknown option %s\n", av[i] );
      return( 2 );
    }
    else
    {
      if( input )
      {
        fprintf( stderr, "mcpu-as: more than one input file\n" );
        return( 2 );
      }
      input = av[i];
    }
  }

  if( !input || !output )
  {
    usage( stderr, av[0] );
    return( 2 );
  }

  __mpu_init();
  src = readall( input );
  if( !src )
  {
    perror( input );
    goto out_mpu;
  }
  if( !mpu_utf8valid( (const __mpu_char8_t *)src ) )
  {
    fprintf( stderr, "mcpu-as: error: assembler source is not valid UTF-8\n" );
    goto out_src;
  }

  err[0] = 0;
  if( prepare_source( src, input, &prepared, &map, err, sizeof( err ) ) )
  {
    fprintf( stderr, "mcpu-as: %s\n", err[0] ? err : "cannot process source" );
    goto out_src;
  }

  if( cpu_assemble_object( prepared, &object, err, sizeof( err ) ) )
  {
    remap_assembler_error( err, sizeof( err ), &map, input );
    if( !strncmp( err, "line ", 5 ) || strstr( err, ": error: " ) )
      fprintf( stderr, "mcpu-as: %s\n", err );
    else
      fprintf( stderr, "mcpu-as: error: %s\n", err );
    goto out_map;
  }

  if( cpu_object_write( output, &object, err, sizeof( err ) ) )
  {
    fprintf( stderr, "mcpu-as: error: %s\n", err );
    cpu_asm_object_free( &object );
    goto out_map;
  }

  if( verbose )
    printf( "MELF ET_REL: %zu symbols, %zu relocations -> %s\n",
            object.symbol_count, object.reloc_count, output );
  cpu_asm_object_free( &object );
  rc = 0;

out_map:
  source_map_free( &map );
  free( prepared );
out_src:
  free( src );
out_mpu:
  __mpu_free_context();
  return( rc );
}