#include <efi.h>
#include <efilib.h>
#include <efidef.h>
#include <string.h>

#include "disk.h"
#include "vga.h"
#include "vmmem.h"
#include "ELF.h"
#include "memorytypes.h"
#include "paging_struct.h"
#include "paging.h"
#include "other.h"

// pixel struct information(hardcoded to what qemu exposes)
typedef struct _Pixel
{
  UINT8 B;
  UINT8 G;
  UINT8 R;
  UINT8 Z;
} Pixel;

// struct that is going to be passed to the kernel about general system information
typedef struct _OSDATA
{
  UINT32 Magic; // magic number to check

  UINT32 FBWidth; // with of the framebuffer
  UINT32 FBHeight; // height
  UINT32 PixelSize; // size of each pixel(an rgb pixel might have a bigger size)
  void * FBAddr; // address of the linear framebuffer

  void * MEMMap; // pointer to the system memory map

  void * RAMDisk; // pointer to a ramdisk loaded from the hdd

  void* RSDP;
} OSDATA;

extern void BootDisableInterrupts(void); // asm code is not correct(callee doesn't set the stack correctly)
typedef void (*kfn)(OSDATA *);  // typedef to setup the entry point of the kernel and do a "jump" into it



void
EFIAPI
efi_main (EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *SystemTable)
{
  __asm__("hlt");

  // Memory Map
  UINTN mapsize               = 0;
  UINTN allocsize             = 0;
  EFI_MEMORY_DESCRIPTOR * map = NULL;
  UINTN mapkey                = 0;
  UINTN descriptorsize        = 0;
  UINT32 version              = 0;

  
  InitializeLib(ImageHandle, SystemTable);
  uefi_call_wrapper(ST->ConOut->ClearScreen, 1, ST->ConOut); // clear the screen
  //Pixel * fb = SetVideoMode(1024, 768, 32);
  Pixel * fb = SetVideoMode(800, 600, 24);

  PrintImageAddr(ImageHandle);

  // Print(L"Number of tables: %d\n", SystemTable->NumberOfTableEntries);

  EFI_GUID acpi_10 = ACPI_TABLE_GUID;
  EFI_GUID acpi_20 = ACPI_20_TABLE_GUID;
  void* acpi20table = NULL;


  for(int i = 0; i < SystemTable->NumberOfTableEntries; i++)
  {
    if(AreEqual(&(SystemTable->ConfigurationTable[i].VendorGuid), &acpi_10, sizeof(EFI_GUID)))
    {
      // Print(L"ACPI 1.0 Found\n");
    }
    if(AreEqual(&(SystemTable->ConfigurationTable[i].VendorGuid), &acpi_20, sizeof(EFI_GUID)))
    {
      // Print(L"ACPI 2.0 Found\n");
      acpi20table = SystemTable->ConfigurationTable[i].VendorTable; // <- this should be the RSDP
    }
  }

  // some uefi implementations time out with their default setting, disable the timer
  BS->SetWatchdogTimer(0, 0, 0, NULL);
   

  //Print(L"Firmware Vendor: %s Rev: 0x%08x\n", ST->FirmwareVendor, ST->FirmwareRevision);

  //while(1){};

  // allocate the datat for the kernel(need to specify memory time not to be a generic loader data type)
  OSDATA * osdata = AllocatePool(sizeof(OSDATA));

  if(osdata == NULL)
  {
    Print(L"Os Data allocation failed\n");
  }

  int32_t kernel_size;

  ELF * kernel = LoadFile(L"kernel.bin", MEM_KERNEL, &kernel_size); // we set the memory type to the one from the kernel

  {
    PH* ph = (PH*)(((uint64_t)kernel) + kernel->e_phoff); // use this to know how much and what pages to map
    uint64_t phcount = (((uint64_t)kernel) + kernel->e_phnum);

    //kernel_size = ph->p_filesz;
    uint64_t size = (uint64_t)(kernel_size = ph->p_memsz);
    uint64_t entry = (uint64_t)(kernel + ph->p_offset);

    /*
    Print(L"Kernel Size:   0x%llX\n", size);
    Print(L"Kernel Entry:  0x%llX\n", entry);
    Print(L"Program Header Count:   0x%llX\n", phcount);
    */
  }

  // print general information about the kernel elf header
  //PrintELFInfo(kernel);
  Loaded_ELF* loaded_kernel = LoadELF(kernel);


  // attempt to allocate the memory map, first try is going to be too small
  // as such the firmware will return the correct size
  EFI_STATUS memret     = EFI_SUCCESS;
  EFI_STATUS bootstatus = EFI_SUCCESS;

  uefi_call_wrapper(BS->GetMemoryMap, 5, &mapsize, map, &mapkey, &descriptorsize, &version);


  // as the allocation will probably modify the memory map allocate 4kb more(one page)
  // so that the new memory map probably fits
  allocsize = mapsize + 10*4098;

  mapsize = allocsize;

  uefi_call_wrapper(BS->AllocatePool, 3, EfiLoaderData, allocsize, (void**)&map);

  uefi_call_wrapper(BS->GetMemoryMap, 5, &mapsize, map, &mapkey, &descriptorsize, &version);

  // set virtual addresses in here
  // try some paging
  initCR3();

  //uint64_t address = 0;
  //uint64_t max = 0x400000000ull;//0x20000000;
  // EFI_MEMORY_DESCRIPTOR * mapiterator = map;

  uint64_t elements = mapsize/descriptorsize;

  printCR3();

  //Print(L"elements: %d\n", elements);

  

  for(int entry = 0; entry < elements; ++entry)
  {
    // mapiterator = (EFI_MEMORY_DESCRIPTOR*)(((EFI_PHYSICAL_ADDRESS)mapiterator + descriptorsize));

    EFI_MEMORY_DESCRIPTOR* mapiterator = (EFI_MEMORY_DESCRIPTOR*)(((uint8_t*)map) + entry*descriptorsize);

    // EFI_MEMORY_DESCRIPTOR* mapiterator = &(map[entry]);
    uint64_t page = mapiterator->PhysicalStart;

    uint64_t Pstart = mapiterator->PhysicalStart;
    uint64_t Vstart = mapiterator->VirtualStart;
    uint64_t Npages = mapiterator->NumberOfPages;
    uint64_t Tpage = mapiterator->Type;
    uint64_t Att = mapiterator->Attribute;

    //if(Pstart != Vstart)
    /*
    if(Vstart != 0)
    {
      Print(L"\n");
      Print(L"---------------------------------------\n");
      Print(L"NON MATCHING PHYSICAL AND VIRTUAL PAGES\n");
      Print(L"Physical Start:   0x%llX\n", Pstart);
      Print(L"Virtual Start:    0x%llX\n", Vstart);
      Print(L"Number of Pages:  0x%llX\n", Npages);
      Print(L"Type of Page:     0x%llX\n", Tpage);
      Print(L"---------------------------------------\n");
      Print(L"\n");
    }
    */

    // if((Npages > 0) && (Vstart != 0))
    // if(Vstart != 0)
    // if((Pstart == 0) && (Vstart == 0))
    // if(Pstart == Vstart)
    if(Tpage == EfiConventionalMemory) // free pages
    {
      Print(L"\n");
      Print(L"---------------------------------------\n");
      Print(L"MAPPED MEMORY\n");
      Print(L"Physical Start:    0x%llX\n", Pstart);
      Print(L"Virtual Start:     0x%llX\n", Vstart);
      Print(L"Number of Pages:   0x%llX\n", Npages);
      Print(L"Type of Page:      0x%llX\n", Tpage);
      Print(L"Attribute of Page: 0x%llX\n", Att);
      Print(L"---------------------------------------\n");
      Print(L"\n");
    }

    //Print(L"page: 0x%llX\n", page);
    /* for(int pageentry = 0; pageentry < mapiterator->NumberOfPages; pageentry++)
    {
      //Print(L"SetAddr\n");
      SetVirtualAddress(page, page);
      page += 0x1000;
    } */
  }

  while(1) {};

  uint64_t maxaddr = 0x0000001000000000ull;
  for(uint64_t page = 0x0; page < maxaddr; page += 0x1000)
  {
    SetVirtualAddress(page, page); // <- This is chanigng the memorymap 100%
  }

  // need to map the physical address the kernel is in to -2GB virtual
  //uint64_t startvm = 0xffffffff7fffffffull;
  //uint64_t startvm = kernel->e_entry;
  //uint64_t currentvm = 0;

  //PH* ph = (PH*)(((uint64_t)kernel) + kernel->e_phoff); // use this to know how much and what pages to map

  //kernel_size = ph->p_filesz;
  //kernel_size = ph->p_memsz;

  /*
  while(currentvm < kernel_size)
  {
    SetVirtualAddress((uint64_t)kernel + ph->p_offset + currentvm, startvm + currentvm);

    currentvm += 0x1000;
  }
  */

  Print(L"BEFORE MAPPING KERNEL ENTRY: 0x%llX\n", loaded_kernel->entry);

  for(int i = 0; i < loaded_kernel->ph_num; i++)
  {
    Loaded_PH* ph = &(loaded_kernel->ph[i]);
    Print(L"MAPPING PH: 0x%llX\n", (uint64_t)ph->vaddr);
    uint64_t mapsize = ph->mem_size;
    uint64_t mapped = 0;
    while(mapped < mapsize)
    {
      SetVirtualAddress((uint64_t)ph->data + mapped, (uint64_t)ph->vaddr + mapped);
      mapped += 0x1000;
    }
  }

  // now map already allocated 8KB for the Kernel Stack


  //printCR3();

  //Print(L"kernel: 0x%llx\n", kernel);

  // the call to exit boot services tells the firmware we are ready to take control of the system

  // never ever call Print after the next line
  // after the first call boot services can be partially disabled
  // loop until the firmware reports a successful exit, the specification allows for partial shutdowns
  // so more than one call might be necessary(on qemu with OVMF it is)
  while((bootstatus = uefi_call_wrapper(BS->ExitBootServices, 2, ImageHandle, mapkey)) != EFI_SUCCESS) // <- actually we need an up to date mapkey, so maybe we could actually get a memorymap before, create a OS liked memory map and exit services?
  {

    mapsize = allocsize;

    while((memret = uefi_call_wrapper(BS->GetMemoryMap, 5, &mapsize, map, &mapkey, &descriptorsize, &version)) == EFI_BUFFER_TOO_SMALL)
    {

      allocsize = allocsize + 2*4096; // add 8kb(2 pages)

      uefi_call_wrapper(BS->FreePool, 1, map);

      uefi_call_wrapper(BS->AllocatePool, 3, EfiLoaderData, allocsize, (void**)&map);

      mapsize = allocsize;
    }
  }

  SetCrc(&(SystemTable->Hdr)); // As we exited boot services we need to set the CRC32 again

  // as we can't print to the screen the way of showing the return status of this function
  // is to write a red or green square on the top left corner of the screen
  if(SetVM(mapsize, descriptorsize, version, map, kernel) != EFI_SUCCESS)
  {
    int i = 0;
    int j = 0;

    for(i = 0; i < 256; ++i)
    {
      for(j = 0; j < 256; ++j)
      {
        Pixel p;
        p.R = 255;
        p.G = 0;
        p.B = 0;
        //p.Z = 255;
        fb[j + 800*i] = p;
      }
    }
  } else
  {
    int i = 0;
    int j = 0;

    for(i = 0; i < 64; ++i)
    {
      for(j = 0; j < 64; ++j)
      {
        Pixel p;
        p.R = 0;
        p.G = 255;
        p.B = 0;
        //p.Z = 255;
        fb[j + 800*i] = p;
      }
    }
  }

  // we should now set s  Print(L"Calling ExitBootServices\n");oe virtual mapping
  //BootDisableInterrupts();


  //OSDATA * osdata = (OSDATA*)(640*1024);



  osdata->Magic = 0xDDEE;
  osdata->FBWidth = 1024;
  osdata->FBHeight = 768;
  osdata->FBAddr = fb;
  osdata->PixelSize = 24;
  osdata->MEMMap = map;
  osdata->RAMDisk = NULL;
  osdata->RSDP = acpi20table;

  // this while only serves not to call the kernel for now
  // after calling exit boot services we can't return to the uefi environment because it's been destroyed
  //

  //kernel = (ELF*)0x280000000; // correct virtual pointer to 10GB (Sign: 0  PML4: 0  PDP:10  PD:0  Page:0)

  // kfn kernel_jump = (void*)kernel->e_entry;//(void*)((EFI_PHYSICAL_ADDRESS)kernel + kernel->EntryPoint);
  kfn kernel_jump = (void*)loaded_kernel->entry;//(void*)((EFI_PHYSICAL_ADDRESS)kernel + kernel->EntryPoint);


  writeCR3();
  //while(1){}
  {
    int i = 0;
    int j = 0;

    for(i = 0; i < 256; ++i)
    {
      for(j = 0; j < 256; ++j)
      {
        Pixel p;
        p.R = 255;
        p.G = 255;
        p.B = 255;
        p.Z = 255;
        fb[j + 800*i] = p;
      }
    }
  }

  // disable interrupts
  __asm__("cli");

  //writeCR3();


  //__asm__("hlt");

  kernel_jump(osdata);

  while(1){}

  __asm__("hlt"); // sanity in case the kernel exists, should throw an error somehow
  // not as if the kernel shouldn't have this same code at the end of the main though

  // we should never ever reach this point(if kernel exists, it should shutdown the computer)
}