#include "asm/chunk.h"
#include "asm/assembler.h"
#include "asm/parse.h"
#include "asm/segment.h"
#include "asm/intel64/all_ops.h"

#include "minicc.h"

#include <boost/algorithm/string.hpp>

#include "gmock/gmock.h"
#include "gtest/gtest.h"

#include <algorithm>
#include <cctype>
#include <deque>
#include <map>
#include <memory>
#include <string>
#include <system_error>
#include <utility>
#include <vector>

using namespace std::string_literals;
namespace fs = std::filesystem;

class AsmTest: public ::testing::Test
{
protected:
 AsmTest() {
  //debug = true;
 }
 ~AsmTest() {
 }
 void SetUp(){
 }
 void TearDown(){
 }
};

class AsmParseTest: public ::testing::Test
{
protected:
 AsmParseTest() {
  //debug = true;
 }
 ~AsmParseTest() {
 }
 void SetUp(){
 }
 void TearDown(){
 }
};

TEST_F(AsmTest, Intel64_add) {
 Segment segment;
 Asm::Args args{{Asm::Args::Register32("eax"), Asm::Args::Immediate32(1)}};
 segment.push_back(makeOp("add", args));

 ASSERT_EQ(segment.size(), 1);
 ASSERT_EQ(segment.getCode(), std::vector<uint8_t>({0x05, 0x01, 0x00, 0x00, 0x00}));
}

TEST_F(AsmTest, Intel64_int_0) {
 Segment segment;
 Asm::Args args{{Asm::Args::Immediate8(0)}};
 segment.push_back(makeOp("int", args));

 ASSERT_EQ(segment.size(), 1);
 ASSERT_EQ(segment.getCode(), std::vector<uint8_t>{0xCE});
}

TEST_F(AsmTest, Intel64_int_1) {
 Segment segment;
 Asm::Args args{{Asm::Args::Immediate8(1)}};
 segment.push_back(makeOp("int", args));

 ASSERT_EQ(segment.size(), 1);
 ASSERT_EQ(segment.getCode(), std::vector<uint8_t>{0xF1});
}

TEST_F(AsmTest, Intel64_int_5) {
 Segment segment;
 Asm::Args args{{Asm::Args::Immediate8(5)}};
 segment.push_back(makeOp("int", args));

 ASSERT_EQ(segment.size(), 1);
 ASSERT_EQ(segment.getCode(), std::vector<uint8_t>({0xCD, 0x05}));
}

TEST_F(AsmTest, Intel64_nop) {
 Segment segment;
 segment.push_back(makeOp("nop"));

 ASSERT_EQ(segment.size(), 1);
 ASSERT_EQ(segment.getCode(), std::vector<uint8_t>{0x90});
}

TEST_F(AsmTest, Intel64_ret) {
 Segment segment;
 segment.push_back(makeOp("ret"));

 ASSERT_EQ(segment.size(), 1);
 ASSERT_EQ(segment.getCode(), std::vector<uint8_t>{0xC3});
}

TEST_F(AsmTest, Intel64_multiple) {
 Segment segment;

 segment.push_back(makeOp("nop"));
 Asm::Args args0{{Asm::Args::Immediate8(5)}};
 segment.push_back(makeOp("int", args0));
 segment.push_back(makeOp("ret"));
 segment.push_back(makeLabel("data1"));
 segment.push_back(makeOp("ret"));
 Asm::Args args1{{Asm::Args::Label("data1")}};
 segment.push_back(makeOp("jmp", args1));
 segment.push_back(makeData({1, 2, 3}));

 segment.insertAddresses();

 ASSERT_EQ(segment.size(), 7);
 ASSERT_EQ(segment.getCode(), std::vector<uint8_t>(
  {
   0x90, // nop
   0xCD, 0x05, // int 5
   0xC3, // ret
   // data1:
   0xC3, // ret
   0xE9, 0xFF, 0xFF, 0xFF, 0xFF, // jmp data1
   0x01, 0x02, 0x03 // data
  }));
 
 segment.optimize();
 
 ASSERT_EQ(segment.size(), 7);
 ASSERT_EQ(segment.getCode(), std::vector<uint8_t>(
  {
   0x90, // nop
   0xCD, 0x05, // int 5
   0xC3, // ret
   // data1:
   0xC3, // ret
   0xEB, 0xFF, // jmp data1
   0x01, 0x02, 0x03 // data
  }));
}

TEST_F(AsmParseTest, parse_empty) {
 std::vector<std::shared_ptr<Chunk>> chunks0{parseAsm("")};
 ASSERT_EQ(chunks0.size(), 0);

 std::vector<std::shared_ptr<Chunk>> chunks1{parseAsm("\n\n")};
 ASSERT_EQ(chunks1.size(), 0);
 
 std::vector<std::shared_ptr<Chunk>> chunks2{parseAsm("\n\n")};
 ASSERT_EQ(chunks2.size(), 0);
}

