#include "objectFile.h"
#include "hashmap.h"
#include "graph.h"
#include <stdlib.h>
#include <string.h>
#include <ctype.h>

#define SO_FOUNDNAME     1
#define SO_FOUNDSECTION  2

#define SO_PREFIX_COUNT 3
static const char* prefix[SO_PREFIX_COUNT] = {".text$", ".rdata$", ".data$"};
static hashmap* sectionMap = 0;
static list* unknownSection = 0;

struct s_objectFile
{
  const char* name;
  list* sects;
  unsigned long count;
};

static char* trim(char* src)
{
  while (isspace(*src))
    ++src;
  return src;
}

static void upper(char* src)
{
  while (*src)
  {
    *src = toupper(*src);
    ++src;
  }
}

static unsigned long key(const char* string)
{
  char c;
  unsigned long hashkey = 0;
  while ((c = *string++))
    hashkey = (c<<7) + c + (hashkey<<6) + (hashkey<<16) - hashkey;
  
  return hashkey;
}

static void parseRelocSection(graph* src, FILE* file)
{
  static char buffer[256];
  char *ptr,
       *token;
  
  while (!feof(file))
  {
    fgets(buffer, sizeof(buffer), file);
    ptr = trim(buffer);
    
    /* a blank line quits the table */
    if (!*ptr)
      return;
    
    /* skip OFFSET */
    token = strtok(ptr, " ");
    /* skip TYPE */
    token = strtok(0, " ");
    
    /* process VALUE */
    token = strtok(0, " ");
    
    if (token)
    {
      graph* dest;
      
      /* trim front-end */
      if (*token == '_')
        ++token;
      /* FASTCALL convention */
      else if (*token == '@')
      {
        ++token;
        token = strtok(token, "@");
      }
      /* check for various prefixes */
      else
      {
        int i = SO_PREFIX_COUNT;
        
        while (i--)
          if (!strncmp(token, prefix[i], strlen(prefix[i])))
          {
            token += strlen(prefix[i]);
            break;
          }
      }
      
      /* trim back-end */
      {
        char* i = token + strlen(token);
        while (isspace(*--i));
        i[1] = 0;
        
        /* watchout for STDCALL convention */
        while (isdigit(*--i));
        if (*i == '@')
          *i = 0;
      }
      
      dest = hashmapGet(sectionMap, key(token));
      
      if (dest)
      {
        /* normal dependency spotted - link engaged */
        if (src)
          graphConnect(src, dest);
        /* dependency with unknown section spotted; we need to calculate its
           dependencies as well, because this section will survive for sure */
        else
          listAdd(unknownSection, dest);
      }
    }
  }
}

static void colorizeGraph(graph* seed, unsigned long color)
{
  unsigned long c = graphGetColorNode(seed);
  if ((c | color) != c)
  {
    list* depends = graphGetConnections(seed);
    graphColorNode(seed, c | color);
    
    listStart(depends);
    while (listNext(depends))
      colorizeGraph((graph*)listGet(depends), color);
  }
}

objectFile* objectFileCreate(const char* name)
{
  objectFile* res = (objectFile*)malloc(sizeof(objectFile));
  
  res->name = name;
  res->sects = listCreate();
  res->count = 0;
  
  return res;
}

void objectFileCollect(objectFile* src, FILE* file)
{
  unsigned long progress = 0;
  char buffer[256],
       *ptr,
       *token;
  
  while (!feof(file))
  {
    fgets(buffer, sizeof(buffer), file);
    ptr = trim(buffer);
    
    if (*ptr)
    {
      switch (progress)
      {
        /* search for the filename */
        case 0:
          token = strtok(ptr, ":");
          if (token && !strcmp(token, src->name))
            progress = SO_FOUNDNAME;
          break;
          
        /* search for the section keyword */
        case SO_FOUNDNAME:
          token = strtok(ptr, ":");
          if (token)
          {
            upper(token);
            if (!strcmp(token, "SECTIONS"))
            {
              fgets(buffer, sizeof(buffer), file);
              progress = SO_FOUNDSECTION;
            }
          }
          break;
          
        /* march through the section table */
        case SO_FOUNDSECTION:
          token = strtok(ptr, " ");
          
          if (token)
          {
            token = strtok(0, " ");
            if (token)
            {
              int i = SO_PREFIX_COUNT;
              
              /* lookout for various prefixes */
              while (i--)
                if (!strncmp(token, prefix[i], strlen(prefix[i])))
                {
                  listAdd(src->sects, graphCreate(strdup(token)));
                  ++src->count;
                  break;
                }
            }
          }
      }
    }
    /* a blank line indicates the end of the section table */
    else if (progress == SO_FOUNDSECTION)
      return;
  }
}

