/* Copyright (C) 2008 - 2011 by John Cronin
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:

 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.

 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

/* main entry point from assembler */

/* At this point we are executing in 32-bit mode, there is a gdt set up by the 
 * assembler code and a stack
 * We need to identify the used areas of memory, then identify the kernel 
 * module and load it to a free set of pages.
 * We then need to set up the 64-bit paging structures and map the kernel to 
 * where it should be.
 * Finally, enable paging and jump to 64-bit kernel */

#include "console.h"
#include "mboot.h"
#include "memmgr.h"
#include "elf.h"
#include "cmdline.h"
#include "stdlib.h"
#include "vmem.h"
#include "kif.h"
#include "malloc.h"

#define MB_ID		0x2BADB002
#define DEFAULT_KERNEL		"kernel"

#ifdef EFI
#include <efi.h>
#include <efilib.h>

EFI_SYSTEM_TABLE *gST;
void setup_mm();
int identity_map_tload(EFI_LOADED_IMAGE *li, unsigned long long int *pml4t);
#else
void setup_mm(struct mb_header *mb_header);
int identity_map_tload(struct mb_header *mb_header, unsigned long long int *pml4t);
#endif

const char *profile(const char *msg);
extern void go64(unsigned long int pml4t, unsigned long long int e_point);
extern void *uncompress_bzipped_kernel(void *compressed_data, unsigned int len);

extern char *kernel_name;
extern char bda_addr;
extern char gdt64;

extern unsigned long int heap_end;
extern unsigned long int heap_start;

extern char tload_start;
extern char tload_end;

extern struct Multiboot_Header *kif;
extern Elf64_Shdr *strtab;
extern Elf64_Shdr *symtab;

extern int tysos_debug;

const char def_k_name[] = DEFAULT_KERNEL;

#ifdef EFI
EFI_DEVICE_PATH *make_media_path(const char *s)
{
	/* Convert an ASCII string to a unicode media path structure */
	size_t slen = strlen(s);
	int i;
	uint16_t buf_len = 4 + slen * 2 + 2;

	uint8_t *buf = (uint8_t *)malloc(buf_len);
	buf[0] = 4;		/* type = media path */
	buf[1] = 4;		/* sub-type = file path */
	*(uint16_t *)&buf[2] = buf_len; /* length of entire structure */

	/* Copy ASCII->UCS-2 */
	for(i = 0; i < slen; i++)
	{
		buf[4 + i * 2] = s[i];
		buf[4 + i * 2 + 1] = 0;
	}
	/* Null terminate */
	buf[4 + slen * 2] = 0;
	buf[4 + slen * 2 + 1] = 0;

	return (EFI_DEVICE_PATH *)buf;
}

void efi_pause()
{
	UINTN index;
	EFI_INPUT_KEY key;
	gST->BootServices->WaitForEvent(1, &gST->ConIn->WaitForKey, &index);
	gST->ConIn->ReadKeyStroke(gST->ConIn, &key);
}
#endif


