#ifndef _asm16
#define _asm16

/*

PC i386 NASM LIKE ASSEMBLY LANGUAGE... 

Author: Matthew William Coan
Date: Wed Aug 28 23:24:38 EDT 2013

*/

#include <iostream>
#include <iomanip>
#include <string>
#include <vector>
#include <list>
#include <map>
#include <sstream>
#include <cctype>
#include <cstdlib>

#include "bfstream.h"
#include "csv.h"

namespace asm_tools {

using namespace std;
using namespace csv_reader;

void add_extern_function(const string & name);
void add_global_function(const string & name);
void add_address(const size_t addr);
void add_address2(const size_t addr);
string get_function_name(void * address);

map< string, string > code_map;

class Instruction {
public:
   string label;
   string opcode;
   string arg1;
   string arg2;
 
   bool is_register(const string & reg) {
      bool ret = false;
      if(reg == "bp" || reg == "dx" || reg == "bx" || reg == "ax" || reg == "sp") {
         ret = true;
      }
      return ret;
   }

   string to_string() {
      string ret = opcode;
      if(arg1.size()) {
         if(opcode == "db") {
            ret = arg1;
         }
         else if(isdigit(arg1[0])) {
            ret += " 0x1234";
         }
         else if(opcode == "push" || opcode == "pop" || opcode == "call") {
            ret += " " + arg1;
         }
         else if(opcode == "mov") {
            ret += " " + arg1;
         }
         else {
            ret += " " + arg1;
         }
      } 
      if(arg2.size()) {
         if(isdigit(arg2[0])) {
            ret += ",0x1234";
         }
         else if(opcode == "mov") {
            if(!is_register(arg2))
               ret += ",0x1234";
            else 
               ret += "," + arg2;
         }
         else {
            ret += "," + arg2;
         }
      }
      return ret;
   }

   string to_string2() {
      string ret = opcode;
      if(arg1.size()) {
         ret += " " + arg1;
      }
      if(arg2.size()) {
         ret += "," + arg2;
      }
      if(label.size() && ret.size()) {
         ret = label + ":" + ret;
      }
      else if(label.size()) {
         ret = label + ":";
      }
      return ret;
   }

 
   string to_string(unsigned int ch) {
      char buffer[128];
      sprintf(buffer, "%d", ch);
      return string(buffer);
   }

   string encode(const string & str) {
      string ret;
      for(size_t i = 0; i < str.size(); i++) {
         if(str[i] == '\'') {
            i++;
            while(str[i] != '\'') {
               ret += to_string((int)str[i]) + ",";
               i++;
            }
            i += 2;
            while(i < str.size()) {
               ret += str[i];
               i++;
            }
         }
         else if(isspace(str[i])) {

         }
         else {
            ret += str[i];
         }
      }
      return ret + ",";
   }

   Instruction() { }

   Instruction(const Instruction & inst) {
      label = inst.label;
      opcode = inst.opcode;
      arg1 = inst.arg1;
      arg2 = inst.arg2;
   }

   Instruction(const string & line) {
      stringstream str;
      string temp,fun;
      str << line;
      str >> label;
      size_t n,index;
      if(label == "global") {
         str >> fun;
         add_global_function(fun);
      }
      else if(label == "extern") {
         str >> fun;
         add_extern_function(fun);
      }
      else if(label == "db") {
         opcode = "db";
         str >> arg1;
         label = "";
      }
      else if(label.find(":") != string::npos) {
         label = label.substr(0,label.size()-1);
         str >> opcode;
         temp = opcode;
         if(temp == "db") {
            unsigned char ch = str.get();
            while(str) {
               if(!isspace(ch)) {
                  break; 
               }
               ch = str.get();
            }
            while(str) {
               arg1 += ch;
               ch = str.get();
            }
cout << "\"" << arg1 << "\" -> \"" << encode(arg1) << "\"" << endl;
//cin.get();cin.get();
            code_map[arg1] = encode(arg1);
         }
         else {
            str >> temp;
            if((index = temp.find(",")) != string::npos) {
               n =  temp.size() - (temp.size() - (index));
               arg1 = temp.substr(0,n);
               arg2 = temp.substr(index+1);
            }
            else {
               arg1 = temp;
            }
         }
      }
      else {
         opcode = label;
         label = "";
         str >> temp;
         if((index = temp.find(",")) != string::npos) {
            n =  temp.size() - (temp.size() - (index));
            arg1 = temp.substr(0,n);
            arg2 = temp.substr(index+1);
         }
         else {
            arg1 = temp;
         }
      }
      print(cout);
   }

