#include <shim.h>
#include <lib.h>
#include <stdint.h>
#include <panic.h>

#include "mboot.h"

#define MBOOT_HEADER_MAGIC  0x36D76289

#define MBOOT_CMDLINE       1
#define MBOOT_MEMORY_MAP    6
#define MBOOT_FRAMEBUFFER   8
#define MBOOT_ELF_SYMBOLS   9
#define MBOOT_OLD_RSDP     14
#define MBOOT_NEW_RSDP     15

struct mboot_info {
	uint32_t total_size;
	uint32_t reserved;
	char tags[];
};

struct mboot_tag {
	uint32_t type;
	uint32_t size;
	char data[];
};

struct mboot_tag_elf_sections {
	uint32_t type;
	uint32_t size;
	uint16_t num;
	uint16_t entsize;
	uint16_t shndx;
	uint16_t reserved;
	char sections[];
};

struct mboot_tag_elf_sections_entry {
	uint32_t sh_name;
	uint32_t sh_type;
	uint64_t sh_flags;
	uint64_t sh_addr;
	uint64_t sh_offset;
	uint64_t sh_size;
	uint32_t sh_link;
	uint32_t sh_info;
	uint64_t sh_addralign;
	uint64_t sh_entsize;
};

struct mboot_mmap_entry {
	uint64_t addr;
	uint64_t len;
	uint32_t type;
	uint32_t zero;
};

struct mboot_tag_mmap {
	uint32_t type;
	uint32_t size;
	uint32_t entry_size;
	uint32_t entry_version;
	struct mboot_mmap_entry entries[];
};


struct mboot_tag_old_rsdp {
	uint32_t type;
	uint32_t size;
	char rsdp[];
};

struct mboot_tag_new_rsdp {
	uint32_t type;
	uint32_t size;
	char rsdp[];
};

struct mboot_tag_cmdline {
	uint32_t type;
	uint32_t size;
	uint8_t cmdline[];
};

struct mboot_tag_framebuffer {
	uint32_t type;
	uint32_t size;
	uint64_t framebuffer_addr;
	uint32_t framebuffer_pitch;
	uint32_t framebuffer_width;
	uint32_t framebuffer_height;
	uint8_t framebuffer_bpp;
	uint8_t framebuffer_type;
	uint16_t reserved;
};

static void read_symbols(
		volatile struct boot_info *shim_info,
		volatile struct mboot_tag_elf_sections *sections
) {
	shim_info->symbol_table = (void * volatile) sections->sections;

//	struct mboot_elf_section_header *section =
//		(struct mboot_elf_section_header *) (layout->elf_section_headers);
//
//	for (uint32_t i = 0; i < layout->num; i++) {
//		char buf[20];
//
//		ultoa(i, buf, 10);
//		serial_out_str("[");
//		serial_out_str(buf);
//		serial_out_str("]\t");
//
//		serial_out_str((char *)(kaddr(symtab) + section->sh_name));
//		serial_out('\t');
//
//		ultoa(section->sh_type, buf, 16);
//		serial_out_str("type: 0x");
//		serial_out_str(buf);
//		serial_out('\t');
//
//		ultoa(section->sh_addr, buf, 16);
//		serial_out_str("addr: 0x");
//		serial_out_str(buf);
//		serial_out('\t');
//
//		ultoa(section->sh_offset, buf, 16);
//		serial_out_str("offset: 0x");
//		serial_out_str(buf);
//		serial_out('\n');
//
//		section++;
//	}
}

static void read_cmdline(
		volatile struct boot_info *shim_info,
		volatile struct mboot_tag_cmdline *cmdline
) {
	uint32_t size = cmdline->size - 8;
	if (size >= CMDLINE_MAX)
		size = CMDLINE_MAX; // truncate :(
	memcpyv(shim_info->cmdline, cmdline->cmdline, size);
	shim_info->cmdline[size] = '\0';
}