#ifdef EFI
EFI_STATUS efi_main(EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *SystemTable)
#else
int kmain(unsigned long int mbid, struct mb_header *mb_header)
#endif
{
	void *kmod;
	void *pml4t;
	unsigned long long int e_point;
	unsigned long long int kernel_base;
	unsigned long int kmod_length;
	unsigned long long int kernel_vsize;
	struct System_String *cmd_line;

#ifdef EFI
	gST = SystemTable;
#endif

	profile("Start");

	/* say hi */
	//clear();
	puts("TLoad starting...");

#ifdef EFI
	puts("EFI boot detected");
	efi_pause();

	setup_mm();
#ifdef _DEBUG
	puts("EFI: mm setup");
	efi_pause();
#endif

	/* Attempt to load the kernel */

	/* First, get the filesystem tload was loaded from */
	EFI_FILE *fi;
	EFI_FILE_IO_INTERFACE *sfs;
	EFI_LOADED_IMAGE *li;
	EFI_STATUS s;

	s = gST->BootServices->HandleProtocol(ImageHandle, &LoadedImageProtocol, (void **)&li);
	if(EFI_ERROR(s))
	{
		puts("EFI: HandleProtocol(LoadedImageProtocol) failed");
		efi_pause();
		return s;
	}

	s = gST->BootServices->HandleProtocol(li->DeviceHandle, &FileSystemProtocol, (void **)&sfs);
	if(EFI_ERROR(s))
	{
		puts("EFI: HandleProtocol(FileSystemProtocol) failed");
		efi_pause();
		return s;
	}

#ifdef _DEBUG
	puts("EFI: FileSystemProtocol obtained for boot device");
	efi_pause();
#endif

	/* Obtain root directory */
	EFI_FILE *root;
	s = sfs->OpenVolume(sfs, &root);
	if(EFI_ERROR(s))
	{
		puts("EFI: OpenVolume failed for root device");
		efi_pause();
		return s;
	}

#ifdef _DEBUG
	puts("EFI: Root directory obtained");
	efi_pause();
#endif
	
	/* Now try and open boot/tysos.bin */
	EFI_FILE *kernel;
	root->Open(root, &kernel, L"boot\tysos.bin", EFI_FILE_MODE_READ, 0);
	if(EFI_ERROR(s))
	{
		switch(s)
		{
			case EFI_NOT_FOUND:
				puts("EFI: kernel \"boot\tysos.bin\" not found");
				efi_pause();
				return s;
			default:
				puts("EFI: unknown error opening kernel");
				efi_pause();
				return s;
		}
	}

#ifdef _DEBUG
	puts("EFI: kernel opened successfully");
	efi_pause();
#endif

	/* Get the size of the kernel */
	EFI_FILE_INFO *gfi = (EFI_FILE_INFO *)malloc(sizeof(EFI_FILE_INFO));
	UINTN *gfi_bufsize = (UINTN *)malloc(sizeof(UINTN));
	*gfi_bufsize = sizeof(EFI_FILE_INFO);
	puts("EFI: pre GetInfo");
	s = kernel->GetInfo(kernel, &GenericFileInfo, gfi_bufsize, gfi); 
	puts("EFI: post GetInfo");
	if(EFI_ERROR(s))
	{
		puts("EFI: GetInfo() failed");
		efi_pause();
		return s;
	}

#ifdef _DEBUG
	_puts("EFI: kernel found, size: ", 0);
	puthex((unsigned long)gfi->PhysicalSize);
	puts("");
	efi_pause();
#endif


	UINTN kernel_file_size;
	EFI_DEVICE_PATH *kpath = make_media_path("boot/tysos.bin");
	//s = lfi->LoadFile(lfi, kpath, FALSE, &kernel_file_size, NULL);
	if(EFI_ERROR(s))
	{
		puts("EFI: LoadFile(\"boot/tysos.bin\") not found");
		efi_pause();
		return s;
	}

	/* Allocate memory for it */
	kmod = malloc_align((unsigned long)kernel_file_size, 4096, Tysos, 1);
	if(kmod == NULL)
	{
		puts("EFI: couldn't allocate space for the kernel module");
		efi_pause();
		return s;
	}
	//s = lfi->LoadFile(lfi, kpath, FALSE, &kernel_file_size, kmod);
	if(EFI_ERROR(s))
	{
		puts("EFI: Error loading boot/tysos.bin");
		efi_pause();
		return s;
	}
	kmod_length = (unsigned long)kernel_file_size;
#ifdef _DEBUG
	puts("EFI: kernel loaded");
	efi_pause();
#endif
#else
	/* Check that the multiboot id is correct */
	if(mbid != MB_ID) {
		puts("Invalid multiboot magic number");
		return -1;
	}

#ifdef _DEBUG
	_puts("  Command line: ", 0);
	puts(mb_header->cmdline);
	_puts("  Multiboot flags: ", 0);
	puthex(mb_header->flags);
	_puts("", 1);
#endif

	/* Set up the memory manager */
	setup_mm(mb_header);
	profile("setup_mm");

	/* Parse command line */
	parse_cmd_line(mb_header->cmdline);
	profile("parse_cmd_line");

	/* Find the kernel module */
	if(kernel_name == NULL)
		kernel_name = (char *)def_k_name;
#ifdef _DEBUG
	_puts("Searching for kernel: ", 0);
	puts(kernel_name);
#endif
	kmod = NULL;
	if(mb_header->flags & 0x8) {
		unsigned int i;
		struct mb_mod *mods = (struct mb_mod *)mb_header->mods_addr;
		for(i = 0; i < mb_header->mods_count; i++) {
#ifdef _DEBUG
			_puts("  Trying module: ", 0);
			puts(mods[i].string);
#endif
			if(!strcmp(kernel_name, mods[i].string)) {
				kmod = (void *)mods[i].mod_start;
				kmod_length = mods[i].mod_end - mods[i].mod_start;
			}
		}
	}

	if(kmod == NULL) {
		puts("Error: could not find kernel module");
		return -1;
	}

	profile("find kernel module");
#endif

	/* Create a page map level 4 table */
	pml4t = malloc_align(4096, 4096, PagingStructure, 1);
	if(pml4t == NULL) {
		puts("Error: could not allocate page map level 4 table");
		return -1;
	}
	memset(pml4t, 0, 4096);
	/* point it back at itself */
	((unsigned long long int *)pml4t)[511] = ((unsigned long long int)(unsigned long int)pml4t) | 0x3LL;

	profile("create pml4t");

	/* Load the kernel */
	if(*(char *)kmod == 'B')
		kmod = uncompress_bzipped_kernel(kmod, kmod_length);
	if(load_elf64_kernel(kmod, pml4t, &e_point, &kernel_base, &kernel_vsize) < 0) {
		puts("Error: interpreting elf64 kernel");
		return -1;
	}

	profile("load kernel");

#ifdef EFI
#else
	/* Identity map tload */
	if(identity_map_tload(mb_header, pml4t) < 0) {
		puts("Error: could not identity map tload");
		return -1;
	}
#endif

	profile("identity map tload");

#ifndef EFI
	/* Map the bios data area (physical 0x0) to somewhere away from there so we can both still use it and catch
	NULL pointers */
	create_virtual_mapping_from_page((unsigned long long int)(unsigned long int)&bda_addr, pml4t, PF_R | PF_W, 1,
		(void *)0x0, BiosDataArea);

	profile("map bda");
#endif

	/* Identity map the vga */
	create_virtual_mapping_from_page(0xb8000, pml4t, PF_R | PF_W, 1, (void *)0xb8000, VideoHardware);
	profile("map vga");

#ifdef EFI
#else
	/* Generate the kernel information format header */
	memmgr_init_linked_list();
	profile("memmgr_init_linked_list");
	make_kif();
	kif->max_tysos = kernel_base + kernel_vsize;
	kif->gdt = (unsigned long long int)(unsigned long int)&gdt64;
	kif->tysos_paddr = (unsigned long long int)(unsigned long int)kmod;
	kif->tysos_size = (unsigned long long int)(unsigned long int)kmod_length;
	kif->tysos_virtaddr = (unsigned long long int)(unsigned long int)kernel_base;
	kif->tysos_sym_tab_paddr = (unsigned long long int)(unsigned long int)kmod + symtab->sh_offset;
	kif->tysos_sym_tab_size = symtab->sh_size;
	kif->tysos_sym_tab_entsize = symtab->sh_entsize;
	kif->tysos_str_tab_paddr = (unsigned long long int)(unsigned long int)kmod + strtab->sh_offset;
	kif->tysos_str_tab_size = strtab->sh_size;
	kif->debug = (uint8_t)tysos_debug;
	write_kif_modules(kif, mb_header);
	CreateString(&cmd_line, mb_header->cmdline);
	kif->cmdline = (unsigned long long int)(unsigned long int)cmd_line;
	profile("make_kif");
#endif

	/* Go 64 bit and jump to the kernel */
#ifdef _DEBUG
	_puts("Enabling long mode and jumping to ", 0);
	puthex(e_point);
	puts("");
#endif

#ifdef _HALT
	puts("Halting before entering kernel (do not define _HALT to stop this happening)");
	while(1);
#endif

	go64((unsigned long int)pml4t, e_point);

	return 0;
}