   void print(ostream & out) {
      if(label.size()) {
         out << label << ":" << endl << flush;
      }
      if(arg2 != "") {
         out << opcode << " " << arg1 << "," << arg2 << endl << flush;
      }
      else if(arg1 != "") {
         out << opcode << " " << arg1 << endl << flush;
      }
      else {
         out << opcode << endl << flush;
      }
   }

   ~Instruction() { }

   Instruction & operator=(const Instruction & right) {
      label = right.label;
      opcode = right.opcode;
      arg1 = right.arg1;
      arg2 = right.arg2;
      return *this;
   }

   void write(bfstream & out) {

   }

   void read(bfstream & in) {

   }
};

typedef vector< Instruction > instruction_vector_type;
typedef vector< size_t > address_vector_type;

class Function {
public:
   enum function_type { 
      F_EXTERN,
      F_GLOBAL 
   };
   string name;
   size_t entry_point;
   size_t size;
   vector< size_t > addr_vec;
   size_t addr;
   function_type fun_type;

   Function() { 
      fun_type = F_EXTERN; 
      size = 0;
      entry_point = 0;
      addr = 0;
   }
   Function(const string & fun, const function_type t) { 
      fun_type = t; 
      name = fun;
      size = 0;
      entry_point = 0;
      addr = 0;
   }
   Function(const Function & fun) { 
      name = fun.name;
      entry_point = fun.entry_point;
      size = fun.size;
      fun_type = fun.fun_type;
      addr_vec = fun.addr_vec;
      addr = fun.addr;
   }
   ~Function() { }

   Function & operator=(const Function & right) {
      name = right.name; 
      entry_point = right.entry_point;
      size = right.size;
      fun_type = right.fun_type;
      addr_vec = right.addr_vec;
      addr = right.addr;
      return *this;
   }

   void write(bfstream & out) {
      size_t name_size = name.size()+1;
      out << (short)name_size;
      out.write((char*)name.c_str(), name_size);
cout << "NAME: \"" << name << "\"" << endl;
cout << "NAME_SIZE: \"" << name_size << "\"" << endl;
//cin.get();cin.get();
      size_t size = addr_vec.size();
      out << (short)entry_point << (short)addr << (short)size;
      for(size_t i = 0; i < addr_vec.size(); i++) {
         out << (short)addr_vec[i];
      }
      out.flush();
   }

