#include "uart.h"
#include "strings.h"
#include "psr.h"
#include "translation_table_descriptors.h"
#include "cp_regs.h"

extern char __end;

extern char
  _binary_PL_0_test_img_start,
  _binary_PL_0_test_img_end,
  _binary_PL_0_test_img_size;

void enabling_code_from_the_net(void);
  
void kernel_main(uint32_t r0, uint32_t r1, uint32_t atags)
{
  // Declare as unused
  (void) r0;
  (void) r1;
  (void) atags;

  uart_init();

  // When we attach screen session after loading kernel with socat
  // we miss kernel's greeting... So we'll make the kernel wait for
  // one char we're going to send from within screen
  uart_getc();
  
  uart_puts("Hello, kernel World!\r\n");

  uint32_t ID_MMFR0;
  // get contents of coprocessor register to check for paging support
  asm("mrc p15, 0, %0, c0, c1, 4" : "=r" (ID_MMFR0));

  char *paging;

  switch(ID_MMFR0 & 0xf) /* lowest 4 bits indicate VMSA support */
    {
    case 0 : paging = "no paging\n\r"; break;
    case 1 : paging = "implementation defined paging\n\r"; break;
    case 2 : paging = "VMSAv6, with cache and TLB type registers\n\r"; break;
    case 3 : paging = "VMSAv7, with support for remapping and access flag\n\r"; break;
    case 4 : paging = "VMSAv7 with PXN bit supported\n\r"; break;
    case 5 : paging = "VMSAv7, PXN and long format descriptors. EPAE is supported.\n\r"; break;
    default : paging = "?_? unknown paging ?_?\n\r";
    }

  uart_puts(paging);

  // get content of current program status register to check the current
  // processor mode
  PSR_t CPSR = read_CPSR();

  char *mode_name;
  
  switch(CPSR.fields.PSR_MODE_4_0)
    {
    case MODE_USER       : mode_name = "User (PL0)\r\n"; break;
    case MODE_FIQ        : mode_name = "FIQ (PL1)\r\n"; break;
    case MODE_IRQ        : mode_name = "IRQ (PL1)\r\n"; break;
    case MODE_SUPERVISOR : mode_name = "Supervisor (PL1)\r\n"; break;
    case MODE_MONITOR    : mode_name = "Monitor (PL1)\r\n"; break;
    case MODE_ABORT      : mode_name = "Abort (PL1)\r\n"; break;
    case MODE_HYPERVISOR : mode_name = "Hyp (PL2)\r\n"; break;
    case MODE_UNDEFINED  : mode_name = "Undefined (PL1)\r\n"; break;
    case MODE_SYSTEM     : mode_name = "System (PL1)\r\n"; break;
    default : mode_name = "Unknown mode\r\n"; break;
    }

  uart_puts("current mode: ");
  uart_puts(mode_name);

  uart_puts("setting mode to system (PL1)...\r\n");
  set_system_mode();

  char bits[33]; // for printing uint32_t bit values
  
  // compute translation table base address
  // translation table shall start at first 2^14-bytes aligned
  // address after the kernel image
  uint32_t kernel_end = (uint32_t) &__end;
  uint32_t translation_table_base =
    ((kernel_end - 1) & ~((uint32_t) 0x3fff)) + (uint32_t) 0x4000;

  uint32_to_bits(translation_table_base, bits);
  uart_puts("binary representation of chosen"
	    " lvl1 translation table address: ");
  uart_puts(bits); uart_puts("\n\r");
 
  // flat map all memory
  uart_puts("preparing translation table\n\r");
  short_descriptor_lvl1_t volatile *translation_table =
    (short_descriptor_lvl1_t*) translation_table_base;
  
  for (uint32_t i = 0; i < 4096; i++)
      translation_table[i].section_fields =
	(short_section_descriptor_t) {
	.SECTION_BASE_ADDRESS_31_20  = i,
	.SECTION_OR_SUPERSECTION_BIT = DESCRIBES_SECTION,
	.ACCESS_PERMISSIONS_2        = AP_2_0_MODEL_RW_PL1 >> 2,
	.ACCESS_PERMISSIONS_1_0      = AP_2_0_MODEL_RW_PL1 & 0b011,
	.DESCRIPTOR_TYPE_1           =
	   SHORT_DESCRIPTOR_SECTION_OR_SUPERSECTION >> 1,
	// rest of fields are 0s
      };

  // meddle with domain settings
  uart_puts("setting domain0 to client access"
	    " and blocking other domains\n\r");

  DACR_t DACR = 0;
  DACR = set_domain_permissions(DACR, 0, DOMAIN_CLIENT_ACCESS);
  for (int i = 1; i < 16; i++)
    DACR = set_domain_permissions(DACR, i, DOMAIN_NO_ACCESS);

  // the above should do the same as this:
  // DACR = 1;
  
  asm("mcr p15, 0, %0, c3, c0, 0" :: "r" (DACR));

  // meddle with SCTLR, which determines how some bits in
  // table descriptors work and also controls caches
  // we don't want to use access flag, so we set AFE to 0
  // we don't want TEX remap, so we set TRE to 0
  // we also disable data and instruction caches and the MMU
  
  // some of this is redundant (i.e. MMU should already be disabled)
  uart_puts("setting C, I, AFE and TRE to 0 in SCTLR\n\r");

  SCTLR_t SCTLR;
  asm("mrc p15, 0, %0, c1, c0, 0" : "=r" (SCTLR.raw));

  SCTLR.fields.M   = 0; // disable MMU
  SCTLR.fields.C   = 0; // disable data cache
  SCTLR.fields.I   = 0; // disable instruction cache
  SCTLR.fields.TRE = 0; // disable TEX remap
  SCTLR.fields.AFE = 0; // disable access flag usage
  asm("mcr p15, 0, %0, c1, c0, 0\n\r"
      "isb" :: "r" (SCTLR.raw) : "memory");

  // TODO: move invalidation instructions to some header as inlines
  
  uart_puts("invalidating instruction cache, branch prediction,"
	    " and entire main TLB\n\r");
  
  // invalidate instruction cache
  asm("mcr p15, 0, r0, c7, c5, 0\n\r" // r0 gets ignored
      "isb" ::: "memory");

  // invalidate branch-prediction
  asm("mcr p15, 0, r0, c7, c5, 6\n\r" // r0 - same as above
      "isb" ::: "memory");

  // invalidate main Translation Lookup Buffer
  asm("mcr p15, 0, %0, c8, c7, 0\n\r"
      "isb" :: "r" (0) : "memory");

  // now set TTBCR to use TTBR0 exclusively
  uart_puts("Setting TTBCR.N to 0, so that"
	    " TTBR0 is used everywhere\n\r");
  
  uint32_t TTBCR = 0;
  asm("mcr p15, 0, %0, c2, c0, 2" :: "r" (TTBCR));
  
  // Now do stuff with TTBR0
  TTBR_t TTBR0;
  TTBR0.raw = 0;
  TTBR0.fields.TTBR_TRANSLATION_TABLE_BASE_ADDRESS =
    translation_table_base >> 14;
  // rest of TTBR0 remains 0s
  
  asm("mcr p15, 0, %0, c2, c0, 0" :: "r" (TTBR0.raw));

  // enable MMU
  uart_puts("enabling the MMU\n\r");

  // redundant - we already have SCTLR contents in the variable
  // asm("mrc p15, 0, %0, c1, c0, 0" : "=r" (SCTLR.raw));

  SCTLR.fields.M = 1;

  asm("mcr p15, 0, %0, c1, c0, 0\n\r"
      "isb" :: "r" (SCTLR.raw) : "memory");

  volatile short_section_descriptor_t *PL0_code_section =
    &translation_table[0b101010101010].section_fields;
  volatile short_section_descriptor_t *UART_memory_section =
    &translation_table[((uint32_t) GPIO_BASE) >> 20].section_fields;

  // map code section to the memory right after out kernel
  uint32_t kernel_memory_end = (uint32_t) (translation_table + 4096);
  uint32_t unprivileged_memory_start =
    (((kernel_memory_end - 1) >> 20) + 1) << 20;
  
  PL0_code_section->SECTION_BASE_ADDRESS_31_20 =
    unprivileged_memory_start >> 20;

  // make the selected section and uart section available for PL0
  PL0_code_section->ACCESS_PERMISSIONS_2   =
    AP_2_0_MODEL_RW_ALL >> 2;
  PL0_code_section->ACCESS_PERMISSIONS_1_0 =
    AP_2_0_MODEL_RW_ALL & 0b011;

  UART_memory_section->ACCESS_PERMISSIONS_2   =
    AP_2_0_MODEL_RW_ALL >> 2;
  UART_memory_section->ACCESS_PERMISSIONS_1_0 =
    AP_2_0_MODEL_RW_ALL & 0b011;
  

  // invalidate main Translation Lookup Buffer (just in case)
  asm("mcr p15, 0, %0, c8, c7, 0\n\r"
      "isb" :: "r" (0) : "memory");

  // check that translation works... by copying a string using one
  // mapping and reading it using other :D
  char message[] = "mapped sections for PL0 code\n\r";

  unsigned int i;
  for (i = 0; i < sizeof(message); i++)
    ((volatile char*) unprivileged_memory_start)[i] =  message[i];
  
  uart_puts((char*) (((uint32_t) 0b101010101010) << 20));

  // now paste a userspace program to that section, jump to it and
  // switch to PL0
  for (size_t i = 0; i < (size_t) &_binary_PL_0_test_img_size; i++)
    ((char*) (((uint32_t) 0b101010101010) << 20))[i] =
      (&_binary_PL_0_test_img_start)[i];

  // jump to that copied code (no switch to PL0 yet, just testing)
  ((void(*)(void)) (((uint32_t) 0b101010101010) << 20))();
  
  while(1);
  
  while (1)
    {
      char c;
      switch(c = uart_getc())
	{
	case '\r':
	  uart_putc('\n');
	default:
	  uart_putc(c);
	}
    }
}