#include "aplos.h"

#define DESCRIPTOR_SIZE 8
#define IDT_ENTRIES 256

enum {
	GDT_START,
		NULL_SEGMENT = GDT_START,
		KERNEL_CODE_SEGMENT,
		KERNEL_DATA_SEGMENT,
		TASK_STATE_SEGMENT,
		TASK_STATE_SEGMENT_CONT,
	GDT_END = TASK_STATE_SEGMENT_CONT,

	IDT_START,
	IDT_END = IDT_START + (2*IDT_ENTRIES) - 1,

	TOTAL_ENTRY_COUNT
};

enum {
	TYPE_NULL_SEGMENT = 0x10,
	TYPE_DATA_SEGMENT = 0x12,
	TYPE_CODE_SEGMENT = 0x18,
	TYPE_TSS_SEGMENT  = 0x09,
	TYPE_IDT          = 0x0E,

	USER_TYPE_MASK = 0x10
};

struct descriptor
{
	int type;

	/* TASK_STATE_SEGMENT */
	struct encoded_tss_entry *tss;

	/* IDT */
	void (*isr)(void);
};

struct encoded_descriptor
{
	alignas(DESCRIPTOR_SIZE) uint8_t data[DESCRIPTOR_SIZE];
};
static_assert(sizeof(struct encoded_descriptor) == DESCRIPTOR_SIZE);

struct encoded_tss_entry
{
	uint8_t data[104]; /* Not part of this is used yet... */
};

struct table_reg
{
	uint8_t data[10];
};
static_assert(sizeof(struct table_reg) == 10);

static struct encoded_tss_entry tss;

static struct descriptor entries[TOTAL_ENTRY_COUNT] = {
	[NULL_SEGMENT] = {
		.type = TYPE_NULL_SEGMENT
	},
	[KERNEL_CODE_SEGMENT] = {
		.type = TYPE_CODE_SEGMENT
	},
	[KERNEL_DATA_SEGMENT] = {
		.type = TYPE_DATA_SEGMENT
	},
	[TASK_STATE_SEGMENT] = {
		.type = TYPE_TSS_SEGMENT,
		.tss = &tss,
	},
};

static struct encoded_descriptor encoded[nelem(entries)];

static void encode_descriptors(struct table_reg *, uint64_t, uint64_t);
static void encode_descriptor(struct descriptor *, struct encoded_descriptor *);

void
setup_descriptors(void)
{
	struct table_reg reg;

	disable_interrupts();
	encode_descriptors(&reg, GDT_START, GDT_END);
	set_gdt(&reg, DESCRIPTOR_SIZE*KERNEL_CODE_SEGMENT, DESCRIPTOR_SIZE*KERNEL_DATA_SEGMENT, DESCRIPTOR_SIZE*TASK_STATE_SEGMENT);

	for(uint64_t i = IDT_START; i < IDT_END; i += 2)
		entries[i].type = TYPE_IDT;
	for(uint64_t i = 0; i < nelem(isr_stubs); i++)
		entries[IDT_START+i*2].isr = isr_stubs[i];

	encode_descriptors(&reg, IDT_START, IDT_END);
	set_idt(&reg);

	enable_interrupts();
}

static void
encode_descriptors(struct table_reg *reg, uint64_t start, uint64_t end)
{
	uint64_t count = 1 + end - start;
	struct descriptor *d = entries + start;
	struct encoded_descriptor *e = encoded + start;


	memset(e, 0, sizeof(*e) * count);
	for(uint64_t i = 0; i < count; i++){
		encode_descriptor(d+i, e+i);
		if(!(d[i].type & USER_TYPE_MASK))
			i++; /* System types are twice as big */
	}

	write_uint16_le(reg->data, count*DESCRIPTOR_SIZE - 1);
	write_uint64_le(reg->data+2, (uint64_t)e);
}

static void
encode_descriptor(struct descriptor *d, struct encoded_descriptor *e)
{
	e->data[5] |= d->type&0x1F;		/* Type */
	e->data[5] |= 1<<7;			/* Present */
	e->data[6] |= 1<<5;			/* Long mode */

	uint64_t v;

	switch(d->type){
	case TYPE_TSS_SEGMENT: 
		write_uint16_le(e->data, sizeof(*d->tss));

		v = (uint64_t)d->tss;
		write_uint16_le(e->data+2, v & 0xFFFF);
		e->data[4] = (v >> 16) & 0xFF;
		e->data[7] = (v >> 24) & 0xFF;
		write_uint32_le((uint8_t *)(e+1), (v >> 32) & 0xFFFFFFFF);
		break;
	case TYPE_IDT:
		v = (uint64_t)d->isr;
		write_uint16_le(e->data, v & 0xFFFF);
		write_uint16_le(e->data+6, (v >> 16) & 0xFFFF);
		write_uint32_le(e->data+8, (v >> 32) & 0xFFFFFFFF);

		write_uint16_le(e->data+2, DESCRIPTOR_SIZE*KERNEL_CODE_SEGMENT);
	}	
}