   void read(bfstream & in) {
      addr_vec.clear();
      in >> name >> entry_point >> addr >> size;
      size_t addr;
      for(size_t i = 0; i < size; i++) {
         in >> addr;
         addr_vec.push_back(addr);
      }
   }
};

typedef vector< Function > function_vector_type;

class Program {
   string name;
   function_vector_type function_vec;
   instruction_vector_type code;
   address_vector_type addr_vec;
   address_vector_type addr_vec2;
   unsigned char * code2;
   size_t count;
   size_t code_size;
   map< string, size_t > label_map;
public:
   Program(const string & name) {
      this->name = name;
      code2 = 0;
      count = 0;
      code_size = 0;
   }
   ~Program() {
      if(code2) delete [] code2;
   }
   void clear() {
      function_vec.clear();
      code.clear();
      addr_vec.clear();
      addr_vec2.clear();
      delete [] code2;
      code2 = 0;
      count = 0;
      code_size = 0;
      label_map = map< string, size_t >();
   }
   void register_function(const string & name, void * fun) {
      Function * fptr = find_function(name);
      if(fptr) {
         fptr->addr = (size_t)fun;
      }
   }
   void add_address(const size_t addr) { addr_vec.push_back(addr); }
   void add_address2(const size_t addr) { addr_vec2.push_back(addr); }
   Function * find_function(const string & name) {
      Function * ret = 0;
      for(size_t i = 0; i < function_vec.size(); i++) {
         if(function_vec[i].name == name) {
            ret = &function_vec[i];
            break;
         }
      }
      return ret;
   }
   string disassemble0(unsigned char * code, size_t & n) {
      n = 0;
#include "findif.h"
      string ret = "db " + to_string((unsigned int)code[n]);
      n++;
      return ret;
   }
   string get_machine_code_string(const unsigned char * ptr, size_t n) {
      string ret,temp;
      char buffer[1024];
      for(size_t i = 0; i < n; i++) {
         sprintf(buffer, "%X", (unsigned int)ptr[i]);
         temp = buffer;
         if(temp.size() == 1)
            temp = "0" + temp;
         ret += temp + " ";
      }
      return ret;
   }
   string to_string(const size_t sz) {
      stringstream str;
      str << sz;
      string ret;
      str >> ret;
      return ret;
   }
   void disassemble(ofstream & out) {
cout << "disassemble..." << endl;
      out << "[BITS 32]" << endl << endl;
      for(size_t i = 0; i < function_vec.size(); i++) {
         out << "extern " << function_vec[i].name << endl;
      }
      out << "global main" << endl;
      out << endl;
      unsigned char * ptr = code2;
      size_t n = 0,temp,tt;
      string str,str2;
      char buffer[1024];
      out << "main:" << endl;
      while(n < code_size) {
         temp = 0;
         str = disassemble0(ptr, temp);
         if(str != "") {
// DEBUG
            size_t *p;
            p = (size_t*)(ptr+temp);
            sprintf(buffer, "%d", *p);
            string t = buffer;
            if(temp > 4) tt = temp - 4;
            else tt = temp;
            for(size_t i = 0; i < function_vec.size(); i++) {
               if(function_vec[i].addr == (*(size_t*)(ptr+tt))) {
                  t = function_vec[i].name;
                  break;
               }
            }
            size_t addr;
            size_t * pointer;
            for(size_t i = 0; i < addr_vec.size(); i++) {
               addr = addr_vec[i];
               pointer = (size_t *)(code2 + addr);
               if(pointer == (size_t*)(ptr+temp)) {
                  t = "L" + to_string(i);
                  break;
               }
            }
            for(size_t i = 0; i < addr_vec.size(); i++) {
               addr = addr_vec2[i];
               pointer = (size_t *)(code2 + addr);
               if(pointer == (size_t*)ptr) {
                  out << "L" << i << ":" << endl;
               }
            }

            str2 = replace2(str, "0x1234", t);
            if(str != str2) {
               str = str2;
               temp += 4;
            }

            out << "\t" << setw(15) << left << str << right 
                << "; " << get_machine_code_string(ptr, temp) << endl;
         }
         n += temp;
         ptr += temp;
      }
cout << "done..." << endl;
   }
   void run() {
cout << "TRY TO RUN" << endl;
      typedef int (*function_type)(int, char **);
      function_type main_ptr = (function_type)(code2);
      int argc = 1;
      char * argv[1];
      argv[0] = (char*)name.c_str();
      int rc = main_ptr(argc, argv);
cout << "rc = " << rc << endl;
cout << "ALL DONE" << endl;
   }
   void add_function(const Function & fun) {
      function_vec.push_back(fun); 
   }
   bool read_line(istream & in, string & line) {
      bool ret = false;
      line = "";
      unsigned char ch;
      ch = in.get();
      while(in) {
         ret = true;
         if(ch == '\n') break;
         line += ch;
         ch = in.get();
      }
      return ret;
   }

   string to_string(const int v) {
      char buffer[1024]; 
      sprintf(buffer, "%d", v);
      return buffer;
   }

   string to_code(const string & v) {
      string ret;
      int val = atoi(v.c_str());
      char * p = (char*)&val;
      ret += to_string(p[0]);
      ret += ",";
      ret += to_string(p[1]);
      ret += ",";
/*
      ret += to_string(p[2]);
      ret += ",";
      ret += to_string(p[3]); 
      ret += ",";
*/
      return ret;
   }

   string replace(const string & code, const string & value, const string & value2) {
      string ret;
      string v2 = to_code(value2);
      if(code.size() && value.size() && value2.size()) {
         size_t index = code.find(value);
         if(index != string::npos) {
            size_t n = index;
            string t1 = code.substr(0, n);
            string t2 = code.substr(index+value.size());
            ret = t1 + v2 + t2;
         }
         else {
            ret = code;
         }
      }
      else {
         ret = code;
      }
      return ret;
   }

   string replace2(const string & code, const string & value, const string & value2) {
      string ret;
      string v2 = value2;
      if(code.size() && value.size() && value2.size()) {
         size_t index = code.find(value);
         if(index != string::npos) {
            size_t n = index;
            string t1 = code.substr(0, n);
            string t2 = code.substr(index+value.size());
            ret = t1 + v2 + t2;
         }
         else {
            ret = code;
         }
      }
      else {
         ret = code;
      }
      return ret;
   }

   size_t count_com(const string & str) {
      size_t ret = 0;
      for(size_t i = 0; i < str.size(); i++) {
         if(str[i] == ',') ret++;
      }
      return ret;
   }