#ifdef EFI
#define MEMMAP_SIZE 0x1000
char efi_memmap[MEMMAP_SIZE];

void dump_bitmap();

void setup_mm()
{
	/* Set up the memory manager with details of the memory map provided by EFI */

	/* loop through twice, first setting up the bitmap then the linked list which we'll pass to tysos.  This is because
	  the linked list requires a heap, which requires the bitmap set up to find space for a heap */

	EFI_STATUS s;

	/* Mark the entire bitmap as used */
	mark_used(0, 0x100000000);

	/* First get the EFI memory map */
	UINTN mmapsize = MEMMAP_SIZE;
	UINTN mapkey;
	UINTN descsize;
	UINT32 descver;
	s = gST->BootServices->GetMemoryMap(&mmapsize, (EFI_MEMORY_DESCRIPTOR *)efi_memmap, &mapkey, &descsize, &descver);
	if(s != EFI_SUCCESS)
	{
		puts("EFI: setup_mm: GetMemoryMap() failed");
		return;
	}

	/* Determine the number of entries */
	int e_entries = mmapsize / descsize;
	int i;

	/* Loop first to set up the bitmap */
	for(i = 0; i < e_entries; i++)
	{
		EFI_MEMORY_DESCRIPTOR *desc = (EFI_MEMORY_DESCRIPTOR *)((UINTN)efi_memmap + i * descsize);

		if(desc->Type == EfiConventionalMemory)
			mark_free((unsigned long long)desc->PhysicalStart, (unsigned long long)(desc->NumberOfPages * 4096));
		else
			mark_used((unsigned long long)desc->PhysicalStart, (unsigned long long)(desc->NumberOfPages * 4096));
	}

	/* Mark first page used so we never return it from malloc_align */
	mark_used(0, 4096);

#ifdef _DEBUG
	puts("EFI: mm bitmap setup");
	efi_pause();
	//dump_bitmap();
#endif

	/* Loop again to set up the linked list */
	for(i = 0; i < e_entries; i++)
	{
		EFI_MEMORY_DESCRIPTOR *desc = (EFI_MEMORY_DESCRIPTOR *)((UINTN)efi_memmap + i * descsize);
		int type = Tysos;

		if(desc->Type == EfiConventionalMemory)
			type = Available;

		add_linked_list((unsigned long long)desc->PhysicalStart, (unsigned long long)(desc->NumberOfPages * 4096),
			type);
	}
}
#else
void setup_mm(struct mb_header *mb_header)
{
	/* Set up the memory manager with details of the memory map provided by grub as well as the locations of all the
	modules and elf sections loaded */

	/* loop through twice, first setting up the bitmap then the linked list which we'll pass to tysos.  This is because
	  the linked list requires a heap, which requires the bitmap set up to find space for a heap */

	/* mmap */
	if(mb_header->flags & 0x40) {
		unsigned long int *cur_offset;
		
		cur_offset = (unsigned long int *)mb_header->mmap_addr;

#ifdef _DEBUG
		_puts("mmap_addr: ", 0);
		puthex(mb_header->mmap_addr);
		_puts("  mmap_length: ", 0);
		puthex(mb_header->mmap_length);
		_puts("", 1);
#endif

		while((unsigned long int)cur_offset < (mb_header->mmap_addr + mb_header->mmap_length)) {
			struct mb_mmap *cur_mmap;
			int cur_size = *cur_offset;
			cur_offset++;

			cur_mmap = (struct mb_mmap *)cur_offset;

#ifdef _DEBUG
			_puts("Mmap type: ", 0);
			puthex(cur_mmap->type);
			_puts(" ", 0);
			puthex(cur_mmap->base_addr);
			_puts("-", 0);
			puthex(cur_mmap->base_addr + cur_mmap->length);
			_puts("", 1);
#endif
			if(cur_mmap->type == 1)
				mark_free(cur_mmap->base_addr, cur_mmap->length);
			else
				mark_used(cur_mmap->base_addr, cur_mmap->length);

			cur_offset = (unsigned long int *)(((unsigned long int)cur_offset) + cur_size);
		}
	} else if(mb_header->flags & 0x1) {
		puts("Warning: multiboot header does not specify a memory map");
		mark_free(0x100000, mb_header->mem_upper * 1024);
	} else {
		puts("Warning: no memory information specified in multiboot header");
	}

	/* Always mark the first 1MB as used */
	mark_used(0, 0x100000);

	/* Modules */
	if(mb_header->flags & 0x8) {
		unsigned int i;
		struct mb_mod *mods = (struct mb_mod *)mb_header->mods_addr;
		for(i = 0; i < mb_header->mods_count; i++) {
			mark_used((unsigned long long int)mods[i].mod_start, (unsigned long long int)(mods[i].mod_end -
				mods[i].mod_start));
#ifdef _DEBUG
			_puts("mod: ", 0);
			_puts(mods[i].string, 0);
			_puts("  s: ", 0);
			puthex(mods[i].mod_start);
			_puts("  e: ", 0);
			puthex(mods[i].mod_end);
			_puts("", 1);
#endif
		}
	}

	/* Elf sections */
	if(mb_header->flags & 0x20) {
		unsigned int i;

#ifdef _DEBUG
		_puts("elf section count: ", 0);
		_puthex(mb_header->elf_num);
		_puts("", 1);
#endif

		for(i = 0; i < mb_header->elf_num; i++) {
			Elf32_Shdr *shdr = (Elf32_Shdr *)(mb_header->elf_addr + i * mb_header->elf_size);
			mark_used((unsigned long long int)shdr->sh_addr, (unsigned long long int)shdr->sh_size);
#ifdef _DEBUG
			_puts("elf section: s: ", 0);
			_puthex(shdr->sh_addr);
			_puts(" e: ", 0);
			_puthex(shdr->sh_addr + shdr->sh_size);
			_puts("", 1);
#endif
		}
	}

	/* tload itself */
	{
		unsigned long long int tl_start = (unsigned long long int)(unsigned long int)&tload_start;
		unsigned long long int tl_end = (unsigned long long int)(unsigned long int)&tload_end;
		mark_used(tl_start, tl_end - tl_start);

#ifdef _DEBUG
		_puts("tload: s: ", 0);
		_puthex((unsigned long int)tl_start);
		_puts(" e: ", 0);
		_puthex((unsigned long int)tl_end);
		_puts("", 1);
#endif
	}

	/* loop again to set up the linked list */

	/* mmap */
	if(mb_header->flags & 0x40) {
		unsigned long int *cur_offset;
		
		cur_offset = (unsigned long int *)mb_header->mmap_addr;

		while((unsigned long int)cur_offset < (mb_header->mmap_addr + mb_header->mmap_length)) {
			struct mb_mmap *cur_mmap;
			int cur_size = *cur_offset;
			cur_offset++;

			cur_mmap = (struct mb_mmap *)cur_offset;

			add_linked_list((unsigned long long int)cur_mmap->base_addr, (unsigned long long int)cur_mmap->length,
				(int)cur_mmap->type);

			cur_offset = (unsigned long int *)(((unsigned long int)cur_offset) + cur_size);
		}
	}

	/* Modules */
	if(mb_header->flags & 0x8) {
		unsigned int i;
		struct mb_mod *mods = (struct mb_mod *)mb_header->mods_addr;
		for(i = 0; i < mb_header->mods_count; i++) {
			add_linked_list((unsigned long long int)mods[i].mod_start, (unsigned long long int)(mods[i].mod_end -
				mods[i].mod_start), Module);
		}
	}

	/* Elf sections */
	if(mb_header->flags & 0x20) {
		unsigned int i;

		for(i = 0; i < mb_header->elf_num; i++) {
			Elf32_Shdr *shdr = (Elf32_Shdr *)(mb_header->elf_addr + i * mb_header->elf_size);
			add_linked_list((unsigned long long int)shdr->sh_addr, (unsigned long long int)shdr->sh_size,
				TLoad);
		}
	}

	/* tload itself */
	{
		unsigned long long int tl_start = (unsigned long long int)(unsigned long int)&tload_start;
		unsigned long long int tl_end = (unsigned long long int)(unsigned long int)&tload_end;
		add_linked_list(tl_start, tl_end - tl_start, TLoad);
	}
}
#endif