TEST_F(AsmParseTest, parse_op_0) {
 std::vector<std::shared_ptr<Chunk>> chunks0{parseAsm("nop")};
 ASSERT_EQ(chunks0.size(), 1);
}

TEST_F(AsmParseTest, parse_op_1) {
 std::vector<std::shared_ptr<Chunk>> chunks1{parseAsm("neg edi")};
 ASSERT_EQ(chunks1.size(), 1);
}

TEST_F(AsmParseTest, parse_op_2) {
 std::vector<std::shared_ptr<Chunk>> chunks2{parseAsm("add eax, edx")};
 ASSERT_EQ(chunks2.size(), 1);
}

TEST_F(AsmParseTest, parse_op_3) {
 std::vector<std::shared_ptr<Chunk>> chunks3{parseAsm("add eax, 3")};
 ASSERT_EQ(chunks3.size(), 1);
}

TEST_F(AsmParseTest, parse_op_4) {
 std::vector<std::shared_ptr<Chunk>> chunks4{parseAsm("add [rdi], 3")};
 ASSERT_EQ(chunks4.size(), 1);
}

TEST_F(AsmParseTest, parse_op_4_error) {
 ASSERT_THROW(parseAsm("add [edi], 3"), std::runtime_error);
}

TEST_F(AsmParseTest, parse_op_5) {
 std::vector<std::shared_ptr<Chunk>> chunks5{parseAsm("add byte ptr [rdi], 3")};
 ASSERT_EQ(chunks5.size(), 1);
}

TEST_F(AsmParseTest, parse_op_5_error) {
 ASSERT_THROW(parseAsm("add byte ptr [edi], 3"), std::runtime_error);
}

TEST_F(AsmParseTest, parse_op_6) {
 std::vector<std::shared_ptr<Chunk>> chunks6{parseAsm("add dword ptr[rdi], 3")};
 ASSERT_EQ(chunks6.size(), 1);
}

TEST_F(AsmParseTest, parse_op_6_error) {
 ASSERT_THROW(parseAsm("add dword ptr[al], 3"), std::runtime_error);
}

TEST_F(AsmParseTest, parse_op_7) {
 std::vector<std::shared_ptr<Chunk>> chunks7{parseAsm("add qword ptr[rdi], 3")};
 ASSERT_EQ(chunks7.size(), 1);
}

TEST_F(AsmParseTest, parse_label) {
 std::vector<std::shared_ptr<Chunk>> chunks0{parseAsm("label1:")};
 ASSERT_EQ(chunks0.size(), 1);

 std::vector<std::shared_ptr<Chunk>> chunks1{parseAsm("label2: add eax, 3")};
 ASSERT_EQ(chunks1.size(), 2);
}

TEST_F(AsmParseTest, parse_multiline) {
 std::vector<std::shared_ptr<Chunk>> chunks1{parseAsm("add eax, 3\n")};
 ASSERT_EQ(chunks1.size(), 1);
 
 std::vector<std::shared_ptr<Chunk>> chunks2{parseAsm("label2: add eax, 3\n   mul rdx")};
 ASSERT_EQ(chunks2.size(), 3);
 
 std::vector<std::shared_ptr<Chunk>> chunks3{parseAsm("label2: add eax, 3\n   mul rdx\n")};
 ASSERT_EQ(chunks3.size(), 3);
 
 std::vector<std::shared_ptr<Chunk>> chunks4{parseAsm("label2: add eax, 3\n\n\n   mul rdx")};
 ASSERT_EQ(chunks4.size(), 3);
}

TEST_F(AsmParseTest, parse_comment) {
 std::vector<std::shared_ptr<Chunk>> chunks0{parseAsm("    ; Comment 1")};
 ASSERT_EQ(chunks0.size(), 0);
 
 std::vector<std::shared_ptr<Chunk>> chunks1{parseAsm("label2: add eax, 3 ; A comment")};
 ASSERT_EQ(chunks1.size(), 2);
 
 std::vector<std::shared_ptr<Chunk>> chunks3{parseAsm("label2: add eax, 3 // Another comment")};
 ASSERT_EQ(chunks3.size(), 2);
 
 std::vector<std::shared_ptr<Chunk>> chunks4{parseAsm("label2: add eax, 3 // Another comment\nadd rax, rdi")};
 ASSERT_EQ(chunks4.size(), 3);
 
 std::vector<std::shared_ptr<Chunk>> chunks5{parseAsm("label2: add eax, 3 // Another comment\n \t  add rax, rdi")};
 ASSERT_EQ(chunks5.size(), 3);
}

TEST_F(AsmParseTest, parse_error) {
 ASSERT_THROW(parseAsm(" add label1:"), std::runtime_error);
 
 ASSERT_THROW(parseAsm(" add eax , "), std::runtime_error);
 ASSERT_THROW(parseAsm(" add eax,"), std::runtime_error);
}