static void read_framebuffer(
		volatile struct boot_info *shim_info,
		volatile struct mboot_tag_framebuffer *framebuffer
) {
	shim_info->fb.addr = framebuffer->framebuffer_addr;
	shim_info->fb.width = framebuffer->framebuffer_width;
	shim_info->fb.height = framebuffer->framebuffer_height;
	shim_info->fb.pitch = framebuffer->framebuffer_pitch;
	shim_info->fb.bit_depth = framebuffer->framebuffer_bpp;
}

static const char *segment_type[] = {
	"Reserved",
	"Free",
	"Reserved",
	"ACPI Reserved",
	"Hibernation",
	"Defective",
	"Unknown"
};

static void read_memory_map(
		volatile struct boot_info *shim_info,
		volatile struct mboot_tag_mmap *map
) {
	int idx = 0;
	uintptr_t i = (uintptr_t)map->entries;
	kprintf("MEMORY MAP\n");
	char buf[20];
	for (;
			i < (uintptr_t)map->entries + map->size;
			i += map->entry_size, idx++
	) {
		struct mboot_mmap_entry *seg = (struct mboot_mmap_entry *) i;
		const char *type = NULL;
		if (seg->type > 4)
			type = segment_type[6];
		else
			type = segment_type[seg->type];
		kprintf("ADDR: %16p  LEN: %4s  TYPE: %s (%d)\n",
				(void *)seg->addr,
				btoa(seg->len, buf),
				type,
				seg->type
		);
		if (seg->type != 1 || seg->len < 1)
			continue;
		shim_info->map.entries[idx].addr = seg->addr;
		shim_info->map.entries[idx].len = seg->len;
	}
	shim_info->map.entry_count = idx;
}

static void read_old_rsdp(
		volatile struct boot_info *shim_info,
		volatile struct mboot_tag_old_rsdp *rsdp
) {
	if (shim_info->acpi_table != NULL)
		return; // xsdp is newer and has been loaded
	shim_info->acpi_table = (void *volatile) rsdp->rsdp;
}

static void read_new_rsdp(
		volatile struct boot_info *shim_info,
		volatile struct mboot_tag_new_rsdp *rsdp
) {
	shim_info->acpi_table = (void *volatile) rsdp->rsdp;
}

void mboot_load_info(
		long mboot_magic,
		volatile const void *mboot_data_ptr,
		volatile struct boot_info *shim_info
) {

	if (mboot_magic != MBOOT_HEADER_MAGIC)
		panic("invalid multiboot magic");

	memsetv(shim_info, 0, sizeof(struct boot_info));

	struct mboot_info *mboot_info = (struct mboot_info *) mboot_data_ptr;
	const char *mboot_end = ((char *) mboot_info) + mboot_info->total_size;

	char *tag_ptr = mboot_info->tags;

	while (tag_ptr < mboot_end) {
	struct mboot_tag *tag = (struct mboot_tag *) tag_ptr;

		switch (tag->type) {
			case MBOOT_CMDLINE:
				read_cmdline(
						shim_info,
						(struct mboot_tag_cmdline *) tag
				);
				break;
			case MBOOT_FRAMEBUFFER:
				read_framebuffer(
						shim_info,
						(struct mboot_tag_framebuffer *) tag
				);
				break;
			case MBOOT_MEMORY_MAP:
				read_memory_map(
						shim_info,
						(struct mboot_tag_mmap *) tag
				);
				break;
			case MBOOT_ELF_SYMBOLS:
				read_symbols(
						shim_info,
						(struct mboot_tag_elf_sections *) tag
				);
				break;
			case MBOOT_OLD_RSDP:
				read_old_rsdp(
						shim_info,
						(struct mboot_tag_old_rsdp *) tag
				);
				break;
			case MBOOT_NEW_RSDP:
				read_new_rsdp(
						shim_info,
						(struct mboot_tag_new_rsdp *) tag
				);
				break;
			default:
				break;
		}

		int size = tag->size;
		if (size % 8 != 0) {
			size += 8 - (size % 8);
		}

		tag_ptr += size;
	}
}