#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <bastard.h>
#include <extension.h>


#define NORM    "\033[0m"
#define RED     "\033[31m"
#define GREEN   "\033[32m"
#define BROWN   "\033[33m"
#define BLUE    "\033[34m"
#define MAGENTA "\033[35m"
#define CYAN    "\033[36m"
#define GREY    "\033[37m"
#define WHITE   "\033[1m"
char asmColor[]  = WHITE "%n" NORM "%:%^"
                   GREY "%a " MAGENTA  "%8b\t" 
                   GREEN "%m\t" CYAN "%d" NORM "%, " CYAN "%s"NORM"%, "BLUE"%t"
                   BROWN "%c %;%x" NORM;
char asmMono[]   = "%n%:%^%a %8b\t%m\t%d%, %s %c %;%x";
char asmFile[]   = "%n%:%^\t%m\t%d%, %s %c";
char dataColor[] = WHITE "%n" NORM "%:%^%a " BLUE "%8b\t" 
                   BROWN "%c %x" NORM;
char dataMono[]  = "%n%:%^%a %8b\t%c %;%x";
char dataFile[]  = "%n%:%^\t DB %32b\t%c";


char comment_prefix[] = ";", 
     reg_prefix[] = "",
     imm_prefix[] = "",
     local_prefix[] = "@@";

struct EXT__ASM *settings;
#define _EXT_ASM_LINE_LENGTH 2048


/* ---------------------------------------------------- functions */
void ext_asm_init( void *param ) {
   settings = (struct EXT__ASM *)param;

   if (! settings) return;
   /* set Language-Specific information */
   settings->asm_ttyColor  = asmColor;   /* so I'm cheating, big deal... */
   settings->asm_ttyMono   = asmMono;
   settings->asm_file      = asmFile;
   settings->asm_lpr       = asmMono;

   settings->data_ttyColor = dataColor;
   settings->data_ttyMono  = dataMono;
   settings->data_file     = dataFile;
   settings->data_lpr      = dataMono;

   settings->comment = comment_prefix;
   settings->reg_pre = reg_prefix;
   settings->imm_pre = imm_prefix;
   settings->local_pre = local_prefix;

   return;
}

void ext_asm_cleanup(void) {
      return;
}

int sprint_code( long rva, char *line, int len, int output){
      struct address *addr;
      struct code inst = {0};
      struct function func = {0};
      struct name n = {0};
      struct string *str;
      struct comment *cmt;
      char *asmfmt, *datafmt, *tmp, *buf;
      void *state;
      int l;

      
      /* first, set the output format string for later */
      switch (output) {
         case ASM_OUTPUT_PRINTER:
            asmfmt = settings->asm_lpr;
            datafmt = settings->data_lpr;
            break;
         case ASM_OUTPUT_FILE:
            asmfmt = settings->asm_file;
            datafmt = settings->data_file;
            break;
         case ASM_OUTPUT_TTY_COLOR:
            asmfmt = settings->asm_ttyColor;
            datafmt = settings->data_ttyColor;
            break;
         case ASM_OUTPUT_TTY:
         default:
            asmfmt = settings->asm_ttyMono;
            datafmt = settings->data_ttyMono;
            break;
      }

      line[0] = '\0'; /* since the caller may forget this :) */
      tmp = calloc(_EXT_ASM_LINE_LENGTH, 1);
      if (! tmp) return(SetLastError(9050));
      state = DBSaveState();
      addr = GetAddressObject(rva);
      if ( ! addr  ){
         DBRestoreState(state);
         return(SetLastError(9010));
      }
      /* this may not belong here ... hmmm, maybe a call to sprint_func_start */
      if (addr->flags & ADDR_FUNCTION){
         if ( DBIndexFind( FUNCTION_RVA, &rva, &func)) {
            DBIndexFind(NAME_RVA, &rva, &n);
            snprintf(tmp, _EXT_ASM_LINE_LENGTH, 
                  "\n%s -------------------------- Subroutine %s\n", 
                  settings->comment, n.name);
            if ( func.comment && (cmt = GetCommentObject(func.comment)) ){
               strncat(tmp, settings->comment, _EXT_ASM_LINE_LENGTH); 
               strncat(tmp, cmt->text, _EXT_ASM_LINE_LENGTH);
               strncat(tmp, "\n", _EXT_ASM_LINE_LENGTH);
               free(cmt);
            }
            strncat(line, tmp, len);
            memset(tmp, 0, _EXT_ASM_LINE_LENGTH);
         }
      }
      //if (addr-?flags & ADDR_INLINE) { }
      if (addr->flags & ADDR_CODE && asmsprintf(tmp, asmfmt, addr) ){
         ; /* this is OK, do nothing */
      } else {
         /* fall through to data print */
         if (! asmsprintf(tmp, datafmt, addr) ){
            /* this is not just fine */
            free(addr); free(tmp);
            return(0);
         }
         
      }
      strncat(line, tmp, len);
      free(tmp);
      free(addr);

      /* if this is a string, print complete string */
      if ( (str = GetStringObject(rva)) && (GetEntireString(rva, &buf)) ){
         strncat(line, "\t\t; String: \"", len);

         l = strlen(line) + 3;
         if (str->length > len - l) {/* resize if nec */
            //tmp  = realloc(line, l +  str->length);
            //if (tmp) {
             //  line = tmp;
             //  l = str->length;
            //} else 
            l = len - l;
         }   else l = str->length;
         strncat(line, buf, l);
         strcat(line, "\"\n");
         free(buf);
      }

      if (str) free(str);
      DBRestoreState(state);
      return(1);
}

