#include <stdatomic.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>

#include <kernel/_kernel.h>
#include <kernel/x86/idt.h>
#include <kernel/x86/pic.h>
#include <kernel/x86/pit.h>
#include <kernel/x86/io.h>
#include <kernel/x86/keyb.h>

__attribute__((aligned(0x10)))
static idt_entry_t idt[256];

static idtr_t idtr;

static bool vectors[IDT_MAX_DESCRIPTORS];

extern struct pit_state pit;

#define EXTERNAL_BIT (0x1)

#define TBL_GDT (0x0)
#define TBL_LDT (0x2)
#define TBL_IDT (0x1)
#define TBL_IDT_TWO (0x3)

static void examine_selector(uint32_t selector)
{
	printf("The interrupt arised %s.\n", (selector & EXTERNAL_BIT) == EXTERNAL_BIT ? "externally" : "internally");


	printf("The descriptor 0x%x, is located in the ", selector, (selector & 0x0000FFFF) >> 3);
	if ((((selector >> 1) & TBL_IDT) == TBL_IDT)
			|| ((selector >> 1) & TBL_IDT_TWO) == TBL_IDT_TWO) {
		puts("IDT.");

	} else if (((selector >> 1) & TBL_LDT) == TBL_LDT) {
		puts("LDT.");
	} else 	{
		puts("GDT.");
	}
	
	return;	
}

#undef TBL_GDT
#undef TBL_LDT
#undef TBL_IDT
#undef TBL_IDT_TWO

#define PF_P (1 << 0)
#define PF_W (1 << 1)
#define PF_U (1 << 2)
#define PF_R (1 << 3)
#define PF_I (1 << 4)
#define PF_PK (1 << 5)
#define PF_SS (1 << 6)
#define PF_SGX (1 << 15)
static void examine_page_fault(uint32_t error_code)
{
	puts("Page fault:");
	if (error_code & PF_P) {
		puts("Page-protection violation.");
	} else {
		puts("Not present page.");
	}

	if (error_code & PF_W) {
		puts("Write access.");
	} else {
		puts("Read access.");
	}

	if (error_code & PF_U) {
		puts("In CPL 3.");
	} 

	if (error_code & PF_R) { // Only applies if PSE or PAE are set in CR4
		puts("One or more page directory entires contain reserved bits which are set to 1.");
	} 

	if (error_code & PF_I) { // Only applicable when the No-Execute bit is supported and enabled
		puts("Instruction fetch.");
	} 

	if (error_code & PF_PK) {
		puts("Protection-key violation"); //TODO PRKU (user mode access) or PRKS MSR (supervisor mode access) specifies the key rights
	} 

	if (error_code & PF_SS) {
		puts("Shadow stack access.");
	} 

	if (error_code & PF_SGX) { // Unrelated to ordinary paging
		puts("SGX vioilation");
	}

	// TODO, also sets CR2 reg to the virtual address which caused the page fault
}
#undef PF_P
#undef PF_W
#undef PF_U
#undef PF_R
#undef PF_I
#undef PF_PK
#undef PF_SS
#undef PF_SGX

static void dump_cpu_state(struct cpu_state cpu, struct stack_state stack)
{
	kinfo("DUMPING CPU STATE");
	printf("eax: %2\nebx: %2\necx: %2\nedx: %2\nesi: %2\nedi: %22\nebp: %2\n", 
			cpu.eax, cpu.ebx, cpu.ecx, cpu.edx, cpu.esi, cpu.edi, cpu.ebp);
	kinfo("DUMPING STACK STATE FOR INTERRUPT");
	printf("eip: %2\ncs: %2\neflags: %2\n",
			stack.eip, stack.cs, stack.eflags);
}

#define EXCEPTION_LOCATION() printf("Exception occurred at 0x%2\n", stack.eip)

