#include "parse.h"

#include "asm/assembler.h"

#include <boost/algorithm/string.hpp>

#include <exception>
#include <regex>
#include <unordered_set>

using namespace std::string_literals;

namespace {
 
 std::unordered_set<std::string> reg8 {
  "al", "ah",
  "bl", "bh",
  "cl", "ch",
  "dl", "dh",
 };

 std::unordered_set<std::string> reg32 {
  "eax", "esp",
  "ebx", "ebp",
  "ecx", "esi",
  "edx", "edi",
 };

 std::unordered_set<std::string> reg64 {
  "rax", "rsp",
  "rbx", "rbp",
  "rcx", "rsi",
  "rdx", "rdi",
 };

 // skip optional whitespace
 void parseWhitespace(const std::string& asm_code, size_t& pos) {
  std::regex re_whitespace("[ \\t]+", std::regex_constants::ECMAScript);

  std::smatch match;
  if (std::regex_search(asm_code.cbegin() + pos, asm_code.cend(), match, re_whitespace, std::regex_constants::match_continuous)) {
   pos += match[0].length();
  }
 }

 // parse optional label
 bool parseLabel(const std::string& asm_code, size_t& pos, std::string& result) {
  parseWhitespace(asm_code, pos);
  
  std::regex re_label("([[:alpha:]]([[:alnum:]])+):", std::regex_constants::ECMAScript);

  std::smatch match;
  if (std::regex_search(asm_code.cbegin() + pos, asm_code.cend(), match, re_label, std::regex_constants::match_continuous)) {
   pos += match[0].length();
   result = match[1];
   return true;
  }

  return false;
 }

 // parse optional mnemonic
 // return true iff mnemonic found
 bool parseMnemonic(const std::string& asm_code, size_t& pos, std::string& result) {
  parseWhitespace(asm_code, pos);
  
  std::regex re_mnemonic("[[:alpha:]]([[:alnum:]])+", std::regex_constants::ECMAScript);

  std::smatch match;
  if (std::regex_search(asm_code.cbegin() + pos, asm_code.cend(), match, re_mnemonic, std::regex_constants::match_continuous)) {
   std::string name {boost::algorithm::to_lower_copy(match[0].str())};
   pos += name.size();
   result = name;
   return true;
  }

  return false;
 }
 
 bool parseRegister8(const std::string& asm_code, size_t& pos, std::any& result) {
  parseWhitespace(asm_code, pos);
  
  std::regex re_name("[[:alpha:]]+", std::regex_constants::ECMAScript);

  std::smatch match;
  if (std::regex_search(asm_code.cbegin() + pos, asm_code.cend(), match, re_name, std::regex_constants::match_continuous)) {
   std::string name {boost::algorithm::to_lower_copy(match[0].str())};
   if (reg8.contains(name)) {
    pos += name.size();
    result = Asm::Args::Register8(name);
    return true;
   }
  }

  return false;
 }

 bool parseRegister32(const std::string& asm_code, size_t& pos, std::any& result) {
  parseWhitespace(asm_code, pos);
  
  std::regex re_name("[[:alpha:]]+", std::regex_constants::ECMAScript);

  std::smatch match;
  if (std::regex_search(asm_code.cbegin() + pos, asm_code.cend(), match, re_name, std::regex_constants::match_continuous)) {
   std::string name {boost::algorithm::to_lower_copy(match[0].str())};
   if (reg32.contains(name)) {
    pos += name.size();
    result = Asm::Args::Register32(name);
    return true;
   }
  }

  return false;
 }

 bool parseRegister64(const std::string& asm_code, size_t& pos, std::any& result) {
  parseWhitespace(asm_code, pos);
  
  std::regex re_name("[[:alpha:]]+", std::regex_constants::ECMAScript);

  std::smatch match;
  if (std::regex_search(asm_code.cbegin() + pos, asm_code.cend(), match, re_name, std::regex_constants::match_continuous)) {
   std::string name {boost::algorithm::to_lower_copy(match[0].str())};
   if (reg64.contains(name)) {
    pos += name.size();
    result = Asm::Args::Register64(name);
    return true;
   }
  }

  return false;
 }