   string to_code2(const string & str) {
      size_t i = 0;
      while(i < str.size()) {
         if(str[i] == ',') {
            i++;
            break;
         }
         i++;
      }
      return str.substr(i);
   }

   void gen_code() { 
      ifstream fin("code.csv", ios::in);
      if(fin) {
         csv_row_type code_row;
         while(fin >> code_row) {
//            cout << "\"" << code_row[0] << "\" -> \"" << code_row[1] << "\"" << endl;
            code_map[code_row[0]] = code_row[1];
         }
         fin.close();
      }
      string temp_code, the_code, tcode;
      count = 0;
      for(size_t i = 0; i < code.size(); i++) {
         the_code = code[i].to_string();
         if(code[i].label != "") {
            label_map[code[i].label] = count;
         }
         tcode = replace(code_map[the_code], "52,18,", code[i].arg2);
         if(code_map[the_code] != "") {
            temp_code += tcode;
            count += count_com(tcode);
         }
         else {
            temp_code += tcode;
            count += count_com(tcode);
         } 
      }
      temp_code = the_code = tcode = "";
      count = 0;
      Function * p_fun = 0;
      for(size_t i = 0; i < code.size(); i++) {
         the_code = code[i].to_string();
         if((p_fun = find_function(code[i].arg2)) != 0) {
            tcode = code_map[the_code];
            if(tcode.find("52,18,") != string::npos) {
               size_t n = count + count_com(tcode) - 2;
               p_fun->addr_vec.push_back(n);
            }
            tcode = replace(code_map[the_code], "52,18,", code[i].arg2);
            cout << code[i].to_string2() << " ; " << tcode << endl;
            code[i].arg2 = to_string(label_map[code[i].arg2]);
            if(code_map[the_code] != "") {
               temp_code += tcode;
               count += count_com(tcode);
            }
         }
         else if((p_fun = find_function(code[i].arg1)) != 0) {
cout << "fund_function: \"" << code[i].arg1 << "\"" << endl;
            the_code = "call 0x1234";
cout << "the_code: \"" << the_code << "\"" << endl;
            tcode = code_map[the_code];
cout << "tcode1: \"" << tcode << "\"" << endl;
//cin.get();cin.get();
//115,86,52,18,
            if(tcode.find("52,18,") != string::npos) {
cout << "found code..." << endl;
               size_t n = count + count_com(tcode) - 2;
               p_fun->addr_vec.push_back(n);
            }
            tcode = replace(code_map[the_code], "52,18,", "0,0");
cout << "tcode2: \"" << tcode << "\"" << endl;
            cout << code[i].to_string2() << " ; " << tcode << endl;
            code[i].arg1 = to_string(p_fun->addr);
            if(code_map[the_code] != "") {
cout << "found code..." << endl;
               temp_code += tcode;
               count += count_com(tcode);
            }
//cin.get();cin.get();
         }
         else if(code[i].opcode == "db" && code[i].label == "") {
            tcode = code[i].arg1; 
            temp_code += tcode + ",";
            count++;
         }
         else if(label_map[code[i].arg2] != 0) {
            tcode = code_map[the_code];
            if(tcode.find("52,18,") != string::npos) {
              size_t n = count + count_com(tcode) - 2;
              add_address(n);
              n = label_map[code[i].arg2];
              add_address2(n);
            }
            tcode = replace(code_map[the_code], "52,18,", code[i].arg2);
            cout << code[i].to_string2() << " ; " << tcode << endl;
            code[i].arg2 = to_string(label_map[code[i].arg2]);
            if(code_map[the_code] != "") {
               temp_code += tcode;
               count += count_com(tcode);
            }
         }
         else {
            tcode = replace(code_map[the_code], "52,18,", code[i].arg2);
            cout << code[i].to_string2() << " ; " << tcode << endl;
            if(code_map[the_code] != "") {
               temp_code += tcode;
               count += count_com(tcode);
            }   
         }
      }
      code_size = count;
      cout << endl;
      cout << "---START---" << endl;
      cout << "CODE = " << temp_code << endl;
      cout << "---END---" << endl;
      cout << endl;
      char * buffer = new char[temp_code.size()+1];
      strcpy(buffer, temp_code.c_str()); 
      char * ptr = strtok(buffer, ",");
      int temp;
      code2 = new unsigned char[temp_code.size()+1];
      size_t index = 0;
      while(ptr != NULL) {
         temp = atoi(ptr);
         code2[index] = temp;
         index++;
         ptr = strtok(NULL, ",");
      }
      delete [] buffer;
   }