#ifdef EFI
int identity_map_tload(EFI_LOADED_IMAGE *li, unsigned long long int *pml4t)
{
	unsigned long int tl_start = (unsigned long int)li->ImageBase;
	unsigned long int tl_end = tl_start + (unsigned long int)li->ImageSize;
	unsigned long int cur_page = tl_start & (~(PAGESIZE - 1));

	while(cur_page < tl_end)
	{
		create_virtual_mapping_from_page(cur_page, pml4t, PF_R | PF_W | PF_X, 1, (void *)cur_page, TLoad);
		cur_page += PAGESIZE;
	}

	return 0;
}
#else
int identity_map_tload(struct mb_header *mb_header, unsigned long long int *pml4t)
{
	/* Identity map the bootloader so that the bootloader is still mapped once we jump to long mode */
	unsigned long int i;

	for(i = 0; i < mb_header->elf_num; i++) {
		Elf32_Shdr *shdr;
		unsigned long int base;

		shdr = (Elf32_Shdr *)(i * mb_header->elf_size + mb_header->elf_addr);
		base = shdr->sh_addr & (~(PAGESIZE - 1));
		while(base < (shdr->sh_addr + shdr->sh_size)) {
			create_virtual_mapping_from_page(base, pml4t, PF_R | PF_W | PF_X, 1, (void *)base, TLoad);
			base += PAGESIZE;
		}
	}

	/* Identity map the tload heap */
	{
		unsigned long int base;
		
		for(base = heap_start; base < heap_end; base += PAGESIZE)
			create_virtual_mapping_from_page(base, pml4t, PF_R | PF_W | PF_X, 1, (void *)base, TLoad);
	}

	/* Identity map using an alternate algorithm */
	{
		unsigned long int tl_start = (unsigned long int)&tload_start;
		unsigned long int tl_end = (unsigned long int)&tload_end;
		unsigned long int cur_page = tl_start & (~(PAGESIZE - 1));

		while(cur_page < tl_end)
		{
			create_virtual_mapping_from_page(cur_page, pml4t, PF_R | PF_W | PF_X, 1, (void *)cur_page, TLoad);
			cur_page += PAGESIZE;
		}
	}

	return 0;
}
#endif

#ifdef _PROFILE
const char *profile(const char *msg)
{
	static unsigned long int prev_high = 0;
	static unsigned long int prev_low = 0;

	unsigned long int high = 0;
	unsigned long int low = 0;

	unsigned long int delta_high = 0;
	unsigned long int delta_low = 0;

	__asm__ __volatile__ ("rdtsc" : "=d"(high), "=a"(low));

	delta_high = high - prev_high;
	delta_low = low - prev_low;

	prev_high = high;
	prev_low = low;

	_puts(msg, 1);
	_puts("Abs: 0x", 0);
	_puthex(high);
	_puthex(low);
	_puts("  Delta: 0x", 0);
	_puthex(delta_high);
	_puthex(delta_low);
	_puts("", 1);

	return msg;
}
#else
const char * profile(const char *msg)
{
	return msg;		/* prevent gcc complaining about unused parameters in a portable way */
}
#endif
