#include "memory.h"
#include "config.h"

#include "driver_log.h"

#include <tee_internal_api.h>

static const PhysicalAddress kPaConfigSymOffsetTextOffset = 0x28;
static const PhysicalAddress kKernelMagicTextOffset = 0x38;
#ifdef CONFIG_KERNEL510
static const PhysicalAddress kTextSymPhysOffset = 0x0000000000000000ULL;
#else
static const PhysicalAddress kTextSymPhysOffset = 0x0000000000080000ULL;
#endif

PaTzResult PlatformMapRegion(PhysicalAddress phys, size_t size,
                             MemoryAccessType type, void **virt) {
  uintptr_t real_virt = 0;
  uint32_t offset = PageOffset(phys);

  // Calculating size aligned on page size
  size_t size_aligned = AlignToPageUp(size + offset);

  PhysicalAddress phys_aligned = AlignToPageDown(phys);

  if (PlatformSysMap(phys_aligned, size_aligned,
                     type, (void **)&real_virt) != PA_TZ_SUCCESS) {
    LOG_E("Failed mapping a region\n");
    return PA_TZ_GENERAL_ERROR;
  }

  // Finally to set start address of region pages
  *virt = (void *)(real_virt + offset);

  return PA_TZ_SUCCESS;
}

PaTzResult PlatformUnmapRegion(void *virt, size_t size) {
  uint32_t offset = PageOffset((uintptr_t)virt);

  // Calculating size aligned on page size
  uint32_t size_aligned = AlignToPageUp(size + offset);

  if (PlatformSysUnmap(virt, size_aligned) != PA_TZ_SUCCESS) {
    LOG_E("Failed unmapping a region\n");
    return PA_TZ_GENERAL_ERROR;
  }

  return PA_TZ_SUCCESS;
}

PaTzResult PhysicalGetBytes(PhysicalAddress phys, size_t size, void *out) {
  if (!out) {
    LOG_E("Invalid arguments.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  void *addr = NULL;
  PaTzResult result = PlatformSysMap(phys, size, kMemoryAccessRead, &addr);
  if (result != PA_TZ_SUCCESS) {
    LOG_E("Can not map paddr to read.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  if (addr) {
    TEE_MemMove(out, addr, size);
  } else {
    LOG_E("Mapping return incorrect address.\n");
  }

  result = PlatformSysUnmap(addr, size);
  if (result != PA_TZ_SUCCESS) {
    LOG_E("Can not unmap\n");
  }

  return result;
}

#define BIT(nr)         (1ULL << (nr))

/**
 * sign_extend64 - sign extend a 64-bit value using specified bit as sign-bit
 * @value: value to sign extend
 * @index: 0 based bit index (0<=index<64) to sign bit
 */
static inline int64_t sign_extend64(uint64_t value, int index)
{
	uint8_t shift = 63 - index;
	return (int64_t)(value << shift) >> shift;
}

/**
 * ARMv8.5 based processors introduce the Memory Tagging Extension (MTE)
 * feature. MTE is built on top of the ARMv8.0 virtual address tagging TBI (Top
 * Byte Ignore) feature and allows software to access a 4-bit allocation tag for
 * each 16-byte granule in the physical address space. Such memory range must be
 * mapped with the Normal-Tagged memory attribute. A logical tag is derived from
 * bits 59-56 of the virtual address used for the memory access.
 * Solution based on arch/arm64/include/asm/memory.h (5.10)
**/
static KernelAddress MaskHighBits(KernelAddress virt) {
  return (KernelAddress)sign_extend64((uint64_t)(virt), 55);
}

PhysicalAddress KernelVirtToPhys(KernelAddress kernel_virt) {
  if (!kernel_virt) {
    return kernel_virt;
  }

  kernel_virt = MaskHighBits(kernel_virt);

  return (
#ifndef CONFIG_KERNEL54
    kernel_virt & BIT(GetConfig()->memory_conf.va_bits - 1) ?
      (kernel_virt & ~GetConfig()->memory_conf.page_offset) +
                                  KernelGetConfigPhysOffset() :
      (kernel_virt - GetConfig()->memory_conf.kimage_voffset));
#else
    (!(kernel_virt & BIT(GetConfig()->memory_conf.va_bits - 1))) ?
      (kernel_virt - GetConfig()->memory_conf.page_offset) +
                                  KernelGetConfigPhysOffset() :
      (kernel_virt - GetConfig()->memory_conf.kimage_voffset));
#endif
}

//kernel virtual addresses start at va_start and end at MAX_UINT64
uint32_t IsKernelVirtualAddress(KernelAddress kernel_virt) {
  return MaskHighBits(kernel_virt) >= GetConfig()->memory_conf.va_start ? 1 : 0;
}

uint32_t IsPhysicalAddress(PhysicalAddress phys) {
  for (size_t i = 0; i < GetConfig()->memory_conf.sys_ram_ranges_num; ++i) {
    if (phys >= GetConfig()->memory_conf.sys_ram_ranges[i].start &&
        phys < GetConfig()->memory_conf.sys_ram_ranges[i].end) {
          return 1;
    }
  }
  return 0;
}

PhysicalAddress KernelGetPaConfigOffsetSymAddress(void) {
  return KernelGetTextSymAddress() + kPaConfigSymOffsetTextOffset;
}

PhysicalAddress KernelGetTextSymAddress(void) {
  return KernelGetStartPhysAddr() + kTextSymPhysOffset;
}

PhysicalAddress KernelGetKernelMagicAddress(void) {
  return KernelGetTextSymAddress() + kKernelMagicTextOffset;
}