   void load(const string & infile) {
      ifstream in(infile.c_str(), ios::in);
      if(in) {
         string line;
         read_line(in, line);
         size_t comment;
         while((in)?(read_line(in, line)):(false)) {
            comment = line.find(";");
            if(comment != string::npos) {
               line = line.substr(0, comment);
            }
            if(line != "") {
               Instruction temp(line);
               code.push_back(temp);
            }
         }
         in.close();
      }
   }

   void load_code(const string & name) {
cout << "LOAD CODE: " << name << endl;
      bfstream fin(name.c_str(), ios::in | ios::binary);
      if(fin) {
/****
         size_t size = 0;
         fin >> size;
cout << "FUNCTION_VECTOR_SIZE=" << size << endl;
         for(size_t i = 0; i < size && fin; i++) {
            Function fun;
            fun.read(fin);
cout << "FUNCTION: \"" << fun.name << "\"" << endl;
cout << "ADDRESS: " << fun.addr << endl;
for(size_t i = 0; i < fun.addr_vec.size(); i++) {
cout << "USE ADDRESS: " << fun.addr_vec[i] << endl;
}
            function_vec.push_back(fun);
         }
         fin >> size;
cout << "VECTOR_SIZE=" << size << endl;
         size_t addr1 = 0;
         size_t addr2 = 0;
         for(size_t i = 0; i < size; i++) {
            fin >> addr1 >> addr2;
cout << addr1 << " -> " << addr2 << endl;
            addr_vec.push_back(addr1);
            addr_vec2.push_back(addr2);
         }
         fin >> size;
         code_size = size;
cout << "code_size = " << code_size << endl;
****/
         code_size = 69886;
         code2 = new unsigned char[code_size];
         fin.read((char*)code2, code_size);
         fin.close();
      }
   }

   void dump(ostream & out) {
      unsigned char ch;
      out << "---START---" << endl;
      out << "CODE = ";
      for(size_t i = 0; i < code_size; i++) {
         ch = code2[i];
         out << (unsigned int)ch << "," << flush;
      }
      out << endl;
      out << "---END---" << endl;
   }

   void emit(const string & outfile) {
      bfstream out(outfile.c_str(), ios::binary | ios::out | ios::trunc);
      out << (short)function_vec.size();
      for(size_t i = 0; i < function_vec.size(); i++) {
cout << "WRITE FUNCTION: \"" << function_vec[i].name << "\"" << endl;
         //for(size_t j = 0; j < function_vec[i].addr_vec.size(); j++) {
//cout << "ADDRESS[" << j << "]: " << function_vec[i].addr_vec[j] << endl;
            function_vec[i].write(out);
         //}
      } 
      out << (short)addr_vec.size();
      for(size_t i = 0; i < addr_vec.size(); i++) {
cout << "WRITE ADDRESS1: \"" << addr_vec[i] << "\"" << endl;
cout << "WRITE ADDRESS2: \"" << addr_vec2[i] << "\"" << endl;
         out << (short)addr_vec[i];
         out << (short)addr_vec2[i];
      }
      out.flush();
      out << (short)code_size;
      out.write((char*)code2, code_size);
      out.flush();
      out.close();
   }

   void link() {
      size_t * pointer;
      size_t addr;
      for(size_t i = 0; i < function_vec.size(); i++) {
         if(function_vec[i].fun_type == Function::F_EXTERN) {
            for(size_t j = 0; j < function_vec[i].addr_vec.size(); j++) {
               addr = function_vec[i].addr_vec[j];
               pointer = (size_t *)(code2 + addr);
               *pointer = (function_vec[i].addr); 
            }
         }
      }
      for(size_t i = 0; i < addr_vec.size(); i++) {
         addr = addr_vec[i];
         pointer = (size_t *)(code2 + addr);
         *pointer = ((size_t)code2 + addr_vec2[i]);
      }
   }
};

Program * program_ptr = 0;

inline
void 
add_extern_function(const string & name)
{
   program_ptr->add_function(Function(name,Function::F_EXTERN));
}

inline
void 
add_global_function(const string & name)
{
   //program_ptr->add_function(Function(name,Function::F_GLOBAL));
}

inline
void 
add_address(const size_t addr)
{
   program_ptr->add_address(addr);
}

inline
void 
add_address2(const size_t addr)
{
   program_ptr->add_address2(addr);
}


}

#endif /* _asm16 */
