/*
** PE Tabulator
*/

#include "mr4th_base.h"
#include "mr4th_stdio.h"
#include "mr4th_cmdln.h"
#include "mr4th_ascii.h"

#include "tinker/tinker_coff.h"

#include "mr4th_pe.table.h"

#include "mr4th_base.c"
#include "mr4th_stdio.c"
#include "mr4th_cmdln.c"
#include "mr4th_ascii.c"
#include "tinker/tinker_coff.c"

////////////////////////////////
// PE Tabulator Argument Functions

MR4TH_SYMBOL_STATIC PETAB_Args*
petab_args(Arena *arena, String8List *command_line_args){
  // parse
  CMDLN *cmdln = cmdln_from_args(arena, command_line_args);
  
  // setup args
  PETAB_Args *args = push_array(arena, PETAB_Args, 1);
  
  // get input file name
  args->input_file_name = cmdln_input_from_idx(cmdln, 0);
  
  return(args);
}

////////////////////////////////
// Entry Point

int main(int argc, char **argv){
  os_main_init(argc, argv);
  Arena *arena = arena_alloc();
  
  // arguments
  String8List command_line_args = os_command_line_arguments();
  PETAB_Args *args = petab_args(arena, &command_line_args);
  
  // load input file
  String8 data = os_file_read(arena, args->input_file_name);
  
  /////////////////
  ///// PARSE /////
  /////////////////
  
  // extract dos header
  PE_DosHeader *dos_header = 0;
  if (sizeof(*dos_header) <= data.size){
    dos_header = (PE_DosHeader*)(data.str);
    if (dos_header->magic != PE_MSDOS_MAGIC){
      dos_header = 0;
    }
  }
  
  // extract coff file offset
  U32 coff_file_offset = 0;
  if (dos_header != 0){
    coff_file_offset = dos_header->coff_file_offset;
  }
  
  // extract pe signature
  U32 pe_signature = 0;
  U32 coff_header_offset = 0;
  if (0 < coff_file_offset &&
      coff_file_offset + sizeof(pe_signature) <= data.size){
    pe_signature = *(U32*)(data.str + coff_file_offset);
    if (pe_signature == PE_SIGNATURE){
      coff_header_offset = coff_file_offset + sizeof(pe_signature);
    }
  }
  
  // extract coff header
  PE_CoffHeader *coff_header = 0;
  U64 coff_header_opl = 0;
  if ((dos_header == 0 || 0 < coff_header_offset) &&
      coff_header_offset + sizeof(*coff_header) <= data.size){
    coff_header = (PE_CoffHeader*)(data.str + coff_header_offset);
    coff_header_opl = coff_header_offset + sizeof(*coff_header);
  }
  
  // optional header
  U64 optional_header_opl = coff_header_opl;
  if (coff_header != 0 &&
      coff_header->optional_header_size > 0){
    optional_header_opl = coff_header_opl + coff_header->optional_header_size;
  }
  
  // determine if this is an image file
  B32 is_image = (coff_header_offset != 0);
  
  // parse section table
  PE_SectionHeader *section_table = 0;
  U64 section_count = 0;
  if (coff_header != 0){
    U64 section_table_offset = optional_header_opl;
    U64 section_table_opl = section_table_offset + coff_header->section_count*sizeof(PE_SectionHeader);
    if (section_table_opl <= data.size){
      section_table = (PE_SectionHeader*)(data.str + section_table_offset);
      section_count = coff_header->section_count;
    }
  }
  
  // parse symbol table
  PETAB_CoffSymbol *symbols = 0;
  U32 symbol_count = 0;
  
  if (coff_header != 0){
    U64 table_foff = coff_header->symbol_table_foff;
    U64 table_size = coff_header->symbol_count*sizeof(PE_CoffSymbolRecord);
    if (table_foff + table_size + 4 <= data.size){
      
      U32 string_table_size = *(U32*)(data.str + table_foff + table_size);
      if (table_foff + table_size + string_table_size <= data.size){
        
        // allocate space for processed symbols
        symbols = push_array(arena, PETAB_CoffSymbol, coff_header->symbol_count);
        
        // fill string table
        U32 string_table_foff = table_foff + table_size;
        U8* coff_string_table = data.str + string_table_foff;
        U32 coff_string_table_size = string_table_size;
        
        // allocate output
        U32 max_count = coff_header->symbol_count;
        
        // setup parse loop
        PE_CoffSymbolRecord *record = (PE_CoffSymbolRecord*)(data.str + table_foff);
        PE_CoffSymbolRecord *opl = record + coff_header->symbol_count;
        
        PETAB_CoffSymbol *symbol_ptr = symbols;
        PETAB_CoffSymbol *symbol_opl = symbols + coff_header->symbol_count;
        
        for (;record < opl;){
          // extract name
          String8 name = {0};
          if (record->name.zero == 0){
            name = str8_cstring_capped(coff_string_table + record->name.string_table_off,
                                       coff_string_table + string_table_size);
          }
          else{
            name = str8_cstring_capped(record->name.raw, record->name.raw + 8);
          }
          
          // save processed symbol
          symbol_ptr->name = name;
          symbol_ptr->val = record->val;
          symbol_ptr->sec_number = record->sec_number;
          symbol_ptr->type = record->type;
          symbol_ptr->storage_class = record->storage_class;
          symbol_ptr->aux_symbol_count = record->aux_symbol_count;
          symbol_ptr += 1;
          
          // increment
          record += 1 + record->aux_symbol_count;
        }
        
        // save final count & put back unused memory
        symbol_count = (symbol_ptr - symbols);
        arena_pop_amount(arena, (symbol_opl - symbol_ptr)*sizeof(*symbol_ptr));
      }
      
    }
  }
  
  /////////////////
  ///// PRINT /////
  /////////////////
  
  STREAM *stream = stream_new();
  STREAM *tempstr = stream_new();
  
  
  // sections
  {
    TABLE *table = table_begin(arena);
    
    {
      table_default_align(table, TABLE_CellFlag_CenterAlign);
      
      // TODO(allen): conditionals on columns
      table_row(arena, table);
      table_cellf(arena, table, "#");
      table_cellf(arena, table, "Name");
      table_cellf(arena, table, "V.Size");
      table_cellf(arena, table, "V.Off");
      table_cellf(arena, table, "F.Size");
      table_cellf(arena, table, "F.Off");
      table_cellf(arena, table, "Mem");
      table_cellf(arena, table, "#Relocs");
      table_cellf(arena, table, "#Lines");
      table_hz_separator(arena, table);
      
      table_default_align(table, TABLE_CellFlag_LeftAlign);
    }
    
    {
      PE_SectionHeader *sec = section_table;
      for (U64 i = 0; i < section_count; i += 1, sec += 1){
        table_row(arena, table);
        
        table_cellf(arena, table, "%u", i + 1);
        table_cell_align(table, TABLE_CellFlag_RightAlign);
        
        table_cellf(arena, table, "%.*s", ArrayCount(sec->name), sec->name);
        
        table_default_align(table, TABLE_CellFlag_RightAlign);
        table_cellf(arena, table, "%$u",  sec->vsize);
        table_cellf(arena, table, "0x%x", sec->voff);
        table_cellf(arena, table, "%$u",  sec->fsize);
        table_cellf(arena, table, "0x%x", sec->foff);
        table_default_align(table, TABLE_CellFlag_LeftAlign);
        
        // memory flags
        // TODO(allen): here we want to be able to write this cell
        // by writing into a stream
        {
          stream_clear(tempstr);
          stream_printf(tempstr, (sec->flags & PE_SectionFlag_MEM_READ   )?"r":"-");
          stream_printf(tempstr, (sec->flags & PE_SectionFlag_MEM_WRITE  )?"w":"-");
          stream_printf(tempstr, (sec->flags & PE_SectionFlag_MEM_EXECUTE)?"x":"-");
          table_cell(arena, table, stream_read(arena, tempstr));
        }
        
        table_default_align(table, TABLE_CellFlag_RightAlign);
        table_cellf(arena, table, "%u", sec->relocation_count);
        table_cellf(arena, table, "%u", sec->line_number_count);
        table_default_align(table, TABLE_CellFlag_LeftAlign);
      }
    }
    
    stream_printf(stream, "SECTIONS: (Total %u)\n", section_count);
    table_print(stream, table, &table_wire_style);
    stream_printf(stream, "\n");
  }
  
  
  // symbols
  if (symbol_count > 0){
    TABLE *table = table_begin(arena);
    
    {
      table_default_align(table, TABLE_CellFlag_CenterAlign);
      
      table_row(arena, table);
      table_cellf(arena, table, "Name");
      table_cellf(arena, table, "Val");
      table_cellf(arena, table, "Sec#");
      table_cellf(arena, table, "Sym.Type");
      table_cellf(arena, table, "Storage");
      table_cellf(arena, table, "Aux#");
      table_hz_separator(arena, table);
      
      table_default_align(table, TABLE_CellFlag_LeftAlign);
    }
    
    {
      PETAB_CoffSymbol *symbol = symbols;
      for (U64 i = 0; i < symbol_count; i += 1, symbol += 1){
        table_row(arena, table);
        
        table_cellf(arena, table, "%S", &symbol->name);
        
        table_default_align(table, TABLE_CellFlag_RightAlign);
        table_cellf(arena, table, "%u", symbol->val);
        table_cellf(arena, table, "%u", symbol->sec_number);
        table_default_align(table, TABLE_CellFlag_LeftAlign);
        
        String8 type_str = pe_str8_from_coff_symbol_type(arena, symbol->type);
        table_cellf(arena, table, "%S", &type_str);
        
        String8 storage_str = pe_str8_from_coff_symbol_storage_class(symbol->storage_class);
        table_cellf(arena, table, "%S", &storage_str);
        
        table_cellf(arena, table, "%u", symbol->aux_symbol_count);
      }
    }
    
    stream_printf(stream, "SYMBOLS: (Total %u)\n", symbol_count);
    table_print(stream, table, &table_wire_style);
    stream_printf(stream, "\n");
  }
  
  stream_fprint(stdout, stream);
  
  return(0);
}
