#include <gtest/gtest.h>
#include <list>

extern "C" {
  #include "tlb.h"
}


typedef struct {
  PhysicalAddress pud, pmd, pte;
} TableIndexes;

typedef enum {
  kPud = 0,      //!< Page upper directory
  kPmd = 1,      //!< Page middle directory
  kPte = 2,      //!< Page table entry
  kPgd = 3,      //!< Page global directory
  kLastType      //!< Last element
} PageTableEntryType;

static const ProcessAddress kValidVirtAddr = 0x000060065777;

static TlbConverterInfo g_context = {0};
static KernelAddress *g_pgd_addr = NULL;
static KernelAddress *g_page = NULL;
static ProcessAddress g_virt_addr = kValidVirtAddr;
static std::list<PhysicalAddress *> g_tables_list;
static uint32_t g_is_physical_address_return = 1;
static PaTzResult g_physical_get_bytes_return = PA_TZ_SUCCESS;

extern "C" PaTzResult PhysicalGetBytes(PhysicalAddress phys, size_t size, void *out) {
  memcpy(out, (void *)phys, size);

  return g_physical_get_bytes_return;
}

extern "C" PhysicalAddress KernelVirtToPhys(KernelAddress kernel_virt) {
  return (PhysicalAddress)kernel_virt;
}

extern "C" uint32_t IsPhysicalAddress(PhysicalAddress phys) {
  return g_is_physical_address_return;
}

static inline PhysicalAddress AddrLevelOffset(PhysicalAddress vaddr, 
                                              PhysicalAddress level_shift) {
  return ((vaddr >> level_shift) & ADDR_LEVEL_MASK);
}

static TableIndexes CalcTableIndexes(PhysicalAddress virt_addr) {
  TableIndexes index = {0};
  index.pud = AddrLevelOffset(virt_addr, PUD_LEVEL_SHIFT);
  index.pmd = AddrLevelOffset(virt_addr, PMD_LEVEL_SHIFT);
  index.pte = AddrLevelOffset(virt_addr, PTE_LEVEL_SHIFT);

  return index;
}

static void AddTable(PhysicalAddress *table) {
  g_tables_list.push_front(table);
  memset(table, 0, PAGE_SIZE);
}

static void FreeTables() {
  while (!g_tables_list.empty()) {
    free(g_tables_list.front());
    g_tables_list.pop_front();
  }
}

static PhysicalAddress *AllocatePud() {
  PhysicalAddress *pud = (PhysicalAddress *)aligned_alloc(PAGE_SIZE, PAGE_SIZE);
  AddTable(pud);

  return pud;
}

static PhysicalAddress *AllocateTables(PhysicalAddress *pud, PhysicalAddress virt_addr) {
  TableIndexes indexes;
  PhysicalAddress *pmd = NULL;
  PhysicalAddress *pte = NULL;
  PhysicalAddress *page = NULL;

  indexes = CalcTableIndexes(virt_addr);

  if(!pud[indexes.pud]) {
    pud[indexes.pud] = (PhysicalAddress)(PhysicalAddress *)aligned_alloc(PAGE_SIZE, PAGE_SIZE);
    AddTable((PhysicalAddress *)pud[indexes.pud]);
  }

  pmd = (PhysicalAddress *)pud[indexes.pud];

  if(!pmd[indexes.pmd]) {
    pmd[indexes.pmd] = (PhysicalAddress)(PhysicalAddress *)aligned_alloc(PAGE_SIZE, PAGE_SIZE);
    AddTable((PhysicalAddress *)pmd[indexes.pmd]);
  }

  pte = (PhysicalAddress *)pmd[indexes.pmd];

  if(!pte[indexes.pte]) {
    pte[indexes.pte] = (PhysicalAddress)(PhysicalAddress *)aligned_alloc(PAGE_SIZE, PAGE_SIZE);
    AddTable((PhysicalAddress *)pte[indexes.pte]);
  }

  page = (PhysicalAddress *)pte[indexes.pte];

  return page;
}

