//! Global Descriptor Table (GDT) and Task State Segment (TSS) for x86_64 SMP
//! Supports up to 4 CPU Cores with dedicated TSS descriptors.

use core::mem::size_of;

pub const MAX_CPUS: usize = 4;

#[repr(C, packed)]
pub struct TaskStateSegment {
    _reserved1: u32,
    pub rsp0: u64,
    pub rsp1: u64,
    pub rsp2: u64,
    _reserved2: u64,
    pub ist: [u64; 7],
    _reserved3: u64,
    _reserved4: u16,
    pub iomap_base: u16,
}

impl TaskStateSegment {
    pub const fn new() -> Self {
        Self {
            _reserved1: 0,
            rsp0: 0,
            rsp1: 0,
            rsp2: 0,
            _reserved2: 0,
            ist: [0; 7],
            _reserved3: 0,
            _reserved4: 0,
            iomap_base: size_of::<TaskStateSegment>() as u16,
        }
    }
}

#[repr(C, packed)]
struct GdtDescriptor {
    limit: u16,
    base: u64,
}

#[no_mangle]
pub static mut TSS: TaskStateSegment = TaskStateSegment::new();

#[no_mangle]
pub static mut PER_CPU_TSS: [TaskStateSegment; MAX_CPUS] = [const { TaskStateSegment::new() }; MAX_CPUS];

// GDT entries:
// 0: Null
// 1: Kernel Code 64 (0x08)
// 2: Kernel Data 64 (0x10)
// 3: User Data 64   (0x18)
// 4: User Code 64   (0x20)
// 5,6: TSS CPU 0    (0x28)
// 7,8: TSS CPU 1    (0x38)
// 9,10: TSS CPU 2   (0x48)
// 11,12: TSS CPU 3  (0x58)
// 13,14,15: padding
static mut GDT: [u64; 16] = [
    0x0000000000000000, // 0x00: Null
    0x00af9a000000ffff, // 0x08: Kernel Code 64 (Ring 0)
    0x00cf92000000ffff, // 0x10: Kernel Data 64 (Ring 0)
    0x00cff2000000ffff, // 0x18: User Data 64   (Ring 3)
    0x00affa000000ffff, // 0x20: User Code 64   (Ring 3)
    0, 0,               // 0x28: TSS CPU 0
    0, 0,               // 0x38: TSS CPU 1
    0, 0,               // 0x48: TSS CPU 2
    0, 0,               // 0x58: TSS CPU 3
    0, 0, 0,            // Padding
];

pub unsafe fn init() {
    let gdt_ptr = core::ptr::addr_of_mut!(GDT) as *mut u64;

    // Populate TSS descriptors for all supported CPU cores
    for cpu in 0..MAX_CPUS {
        let tss_addr = if cpu == 0 {
            core::ptr::addr_of!(TSS) as u64
        } else {
            core::ptr::addr_of!(PER_CPU_TSS[cpu]) as u64
        };
        let tss_size = (size_of::<TaskStateSegment>() - 1) as u64;

        let tss_low = (tss_size & 0xffff)
            | ((tss_addr & 0xffff) << 16)
            | (((tss_addr >> 16) & 0xff) << 32)
            | (0x89u64 << 40) // Present, 64-bit TSS (Available)
            | (((tss_size >> 16) & 0xf) << 48)
            | (((tss_addr >> 24) & 0xff) << 56);
        let tss_high = tss_addr >> 32;

        let entry_idx = 5 + cpu * 2;
        *gdt_ptr.add(entry_idx) = tss_low;
        *gdt_ptr.add(entry_idx + 1) = tss_high;
    }

    let descriptor = GdtDescriptor {
        limit: (size_of::<[u64; 16]>() - 1) as u16,
        base: gdt_ptr as u64,
    };

    core::arch::asm!(
        "lgdt [{0}]",
        "push 0x08",
        "lea {tmp}, [2f + rip]",
        "push {tmp}",
        "retfq",
        "2:",
        "mov {tmp:e}, 0x10",
        "mov ds, {tmp:x}",
        "mov es, {tmp:x}",
        "mov ss, {tmp:x}",
        "mov fs, {tmp:x}",
        "mov gs, {tmp:x}",
        "ltr {tss_sel:x}",
        in(reg) &descriptor,
        tmp = out(reg) _,
        tss_sel = in(reg) (0x28u16),
        options(preserves_flags)
    );
}

/// Initialize GDT and load dedicated TSS for an Application Processor (AP)
pub unsafe fn init_ap(cpu_id: usize) {
    let gdt_ptr = core::ptr::addr_of_mut!(GDT) as *mut u64;
    let descriptor = GdtDescriptor {
        limit: (size_of::<[u64; 16]>() - 1) as u16,
        base: gdt_ptr as u64,
    };

    let tss_sel = (0x28 + (cpu_id.min(MAX_CPUS - 1) * 16)) as u16;

    core::arch::asm!(
        "lgdt [{0}]",
        "push 0x08",
        "lea {tmp}, [2f + rip]",
        "push {tmp}",
        "retfq",
        "2:",
        "mov {tmp:e}, 0x10",
        "mov ds, {tmp:x}",
        "mov es, {tmp:x}",
        "mov ss, {tmp:x}",
        "mov fs, {tmp:x}",
        "mov gs, {tmp:x}",
        "ltr {tss_sel:x}",
        in(reg) &descriptor,
        tmp = out(reg) _,
        tss_sel = in(reg) tss_sel,
        options(preserves_flags)
    );
}

#[no_mangle]
pub extern "C" fn set_kernel_stack(stack_top: u64) {
    unsafe {
        let tss_ptr = core::ptr::addr_of_mut!(TSS);
        (*tss_ptr).rsp0 = stack_top;
    }
}