void objectFileCompute(list* oFiles, FILE* file)
{
  objectFile* cFile = 0;
  graph*      cGraph = 0;
  
  /* check for older runs... just to be clean */
  if (sectionMap)
    hashmapDelete(sectionMap);
  if (unknownSection)
    listDelete(unknownSection);
  
  sectionMap = hashmapCreate(0);
  unknownSection = listCreate();
  
  /* insert all functions and data */
  listStart(oFiles);
  while (listNext(oFiles))
  {
    cFile = (objectFile*)listGet(oFiles);
    if (cFile->count)
    {
      listStart(cFile->sects);
      while (listNext(cFile->sects))
      {
        char* token;
        int i = SO_PREFIX_COUNT;
        cGraph = (graph*)listGet(cFile->sects);
        
        token = (char*)graphGetNameNode(cGraph);
        
        /* skip the prefix for key generation
           (needed when reading the relocation table)*/
        while (i--)
          if (!strncmp(token, prefix[i], strlen(prefix[i])))
          {
            token += strlen(prefix[i]);
            break;
          }
        
        hashmapInsert(sectionMap, cGraph, key(token));
      }
    }
  }
  
  /* parse the file */
  {
    char buffer[256],
         *ptr,
         *token;
    
    while (!feof(file))
    {
      fgets(buffer, sizeof(buffer), file);
      ptr = trim(buffer);
      
      if (*ptr)
      {
        token = strtok(ptr, ":");
        if (token)
        {
          /* look for the relocation keyword */
          if (!strncmp(token, "RELOCATION", sizeof("RELOCATION")-1))
          {
            int i;
            
            /* get the name */
            token += sizeof("RELOCATION");
            while (*token && *token++ != '[');
            
            /* just to be sure */
            if (!*token)
            {
              fprintf(stderr, "ERROR: file with relocation table "
                              "has invalid format!\n");
              return;
            }
            
            /* we could easily search for the '$', but that would create
               more dependencies to the compilers naming conventions */
            i = SO_PREFIX_COUNT;
            while (i--)
              if (!strncmp(token, prefix[i], strlen(prefix[i])))
              {
                token += strlen(prefix[i]);
                break;
              }
            
            token = strtok(token, "]");
            
            /* just to be sure */
            if (!token)
            {
              fprintf(stderr, "ERROR: file with relocation table "
                              "has invalid format!\n");
              return;
            }
            
            cGraph = (graph*)hashmapGet(sectionMap, key(token));
            
            /* skip the tables caption */
            fgets(buffer, sizeof(buffer), file);
            parseRelocSection(cGraph, file);
          }
        }
      }
    }
  }
  
  /* colorize the unknown section dependencies */
  
  /* this part may be optimized a bit by calculating their dependencies
     as well, but I'm unsure if that wouldn't be too aggressive, removing
     sections that may be needed somehow (when the linker adds code to the
     final exe for example), so I left them in to be sure */
  
  listStart(unknownSection);
  while (listNext(unknownSection))
    colorizeGraph((graph*)listGet(unknownSection), 0x80000000);
}

void objectFileColorize(const char* seed, unsigned long color)
{
  graph* start = (graph*)hashmapGet(sectionMap, key(seed));
  
  if (start)
    colorizeGraph(start, color);
}

list* objectFileGetUsed(objectFile* src)
{
  list* res = listCreate();
  
  listStart(src->sects);
  while (listNext(src->sects))
  {
    graph* curr = (graph*)listGet(src->sects);
    if (graphGetColorNode(curr))
      listAdd(res, graphGetNameNode(curr));
  }
  
  return res;
}

list* objectFileGetUnused(objectFile* src)
{
  list* res = listCreate();
  
  listStart(src->sects);
  while (listNext(src->sects))
  {
    graph* curr = (graph*)listGet(src->sects);
    if (!graphGetColorNode(curr))
      listAdd(res, graphGetNameNode(curr));
  }
  
  return res;
}

const char* objectFileGetName(objectFile* src)
{
  return src->name;
}

void objectFileDumpMap(list* oFiles)
{
  /* yay ^_^, what a loop */
  printf("\n<MAP>\n");
  listStart(oFiles);
  while (listNext(oFiles))
  {
    objectFile* oFile = (objectFile*)listGet(oFiles);
    printf("\t<FILE name=\"%s\">\n", oFile->name);
    
    listStart(oFile->sects);
    while (listNext(oFile->sects))
    {
      graph* sect = (graph*)listGet(oFile->sects);
      list* depend = graphGetConnections(sect);
      printf("\t\t<SECTION name=\"%s\" color=%lu>\n",
             (char*)graphGetNameNode(sect), graphGetColorNode(sect));
      
      listStart(depend);
      while (listNext(depend))
      {
        graph* d = (graph*)listGet(depend);
        printf("\t\t\t<DEPENDS>%s</DEPENDS>\n",
               (char*)graphGetNameNode(d));
      }
      
      printf("\t\t</SECTION>\n");
    }
    printf("\t</FILE>\n");
  }
  printf("\n</MAP>\n");
}

void objectFileDumpUsed(list* oFiles)
{
  printf("\n<USED>\n");
  
  listStart(oFiles);
  while (listNext(oFiles))
  {
    objectFile* obj = (objectFile*)listGet(oFiles);
    list* curr = objectFileGetUsed(obj);
    
    printf("\t<FILE name=\"%s\">\n", obj->name);
    
    listStart(curr);
    while (listNext(curr))
      printf("\t\t<SECTION>%s</SECTION>\n", (char*)listGet(curr));
    
    listDelete(curr);
    
    printf("\t</FILE>\n");
  }
  printf("</USED>\n");
}

void objectFileDumpUnused(list* oFiles)
{
  printf("\n<UNUSED>\n");
  
  listStart(oFiles);
  while (listNext(oFiles))
  {
    objectFile* obj = (objectFile*)listGet(oFiles);
    list* curr = objectFileGetUnused(obj);
    
    printf("\t<FILE name=\"%s\">\n", obj->name);
    
    listStart(curr);
    while (listNext(curr))
      printf("\t\t<SECTION>%s</SECTION>\n", (char*)listGet(curr));
    
    listDelete(curr);
    
    printf("\t</FILE>\n");
  }
  printf("</UNUSED>\n");
}