 bool parseImmediate32(const std::string& asm_code, size_t& pos, std::any& result) {
  parseWhitespace(asm_code, pos);
  
  std::regex re_name("[[:digit:]]+|0x[[:xdigit:]]+", std::regex_constants::ECMAScript);

  std::smatch match;
  if (std::regex_search(asm_code.cbegin() + pos, asm_code.cend(), match, re_name, std::regex_constants::match_continuous)) {
   int32_t value{};
   try {
    value = stoll(match[0]);
   } catch (...) {
    throw std::runtime_error("Assembler parse error: Bad Immediate: "s + match[0].str());
   }
   pos += match[0].length();
   result = Asm::Args::Immediate32(static_cast<uint32_t>(value));
   return true;
  }

  return false;
 }

 // parse optional single operand
 bool parseOperand(const std::string& asm_code, size_t& pos, std::any& result) {
  parseWhitespace(asm_code, pos);

  if (parseRegister8(asm_code, pos, result))
   return true;
  if (parseRegister32(asm_code, pos, result))
   return true;
  if (parseRegister64(asm_code, pos, result))
   return true;

  if (parseImmediate32(asm_code, pos, result))
   return true;

  return false;
 }

 // parse optional multiple operands, separated by commas
 void parseOperands(const std::string& asm_code, size_t& pos, Asm::Args& result) {
  std::any operand;
  if (parseOperand(asm_code, pos, operand)) {
   result.push_back(operand);
   parseWhitespace(asm_code, pos);
   while (pos < asm_code.size() && asm_code[pos] == ',') {
    pos++;
    if (parseOperand(asm_code, pos, operand)) {
     result.push_back(operand);
    } else {
     throw std::runtime_error("Assembler error: expected operand after comma");
    }
    parseWhitespace(asm_code, pos);
   }
  }
 }

 // parse optional comment
 void parseComment(const std::string& asm_code, size_t& pos) {
  parseWhitespace(asm_code, pos);
  
  std::regex re_comment("(;|//).*", std::regex_constants::ECMAScript);

  std::smatch match;
  if (std::regex_search(asm_code.cbegin() + pos, asm_code.cend(), match, re_comment, std::regex_constants::match_continuous)) {
   pos += match[0].length();
  }
 }

 // parse end of line (or whole code)
 bool parseEol(const std::string& asm_code, size_t& pos) {
  parseWhitespace(asm_code, pos);

  if (pos < asm_code.size() && asm_code[pos] != 0x0a && asm_code[pos] != 0x0d)
   return false; // this is the only case where parseEol() doesn't work

  while (pos < asm_code.size()) {
   char c { asm_code[pos] };
   if (c == 0x0a || c == 0x0d) {
    pos++;
   } else {
    break;
   }
  }

  return true;
 }

 // parse single line
 void parseLine(const std::string& asm_code, size_t& pos, std::vector<std::shared_ptr<Chunk>>& result) {
  // all optional:
  // label: mnemonic operands... ;comment <eol>

  std::string result_string;
  if (parseLabel(asm_code, pos, result_string))
   result.emplace_back(std::make_shared<Label>(result_string));

  if (parseMnemonic(asm_code, pos, result_string)) {
   Asm::Args args;
   parseOperands(asm_code, pos, args);
   result.emplace_back(makeOp(result_string, args));
  }

  parseComment(asm_code, pos);

  if (!parseEol(asm_code, pos))
   throw std::runtime_error("Assembler error at pos "s + std::to_string(pos));
 }

} // namespace

std::vector<std::shared_ptr<Chunk>> parseAsm(const std::string& asm_code)
{
 std::vector<std::shared_ptr<Chunk>> result;
 size_t pos{0};

 while (pos != asm_code.size()) {
  parseLine(asm_code, pos, result);
 }

 return result;
}