#include "gdt.h"
#include "arch.h"
#include "fb.h"
#include "dbg.h"

struct gdt_entry {
    uint16_t limit_low;
    uint16_t base_low;
    uint8_t  base_mid;
    uint8_t  access;
    uint8_t  flags_limit;
    uint8_t  base_high;
} __attribute__((packed));

struct tss_descriptor {
    uint16_t limit_low;
    uint16_t base_low;
    uint8_t  base_mid;
    uint8_t  access;
    uint8_t  flags_limit;
    uint8_t  base_high;
    uint32_t base_upper;
    uint32_t reserved;
} __attribute__((packed));

struct gdt_ptr {
    uint16_t limit;
    uint64_t base;
} __attribute__((packed));

struct tss {
    uint32_t reserved0;
    uint64_t rsp[3];
    uint64_t reserved1;
    uint64_t ist[7];
    uint64_t reserved2;
    uint16_t reserved3;
    uint16_t io_map_base;
} __attribute__((packed));

static struct tss tss __attribute__((aligned(16))) = {0};

/* IST stack for crash/exception handlers — survives kernel stack overflow */
static uint64_t ist_crash_stack[4096] __attribute__((aligned(16)));
#define IST_CRASH_TOP ((uint64_t)&ist_crash_stack[4096])

/* Нужна GDT как массив 8-байтовых записей плюс дескриптор TSS (16 байт) */
/* Итого: 6 записей × 8 = 48 байт для 6 дескрипторов, но TSS требует 16 байт */
/* Используем плоский массив и аккуратно вычисляем смещения */

#define GDT_ENTRIES 7
static uint64_t gdt[GDT_ENTRIES] __attribute__((aligned(16)));

static void gdt_set_entry(int idx, uint32_t base, uint32_t limit, uint8_t access, uint8_t flags) {
    struct gdt_entry *e = (struct gdt_entry*)&gdt[idx];
    e->limit_low = limit & 0xFFFF;
    e->base_low = base & 0xFFFF;
    e->base_mid = (base >> 16) & 0xFF;
    e->access = access;
    e->flags_limit = (flags << 4) | ((limit >> 16) & 0x0F);
    e->base_high = (base >> 24) & 0xFF;
}

static void gdt_set_tss(int idx, uint64_t base, uint32_t limit, uint8_t access) {
    struct tss_descriptor *t = (struct tss_descriptor*)&gdt[idx];
    t->limit_low = limit & 0xFFFF;
    t->base_low = base & 0xFFFF;
    t->base_mid = (base >> 16) & 0xFF;
    t->access = access;
    t->flags_limit = ((limit >> 16) & 0x0F);
    t->base_high = (base >> 24) & 0xFF;
    t->base_upper = (base >> 32) & 0xFFFFFFFF;
    t->reserved = 0;
}

void gdt_init(void) {
    /* NULL-дескриптор */
    gdt[0] = 0;

    /* Код ядра: base=0, limit=0xFFFFF, access=0x9A (присутствует, ring0, код, exec/read), flags=0xA (64-бита, гранулярность 4K) */
    gdt_set_entry(1, 0, 0xFFFFF, 0x9A, 0xA);

    /* Данные ядра: access=0x92 (присутствует, ring0, данные, read/write), flags=0xA */
    gdt_set_entry(2, 0, 0xFFFFF, 0x92, 0xA);

    /* Код пользователя: access=0xFA (присутствует, ring3, код, exec/read), flags=0xA */
    gdt_set_entry(3, 0, 0xFFFFF, 0xFA, 0xA);

    /* Данные пользователя: access=0xF2 (присутствует, ring3, данные, read/write), flags=0xA */
    gdt_set_entry(4, 0, 0xFFFFF, 0xF2, 0xA);

    /* TSS: 104 байта */
    tss.io_map_base = 0xFFFF;
    tss.ist[0] = IST_CRASH_TOP; /* IST1: dedicated stack for exception handlers */
    uint64_t tss_base = (uint64_t)&tss;
    gdt_set_tss(5, tss_base, sizeof(struct tss) - 1, 0x89);

    dbg_printf("GDT: TSS at %p\n", tss_base);

    struct gdt_ptr ptr = {
        .limit = sizeof(gdt) - 1,
        .base = (uint64_t)gdt
    };

    dbg_printf("GDT: loading...\n");

    /* Отключаем прерывания перед загрузкой новой GDT, чтобы
       записи IDT (ссылающиеся на UEFI CS) не сработали с нашей GDT */
    cli();
    lgdt(&ptr);

    dbg_printf("GDT: lgdt OK (entries=%d)\n", GDT_ENTRIES);

    /* Перезагружаем CS через дальний возврат — обязаны использовать отдельные регистры */
    {
        uint64_t ret_addr;
        uint64_t cs_val = (uint64_t)GDT_KERNEL_CODE;
        __asm__ volatile(
            "lea 1f(%%rip), %0\n"
            "push %1\n"
            "push %0\n"
            "lretq\n"
            "1:\n"
            : "=&r"(ret_addr) : "r"(cs_val) : "memory"
        );
    }

    dbg_printf("GDT: CS reloaded (cs=0x%x)\n", GDT_KERNEL_CODE);

    /* Перезагружаем сегменты данных */
    __asm__ volatile(
        "mov %0, %%ds\n"
        "mov %0, %%es\n"
        "mov %0, %%ss\n"
        "mov %0, %%fs\n"
        "mov %0, %%gs\n"
        : : "r"((uint64_t)GDT_KERNEL_DATA) : "memory"
    );

    dbg_printf("GDT: segments reloaded (ds=0x%x)\n", GDT_KERNEL_DATA);

    /* Загружаем TSS */
    ltr(GDT_TSS);

    dbg_printf("GDT: TSS loaded\n");
}