static void UpdateContex(PhysicalAddress virt_addr, TlbConverterInfo *context) {
  TableIndexes indexes;

  context->pgd = (KernelAddress)g_pgd_addr;
  context->phys_pgd = (PhysicalAddress)g_pgd_addr;

  indexes = CalcTableIndexes(virt_addr);
  context->index_table[kPud] = indexes.pud;
  context->index_table[kPmd] = indexes.pmd;
  context->index_table[kPte] = indexes.pte;

  context->buffer[kPud][indexes.pud] = g_pgd_addr[indexes.pud];
  context->buffer[kPmd][indexes.pmd] = 
    ((PhysicalAddress *)context->buffer[kPud][indexes.pud])[indexes.pmd];
  context->buffer[kPte][indexes.pte] = 
    ((PhysicalAddress *)context->buffer[kPmd][indexes.pmd])[indexes.pte];
}

class TaskAddressToPhysicalTest : public ::testing::Test {
protected:
  virtual void SetUp() {
    g_virt_addr = kValidVirtAddr;
    g_pgd_addr = AllocatePud();
    g_page = AllocateTables(g_pgd_addr, g_virt_addr);
    g_context.pgd = 0x0;
    g_context.phys_pgd = 0x0;
    g_is_physical_address_return = 1;
    g_physical_get_bytes_return = PA_TZ_SUCCESS;
    memset(g_context.index_table, 0, sizeof(g_context.index_table));
    memset(g_context.buffer, 0, sizeof(g_context.buffer));
  }

  virtual void TearDown() {
    FreeTables();
  }
};


TEST_F(TaskAddressToPhysicalTest, AllParamAreValid) {
  PaTzResult result;
  PhysicalAddress phys_addr = 0x0;

  result = TaskAddressToPhysical(&g_context, (KernelAddress)g_pgd_addr, g_virt_addr, 
                                 &phys_addr, NULL);
  EXPECT_TRUE(PA_TZ_SUCCESS == result);

  phys_addr = (phys_addr & ~(PAGE_SIZE - 1));
  EXPECT_TRUE(phys_addr == (PhysicalAddress)g_page);
}

TEST_F(TaskAddressToPhysicalTest, AllParamAreValidAndNonNullFlags) {
  PaTzResult result;
  PhysicalAddress phys_addr = 0x0;

  AccessPermissionFlags flags;
  result = TaskAddressToPhysical(&g_context, (KernelAddress)g_pgd_addr, g_virt_addr,
                                 &phys_addr, &flags);

  EXPECT_TRUE(PA_TZ_SUCCESS == result);

  phys_addr = (phys_addr & ~(PAGE_SIZE - 1));
  EXPECT_EQ((PhysicalAddress)g_page, phys_addr);
}

TEST_F(TaskAddressToPhysicalTest, MinVirtAddr) {
  PaTzResult result;
  g_virt_addr = 0x000000001000;  
  PhysicalAddress phys_addr = 0x0;

  FreeTables();
  g_pgd_addr = AllocatePud();
  g_page = AllocateTables(g_pgd_addr, g_virt_addr);

  result = TaskAddressToPhysical(&g_context, (KernelAddress)g_pgd_addr, g_virt_addr, 
                                 &phys_addr, NULL);
  EXPECT_TRUE(PA_TZ_SUCCESS == result); 

  phys_addr = (phys_addr & ~(PAGE_SIZE - 1));
  EXPECT_TRUE(phys_addr == (PhysicalAddress)g_page);
}

TEST_F(TaskAddressToPhysicalTest, MaxVirtAddr) {
  PaTzResult result;
  g_virt_addr = 0x007FFFFFFFFC;  
  PhysicalAddress phys_addr = 0x0;

  FreeTables();
  g_pgd_addr = AllocatePud();
  g_page = AllocateTables(g_pgd_addr, g_virt_addr);

  result = TaskAddressToPhysical(&g_context, (KernelAddress)g_pgd_addr, g_virt_addr, 
                                 &phys_addr, NULL);
  EXPECT_TRUE(PA_TZ_SUCCESS == result);

  phys_addr = (phys_addr & ~(PAGE_SIZE - 1));
  EXPECT_TRUE(phys_addr == (PhysicalAddress)g_page);
}