void exception_handler(struct cpu_state __attribute__((unused)) cpu, uint32_t interrupt, struct stack_state stack)
{
	uint32_t inbyte;
	switch (interrupt) {
		/** EXCEPTIONS BEGIN **/
		case EXCEPT_DIV_ERR:
			kerror("EXCEPTION: DIV BY ZERO");
			EXCEPTION_LOCATION();
			break;
		case EXCEPT_DEBUG:
			kinfo("EXCEPTION: DEBUG"); // TODO: this one has more specific requirements
			break;
		case EXCEPT_NMI:
			kinfo("EXCEPTION: NON-MASKABLE INTERRUPT"); // TODO
			break;
		case EXCEPT_BREAKPOINT:
			kinfo("EXCEPTION: BREAKPOINT");
			// TODO
			break;
		case EXCEPT_OVERFLOW:
			kerror("EXCEPTION: OVERFLOW");
			break;
		case EXCEPT_BOUND_RANGE_EXCEEDED:
			kerror("EXCEPTION: BOUND RANGE EXCEEDED");
			EXCEPTION_LOCATION();
			break;
		case EXCEPT_INVALID_OPCODE:
			kerror("EXCEPTION: INVALID OPCODE");
			EXCEPTION_LOCATION();
			break;
		case EXCEPT_DEVICE_NOT_AVAILABLE:
			kerror("EXCEPTION: DEVICE NOT AVAILABLE");
			printf("Instructions expecting an FPU were used, but no FPU was found. Is it disabled in CR0? Do you have an FPU?\n");
			EXCEPTION_LOCATION();
			break;
		case EXCEPT_DOUBLE_FAULT:
			kerror("EXCEPTION: DOUBLE FAULT");
			__asm__ volatile ("cli; hlt"); // boned
			break;
		case EXCEPT_INVALID_TSS:
			kerror("EXCEPTION: INVALID TSS");
			examine_selector(stack.error_code);
			// TODO: when i implement the TSS, the EIP will be one of two things
			// if this occurs before loading segmenet selectors from the TSS, EIP will be the instruction that caused the instruction
			// otherwise its the first instruction in the new task
			break;
		case EXCEPT_SEG_NOT_PRESENT:
			kerror("EXCEPTION: SEGMENT NOT PRESENT");
			EXCEPTION_LOCATION();
			examine_selector(stack.error_code);
			break;
		case EXCEPT_STACK_SEG_FAULT:
			kerror("EXCEPTION: STACK-SEGMENT FAULT");
			if (stack.error_code == 0) {
				EXCEPTION_LOCATION();
			} else {
				printf("Tried loading non-present stack segment during hardware task switch.\n");
				examine_selector(stack.error_code);
			}
			break;
		case EXCEPT_GENERAL_PROTECTION:
			kerror("EXCEPTION: GENERAL PROTECTION FAULT");
			EXCEPTION_LOCATION();
			if (stack.error_code != 0)
				examine_selector(stack.error_code);
			break;
		case EXCEPT_PAGE_FAULT:
			kerror("EXCEPTION: PAGE FAULT");
			EXCEPTION_LOCATION();
			examine_page_fault(stack.error_code);
			break;
		case EXCEPT_FLOATING_POINT_ERR_FPU:
			kerror("EXCEPTION: FLOATING POINT ERROR FPU");
			// TODO: this requires a custom handler, it needs another register
			break;
		case EXCEPT_ALIGNMENT_CHECK:
			kerror("EXCEPTION: ALIGNMENT CHECK");

			// TODO: only in CPL3
			break;
		case EXCEPT_MACHINE_CHECK:
			kerror("EXCEPTION: MACHINE CHECK");
			// TODO disabled by default
			break;
		case EXCEPT_FLOATING_POINT_ERR_SIMD:
			kerror("EXCEPTION: FLOATING POINT ERROR SIMD");
			// TODO disabled by default
			break;
		case EXCEPT_VIRT:
			kerror("EXCEPTION: VIRTUALIZATION");
			break;
		case EXCEPT_CTRL_PROT:
			kerror("EXCEPTION: CONTROL PROTECTION");
			printf("Error code: 0x%2\n", stack.error_code);
			break;
		case EXCEPT_HYPERVISOR_INJECTION:
			kerror("EXCEPTION: HYPERVISOR INJECTION");
			break;
		case EXCEPT_VMM_COMMUNICATION:
			kerror("EXCEPTION: VMM COMMUNICATION");
			printf("Error code: 0x%2\n", stack.error_code);
			break;
		case EXCEPT_SECURITY_EXCEPTION:
			kerror("EXCEPTION: SECURITY");
			printf("Error code: 0x%2\n", stack.error_code);
			break;
		/** EXCEPTIONS END **/

		/** PIC BEGIN **/
		case PIC_PIT:
			pit.interrupts++;
			PIC_sendEOI(0);
			break;
		case PIC_KEYB:
			inbyte = inb(0x60);
			do_keypress(decode_scancode(inbyte));
			PIC_sendEOI(1);
			break;

		/** PIC END **/
		default:
			kerror("EXCEPTION: UNHANDLED EXCEPTION OR INTERRUPT");
			printf("Error code: 0x%2\n", stack.error_code);
			dump_cpu_state(cpu, stack);

			// TODO i would love to track spurious interrupts as well, would be useful to dump their state and specify that its spurious (if possible)
			break;
	}
}

void idt_set_descriptor(uint8_t vector, void *isr, uint8_t flags)
{
	idt_entry_t* descriptor = &idt[vector];

	descriptor->isr_low = (uint32_t) isr & 0xFFFF;
	descriptor->kernel_cs = 0x08; // this is whatever the kernel code selector is in the GDT
	descriptor->attributes = flags;
	descriptor->isr_high = (uint32_t) isr >> 16;
	descriptor->reserved = 0;
}

extern void* isr_stub_table[];
void idt_init(void)
{
#ifdef __TESTING__
	kinfo("Initializing the IDT");	
#endif
	idtr.base = (uintptr_t)&idt[0];
	idtr.limit = (uint16_t) sizeof(idt_entry_t) * IDT_MAX_DESCRIPTORS - 1;

	for (uint8_t vector = 0; vector < IDT_MAX_DESCRIPTORS; vector++) {
		idt_set_descriptor(vector, isr_stub_table[vector], 0x8E);
		vectors[vector] = true;
	}

	// The "m" indicates actual data, not a pointer
	__asm__ volatile("lidt %0" : : "m"(idtr)); // load the new IDT
						  
	PIC_remap(0x20, 0x28);
	pic_disable(); // mask everything
	IRQ_clear_mask(0);
	IRQ_clear_mask(1);
	
	__asm__ volatile("sti"); // set the interrupt flag
#ifdef __TESTING__
	kinfo("Initialized the IDT");	
#endif			 
}