// ---------------------------------------------------------------------------
// GDT + TSS — kernel-owned segment tables and interrupt stack tables (IST).
//
// The kernel boots on UEFI's GDT.  We take over with our own GDT that
// mirrors the current selector layout (flat ring-0 code/data descriptors
// at the same indices UEFI chose), so CS/DS/SS/ES keep working without
// any far-jump or segment-reload gymnastics.  On top of that we add a
// TSS with two interrupt stack tables:
//
//   - IST1: page-fault (vector 14) stack.  A stack-overflowing task
//     faults inside its unmapped guard page; the CPU pushes the fault
//     frame onto this dedicated stack instead of the faulting (also
//     unmapped) stack, so we can kill the task gracefully instead of
//     double/triple-faulting into a machine reset.
//   - IST2: double-fault (vector 8) stack.  If anything else goes wrong
//     during exception handling we land here and halt cleanly.
//
// We never change privilege level, so the TSS's RSP0..2 are unused —
// only the IST pointers matter.
// ---------------------------------------------------------------------------
use core::arch::asm;

const GDT_ENTRIES: usize = 64;
const TSS_INDEX: usize = GDT_ENTRIES - 1; // 63
const TSS_SEL: u16 = (TSS_INDEX << 3) as u16; // 0x1F8

/// 8-byte segment descriptor.
#[repr(C, packed)]
#[derive(Clone, Copy)]
struct GdtEntry {
    limit_lo: u16,
    base_lo: u16,
    base_mid: u8,
    access: u8,
    flags_limit: u8, // limit 19:16 | flags nibble (G D/B L AVL)
    base_hi: u8,
}

impl GdtEntry {
    const fn null() -> Self {
        Self {
            limit_lo: 0,
            base_lo: 0,
            base_mid: 0,
            access: 0,
            flags_limit: 0,
            base_hi: 0,
        }
    }
    /// Flat ring-0 64-bit code segment (present, exec/read, L=1).
    const fn code() -> Self {
        Self {
            limit_lo: 0xFFFF,
            base_lo: 0,
            base_mid: 0,
            access: 0x9A,
            flags_limit: 0xAF, // G=1, D/B=0, L=1, AVL=0; limit 0xF
            base_hi: 0,
        }
    }
    /// Flat ring-0 data segment (present, read/write).
    const fn data() -> Self {
        Self {
            limit_lo: 0xFFFF,
            base_lo: 0,
            base_mid: 0,
            access: 0x92,
            flags_limit: 0xCF, // G=1, D/B=1, L=0, AVL=0; limit 0xF
            base_hi: 0,
        }
    }
}

/// 64-bit TSS: only the IST pointers are used (no privilege changes).
#[repr(C, packed)]
#[derive(Clone, Copy)]
struct Tss {
    _res1: u32,
    rsp0: u64,
    rsp1: u64,
    rsp2: u64,
    _res2: u64,
    ist1: u64,
    ist2: u64,
    ist3: u64,
    ist4: u64,
    ist5: u64,
    ist6: u64,
    ist7: u64,
    _res3: u64,
    _res4: u16,
    iomap_base: u16,
}

const TSS_SIZE: usize = 104;

impl Tss {
    const fn zero() -> Self {
        Self {
            _res1: 0,
            rsp0: 0,
            rsp1: 0,
            rsp2: 0,
            _res2: 0,
            ist1: 0,
            ist2: 0,
            ist3: 0,
            ist4: 0,
            ist5: 0,
            ist6: 0,
            ist7: 0,
            _res3: 0,
            _res4: 0,
            iomap_base: 0,
        }
    }
}

/// IST stack size: the fault path (fault dump + best-effort crash-log
/// FAT write) is deep, so each IST stack is 32 KiB, allocated from the
/// heap at init.  They must NOT live in .bss: a 4 KiB .bss stack placed
/// next to kernel statics silently overwrites them on overflow (the FAT
/// ACTIVE_VOL was clobbered this way — a "no volume at slot 8" crash in
/// the fault path).  The heap region is far from the kernel image.
const IST_STACK_SIZE: usize = 32 * 1024;

static mut TSS: Tss = Tss::zero();
// GDT slots are 8 bytes except the TSS descriptor at index 63, which is
// 16 bytes — total 64*8 + 16.
static mut GDT: [u8; GDT_ENTRIES * 8 + 16] = [0; GDT_ENTRIES * 8 + 16];
static mut GDTR: [u16; 5] = [0; 5];

unsafe fn read_seg(which: u8) -> u16 {
    // `which` picks the segment register to read (0=cs, 1=ds, 2=ss, ...).
    let v: u64;
    unsafe {
        match which {
            0 => asm!("mov {0}, cs", out(reg) v, options(nostack, preserves_flags)),
            1 => asm!("mov {0}, ds", out(reg) v, options(nostack, preserves_flags)),
            2 => asm!("mov {0}, ss", out(reg) v, options(nostack, preserves_flags)),
            3 => asm!("mov {0}, es", out(reg) v, options(nostack, preserves_flags)),
            4 => asm!("mov {0}, fs", out(reg) v, options(nostack, preserves_flags)),
            _ => asm!("mov {0}, gs", out(reg) v, options(nostack, preserves_flags)),
        }
    }
    v as u16
}