TEST_F(TaskAddressToPhysicalTest, WrongPud) {
  PaTzResult result;
  PhysicalAddress wrong_pud = g_virt_addr ^ (1ULL << 31);
  PhysicalAddress phys_addr = 0x0;

  result = TaskAddressToPhysical(&g_context, (KernelAddress)g_pgd_addr, wrong_pud,
                                 &phys_addr, NULL);
  EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskAddressToPhysicalTest, WrongPmd) {
  PaTzResult result;
  PhysicalAddress wrong_pmd = g_virt_addr ^ (1ULL << 22);
  PhysicalAddress phys_addr = 0x0;

  result = TaskAddressToPhysical(&g_context, (KernelAddress)g_pgd_addr, wrong_pmd, 
                                 &phys_addr, NULL);
  EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskAddressToPhysicalTest, WrongPte) {
  PaTzResult result;
  PhysicalAddress wrong_pte = g_virt_addr ^ (1ULL << 13);
  PhysicalAddress phys_addr = 0x0;

  result = TaskAddressToPhysical(&g_context, (KernelAddress)g_pgd_addr, wrong_pte, 
                                 &phys_addr, NULL);
  EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskAddressToPhysicalTest, ContextNull) {
  PaTzResult result;
  TlbConverterInfo *null_context = NULL;
  PhysicalAddress phys_addr = 0x0;

  result = TaskAddressToPhysical(null_context, (KernelAddress)g_pgd_addr, g_virt_addr, 
                                 &phys_addr, NULL);

  EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskAddressToPhysicalTest, PgdAddrZero) {
  PaTzResult result;
  KernelAddress zero_pgd_addr = 0x0;
  PhysicalAddress phys_addr = 0x0;

  result = TaskAddressToPhysical(&g_context, zero_pgd_addr, g_virt_addr, 
                                 &phys_addr, NULL);
  EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskAddressToPhysicalTest, VirtAddrZero) {
  PaTzResult result;
  ProcessAddress zero_virt_addr = 0x0;
  PhysicalAddress phys_addr = 0x0;

  FreeTables();
  g_pgd_addr = AllocatePud();
  g_page = AllocateTables(g_pgd_addr, zero_virt_addr);

  result = TaskAddressToPhysical(&g_context, (KernelAddress)g_pgd_addr, zero_virt_addr, 
                                 &phys_addr, NULL);
  EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskAddressToPhysicalTest, PhysAddrNull) {
  PaTzResult result;
  PhysicalAddress *null_phys_addr = NULL;

  result = TaskAddressToPhysical(&g_context, (KernelAddress)g_pgd_addr, g_virt_addr, 
                                 null_phys_addr, NULL);
  EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskAddressToPhysicalTest, AddrIsNotPhysical) {
  PaTzResult result;
  PhysicalAddress phys_addr = 0x0;

  g_is_physical_address_return = 0;

  result = TaskAddressToPhysical(&g_context, (KernelAddress)g_pgd_addr, g_virt_addr, 
                                 &phys_addr, NULL);
  EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskAddressToPhysicalTest, FailedPhysicalGetBytes) {
  PaTzResult result;
  PhysicalAddress phys_addr = 0x0;

  g_physical_get_bytes_return = PA_TZ_GENERAL_ERROR;

  result = TaskAddressToPhysical(&g_context, (KernelAddress)g_pgd_addr, g_virt_addr, 
                                 &phys_addr, NULL);
  EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}