int sprint_asm_func_start( char *str, int len, int func_id) {
   struct function f;
   struct name n;

   if (DBIndexFind(FUNCTION_ID, &func_id, &f) && 
       DBIndexFind(NAME_RVA, &f.rva, &n) )
      return(snprintf(str, len, "%s PROC %s\n%s:", 
                        settings->comment, n.name, n.name) );
   return(0);
}
int sprint_asm_func_end( char *str, int len, int func_id) {
   struct function f;
   struct name n;

   if (DBIndexFind(FUNCTION_ID, &func_id, &f) &&
       DBIndexFind(NAME_RVA, &f.rva, &n) )
      return(snprintf(str, len, "%s ENDPROC %s", settings->comment, n.name) );
   return(0);
}

int fprint_asm_func( FILE *f, struct function *func) {
   int cont;
   char line[512];
   struct address a;

   sprint_asm_func_start(line, 512, func->id);
   fprintf(f, "%s\n", line);

   cont = DBIndexFind(ADDRESS_RVA, &func->rva, &a);
   while (cont && a.rva < func->rva + func->size ) {
      if (sprint_code(a.rva, line, 512, ASM_OUTPUT_FILE))
         fprintf(f, "%s\n", line);;   
      cont = DBIndexNext(ADDRESS_RVA, &a); 
   }

   sprint_asm_func_end(line, 512, func->id);
   fprintf(f, "%s\n\n", line);
   return(1);
}

int sprint_asm_struct( char *str, int len, long rva ) {
   struct address a;
   struct structure s;
   struct struct_member m;
   struct data_type t;
   char tmp[128];
   int cont;
   
   if (! DBIndexFind(ADDRESS_RVA, &rva, &a) )
      return(SetLastError(0));
   if (! DBIndexFind(STRUCTURE_ID, &a.structure, &s))
      return(SetLastError(0));
   
   /* print structure  header */
   snprintf(str, len, "STRUCT %s:\t%s %s bytes\n", 
                        s.name, settings->comment, s.size);
   cont = DBIndexFind(STRUCT_MEMBER_STRUCTURE, &s.id, &m);
   while (cont && m.structure == s.id) {
      DBIndexFind(DATA_TYPE_ID, &m.type, &t);
      snprintf(tmp, 128, "\t%s %s\n", t.name, m.name);
      strncat(str, tmp, len);
      cont = DBIndexNext(STRUCT_MEMBER_STRUCTURE, &m);
   }
   /* Print Structure End */
   sprintf(tmp, "%s ENDSTRUCT %s\n", settings->comment, s.name);
   strncat(str, tmp, len);
   return(1);
}

int sprint_section_start( char *str, int len, char *name ){
   return( snprintf(str, len, "SECTION %s", name) );
}

int sprint_section_end( char *str, int len, char *name ){
   return( snprintf(str, len, "%s END SECTION %s", settings->comment, name) );
}

int fprint_section(FILE *f, struct section *s){
   struct address a;
   struct function func;
   int cont;
   char line[512];

   if (sprint_section_start(line, 512, s->name))
      fprintf(f, "%s\n", line);
   line[0] = '\0';

   cont = DBFindClosestNext(ADDRESS_RVA, &s->rva, &a);
   while (cont && a.rva < s->rva + s->size ){
      if (DBIndexFind(FUNCTION_RVA, &a.rva, &func)){
         fprint_asm_func(f, &func);
         a.rva += func.size;
         cont = DBFindClosestNext(ADDRESS_RVA, &a.rva, &a);
         if (a.rva < func.rva) cont = 0; /* wrapped around DB table */
      } else {
         sprint_code(a.rva, line, 512, ASM_OUTPUT_FILE);
         fprintf(f, "%s\n", line);
         line[0]='\0';
         cont = DBIndexNext(ADDRESS_RVA, &a);
      }
   }

   if (sprint_section_end(line, 512, s->name))
      fprintf(f, "%s\n\n", line);
   return(1);
}


/* OK, here is how the Intel Effective Addresses are expressed:
     Direct Addressing
       [disp]
     Register Indirect
       [base]
     Based
       [base] + disp
     Indexed
       (scale * [index]) + disp
     Based-Indexed
       [base] + (scale * [index]) + disp
*/
int sprint_addrexp(char *str, int len, char *scale, char *index, char *base,
                                 char *disp, int sign){
   char sd = '+', idx[16] = {0};
   char tmp[32];

   if (sign & 0x0001) sd = '-';

   if (scale[0] && index[0])
      snprintf(idx, 16, "(%s * [%s])", scale, index);
   else if (index[0]) 
      snprintf(idx, 16, "[%s]", index);

   if (base[0]) {
      snprintf(str, len, "[%s]", base);
      if (idx[0]) {
         strncat(str, " + ", len);
         strncat(str, idx, len);
      }
      if (disp[0]){
         snprintf(tmp, 32, " %c %s", sd, disp);
         strncat(str, tmp, len);
      }
   } else if (idx[0]) {
         snprintf(str, len, "%s %c %s", idx, sd, disp);
   } else {
         snprintf(str, len, "%c %s", sd, disp);
   }

   return(strlen(str));
}

int gen_asm_file(char *path){
   FILE *f;
   struct section s;
   struct DISASM_TGT *t;
   int cont;
   void *state;

   f = fopen(path, "w");
   if ( ! f)
      return(SetLastError(9030));
   
   state = DBSaveState();

   /* print target info */
   t = GetBastardTgt();
   fprintf(f,"%s %s source file for project %s: %s\n\n",
            settings->comment, t->assembler, t->name, path);

   cont = DBIndexFirst(SECTION_RVA, &s);
   while (cont) {
      fprint_section(f, &s);
      cont = DBIndexNext(SECTION_RVA, &s);
   }

   fprintf(f, "\n%s EOF\n", settings->comment);
   fclose(f);
   DBRestoreState(state);
   return(1);
}