/// Install the kernel's own GDT (mirroring UEFI's selectors) and a TSS
/// with IST stacks for page-fault and double-fault handling.
///
/// Returns false if UEFI's selector layout cannot be mirrored (any
/// nonzero selector at an index >= TSS_INDEX) — then we keep UEFI's GDT
/// and the IST-based fault stacks stay off (guard pages still exist but
/// an overflow would fault while pushing the fault frame → machine
/// reset instead of a graceful task death).
///
/// Must run with interrupts disabled and before `interrupts::init` (the
/// IDT's vector-8/14 entries carry the IST indices).
pub fn init() -> bool {
    unsafe {
        asm!("cli", options(nostack));

        let cs = read_seg(0);
        let ds = read_seg(1);
        let ss = read_seg(2);
        let es = read_seg(3);
        let fs = read_seg(4);
        let gs = read_seg(5);
        let sels = [cs, ds, ss, es, fs, gs];

        // Every nonzero selector must land on a descriptor slot below
        // TSS_INDEX (63 holds the 16-byte TSS descriptor).
        for &s in &sels {
            if s != 0 && (s >> 3) as usize >= TSS_INDEX {
                return false;
            }
        }

        let gdt = core::ptr::addr_of_mut!(GDT[0]);
        let write_entry = |idx: usize, e: GdtEntry| {
            let p = gdt.add(idx * 8) as *mut GdtEntry;
            core::ptr::write_unaligned(p, e);
        };

        write_entry(0, GdtEntry::null());
        write_entry((cs >> 3) as usize, GdtEntry::code());
        let cs_idx = (cs >> 3) as usize;
        for &s in &sels {
            let idx = (s >> 3) as usize;
            if s != 0 && idx != 0 && idx != cs_idx {
                write_entry(idx, GdtEntry::data());
            }
        }

        // TSS descriptor (16 bytes) at index 63.
        let tss_base = core::ptr::addr_of_mut!(TSS) as usize as u64;
        // IST stacks come from the heap (see IST_STACK_SIZE): 32 KiB each,
        // well away from the kernel image/.bss so an overflowing fault
        // path can't silently overwrite statics.
        let pf = crate::memory::allocate(IST_STACK_SIZE);
        let df = crate::memory::allocate(IST_STACK_SIZE);
        if pf.is_null() || df.is_null() {
            return false;
        }
        TSS.ist1 = pf as u64 + IST_STACK_SIZE as u64;
        TSS.ist2 = df as u64 + IST_STACK_SIZE as u64;
        TSS.iomap_base = TSS_SIZE as u16; // no I/O bitmap

        let mut td = [0u8; 16];
        td[0..2].copy_from_slice(&((TSS_SIZE - 1) as u16).to_le_bytes());
        td[2..4].copy_from_slice(&(tss_base as u16).to_le_bytes());
        td[4] = (tss_base >> 16) as u8;
        td[5] = 0x89; // present, available 64-bit TSS
        td[6] = 0x00; // limit 19:16 + flags = 0
        td[7] = (tss_base >> 24) as u8;
        td[8..12].copy_from_slice(&((tss_base >> 32) as u32).to_le_bytes());
        // td[12..16] reserved — zero.
        core::ptr::copy_nonoverlapping(td.as_ptr(), gdt.add(TSS_INDEX * 8), 16);

        // GDTR: limit = 63*8 + 16 - 1 = 519.
        GDTR = [
            (TSS_INDEX * 8 + 16 - 1) as u16,
            gdt as u16,
            (gdt as u64 >> 16) as u16,
            (gdt as u64 >> 32) as u16,
            (gdt as u64 >> 48) as u16,
        ];
        asm!("lgdt [{}]", in(reg) core::ptr::addr_of!(GDTR), options(nostack, preserves_flags));
        asm!("ltr ax", in("ax") TSS_SEL, options(nostack, preserves_flags));

        // Re-load the data segments so the CPU re-validates them against
        // the new GDT (CS cannot be reloaded by mov; its descriptor is
        // unchanged at the same index).
        if ds != 0 {
            asm!("mov ds, ax", in("ax") ds, options(nostack, preserves_flags));
        }
        if ss != 0 {
            asm!("mov ss, ax", in("ax") ss, options(nostack, preserves_flags));
        }
        if es != 0 {
            asm!("mov es, ax", in("ax") es, options(nostack, preserves_flags));
        }
        if fs != 0 {
            asm!("mov fs, ax", in("ax") fs, options(nostack, preserves_flags));
        }
        if gs != 0 {
            asm!("mov gs, ax", in("ax") gs, options(nostack, preserves_flags));
        }
        true
    }
}