TEST_F(TaskAddressToPhysicalTest, SamePudCorrectAddr) {
  PaTzResult result;
  PhysicalAddress phys_addr = 0x0;

  g_virt_addr = 0x001C647E6148;
  ProcessAddress same_pud_addr = 0x001C77864355;

  FreeTables();
  g_pgd_addr = AllocatePud();
  g_page = AllocateTables(g_pgd_addr, same_pud_addr);
  g_page = AllocateTables(g_pgd_addr, g_virt_addr);

  UpdateContex(g_virt_addr, &g_context);
  UpdateContex(same_pud_addr, &g_context);
 
  result = TaskAddressToPhysical(&g_context, (KernelAddress)g_pgd_addr, g_virt_addr, 
                                 &phys_addr, NULL);
  EXPECT_TRUE(PA_TZ_SUCCESS == result);

  phys_addr = (phys_addr & ~(PAGE_SIZE - 1));
  EXPECT_TRUE(phys_addr == (PhysicalAddress)g_page);
}

TEST_F(TaskAddressToPhysicalTest, SamePudWrongAddr) {
  PaTzResult result;
  PhysicalAddress phys_addr = 0x0;

  g_virt_addr = 0x001C647E6148; 
  ProcessAddress same_pud_addr = 0x001C77864355;
  ProcessAddress wrong_virt_addr = g_virt_addr ^ (1ULL << 12);

  FreeTables();
  g_pgd_addr = AllocatePud();
  g_page = AllocateTables(g_pgd_addr, same_pud_addr);
  g_page = AllocateTables(g_pgd_addr, g_virt_addr);

  UpdateContex(g_virt_addr, &g_context);
  UpdateContex(same_pud_addr, &g_context);
 
  result = TaskAddressToPhysical(&g_context, (KernelAddress)g_pgd_addr, wrong_virt_addr, 
                                 &phys_addr, NULL);
  EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskAddressToPhysicalTest, SamePudAndPmdCorrectAddr) {
  PaTzResult result;
  PhysicalAddress phys_addr = 0x0;

  ProcessAddress same_pud_pmd_addr = 0x005f07aa6150;
  g_virt_addr = 0x005F07AF6465;

  FreeTables();
  g_pgd_addr = AllocatePud();
  g_page = AllocateTables(g_pgd_addr, same_pud_pmd_addr);
  g_page = AllocateTables(g_pgd_addr, g_virt_addr);

  UpdateContex(g_virt_addr, &g_context);
  UpdateContex(same_pud_pmd_addr, &g_context);
 
  result = TaskAddressToPhysical(&g_context, (KernelAddress)g_pgd_addr, g_virt_addr, 
                                 &phys_addr, NULL);
  EXPECT_TRUE(PA_TZ_SUCCESS == result);

  phys_addr = (phys_addr & ~(PAGE_SIZE - 1));
  EXPECT_TRUE(phys_addr == (PhysicalAddress)g_page);
}

TEST_F(TaskAddressToPhysicalTest, SamePudAndPmdWrongAddr) {
  PaTzResult result;
  PhysicalAddress phys_addr = 0x0;
  ProcessAddress same_pud_pmd_addr = 0x005f07aa6150;
  g_virt_addr = 0x005F07AF6465;
  ProcessAddress wrong_virt_addr = g_virt_addr ^ (1ULL << 12);

  FreeTables();
  g_pgd_addr = AllocatePud();
  g_page = AllocateTables(g_pgd_addr, same_pud_pmd_addr);
  g_page = AllocateTables(g_pgd_addr, g_virt_addr);

  UpdateContex(g_virt_addr, &g_context);
  UpdateContex(same_pud_pmd_addr, &g_context);
 
  result = TaskAddressToPhysical(&g_context, (KernelAddress)g_pgd_addr, wrong_virt_addr, 
                                 &phys_addr, NULL);
  EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskAddressToPhysicalTest, SameVirtAddr) {
  PaTzResult result;
  PhysicalAddress phys_addr = 0x0;

  UpdateContex(g_virt_addr, &g_context);
 
  result = TaskAddressToPhysical(&g_context, (KernelAddress)g_pgd_addr, g_virt_addr, 
                                 &phys_addr, NULL);
  EXPECT_TRUE(PA_TZ_SUCCESS == result);

  phys_addr = (phys_addr & ~(PAGE_SIZE - 1));
  EXPECT_TRUE(phys_addr == (PhysicalAddress)g_page);
}