use core::{
    arch::asm,
    fmt::Debug,
    mem::{MaybeUninit, offset_of},
    ops::{Deref, DerefMut},
};

use bit_field::BitField;

use crate::{
    serial_println,
    sync::LazyLock,
    x86_64::{idt::ToAddress, instructions::*},
};

#[repr(C)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct GlobalDescriptorTable {
    pub null: RawGdtEntry,
    pub kernel_code: RawGdtEntry,
    pub kernel_data: RawGdtEntry,
    pub tss: TssEntry,
}

impl GlobalDescriptorTable {
    pub fn new() -> Self {
        let tss = &*TSS;
        GlobalDescriptorTable {
            null: RawGdtEntry::new(0, 0, 0, 0),
            kernel_code: RawGdtEntry::new(0xfffff, 0, 0, 0)
                .with_segment(true)
                .with_privilege_level(RING0)
                .with_present(true)
                .with_segment_kind(
                    CodeDataSegmentKind::CodeSegment
                        | CodeDataSegmentKind::Executable
                        | CodeDataSegmentKind::ReadWrite
                        | CodeDataSegmentKind::Accessed,
                )
                .with_present(true)
                .with_code_segment64(true)
                .with_granularity(true),
            kernel_data: RawGdtEntry::new(0xfffff, 0, 0, 0)
                .with_segment(true)
                .with_privilege_level(RING0)
                .with_present(true)
                .with_segment_kind(CodeDataSegmentKind::ReadWrite | CodeDataSegmentKind::Accessed)
                .with_code_segment64(true)
                .with_granularity(true),
            tss: TssEntry::new(
                core::mem::size_of::<TaskStateSegment>() as u32 - 1,
                tss as *const TaskStateSegment as u64,
            ),
        }
    }

    pub fn load(&'static self) {
        serial_println!("Loading GDT {:#?}", self);
        unsafe { Self::load_unsafe(self) }
    }

    /// # Safety
    /// The caller must ensure that `this` is a valid pointer to a `GlobalDescriptorTable` that remains valid for the lifetime of the gdt.
    pub unsafe fn load_unsafe(this: *const Self) {
        let gdt_register = GdtRegister {
            limit: (core::mem::size_of::<Self>() - 1) as u16,
            base: this as u64,
        };

        unsafe {
            core::arch::asm!(
                "lgdt [{}]",
                in(reg) &gdt_register,
                options(readonly, nostack, preserves_flags)
            );

            let gs = msr::read_msr(msr::MSR_GS_BASE);

            asm!(
                "push {code_seg}",
                "lea rax, [rip + 2f]",
                "push rax",
                "retfq",
                "2:",
                "mov ax, {data_seg}",
                "mov ds, ax",
                "mov es, ax",
                "mov fs, ax",
                "mov gs, ax",
                "mov ss, ax",
                code_seg = const offset_of!(GlobalDescriptorTable, kernel_code),
                data_seg = const offset_of!(GlobalDescriptorTable, kernel_data),
                lateout("rax") _,
            );

            msr::write_msr(msr::MSR_GS_BASE, gs);

            asm!(
                "mov ax, {tss_seg}",
                "ltr ax",
                tss_seg = const offset_of!(GlobalDescriptorTable, tss),
                lateout("rax") _,
            );
        }
    }
}

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

impl Default for GlobalDescriptorTable {
    fn default() -> Self {
        Self::new()
    }
}

#[repr(u8)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SystemSegmentKind {
    Reserved = 0b0000,
    Ldt = 0b0010,
    TssAvailable = 0b1001,
    TssBusy = 0b1011,
    CallGate = 0b1100,
    InterruptGate = 0b1110,
    TrapGate = 0b1111,
}

pub const RING0: u8 = 0;
pub const RING1: u8 = 1;
pub const RING2: u8 = 2;
pub const RING3: u8 = 3;

bitflags::bitflags! {
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub struct CodeDataSegmentKind: u8 {
        const CodeSegment = 1 << 3;
        const Executable = 1 << 2;
        /// `ReadWrite` is set for code segments if they are readable, and for
        /// data segments if they are writable.
        const ReadWrite = 1 << 1;
        const Accessed = 1 << 0;
    }
}

impl SystemSegmentKind {
    pub fn into_u8(self) -> u8 {
        self as u8
    }
    pub fn from_u8_or_reserved(value: u8) -> Self {
        match value {
            0b0010 => SystemSegmentKind::Ldt,
            0b1001 => SystemSegmentKind::TssAvailable,
            0b1011 => SystemSegmentKind::TssBusy,
            0b1100 => SystemSegmentKind::CallGate,
            0b1110 => SystemSegmentKind::InterruptGate,
            0b1111 => SystemSegmentKind::TrapGate,
            _ => SystemSegmentKind::Reserved,
        }
    }

    /// # Safety
    /// the caller must ensure that the lower 4 bits of `value` are a valid `GdtKind` discriminant.
    pub unsafe fn from_u8_unchecked(value: u8) -> Self {
        unsafe { core::mem::transmute(value & 0b1111) }
    }
}

#[repr(C)]
#[derive(Clone, Copy, PartialEq, Eq)]
pub struct RawGdtEntry {
    pub limit_low: u16,
    pub base_low: u16,
    pub base_middle: u8,
    pub access: u8,
    pub flags: u8,
    pub base_high: u8,
}

fn debug_commont_gdt_fields<'a, 'b>(
    entry: &RawGdtEntry,
    dbg: &'a mut core::fmt::DebugStruct<'a, 'b>,
) -> &'a mut core::fmt::DebugStruct<'a, 'b> {
    dbg.field_with("limit", |f| {
        if entry.granularity() {
            f.write_fmt(format_args!("0x{:x} (4KiB blocks)", entry.limit()))
        } else {
            f.write_fmt(format_args!("0x{:x} (bytes)", entry.limit()))
        }
    })
    .field("segment", &entry.segment());

    if entry.segment() {
        dbg.field("segment_kind", &entry.segment_kind());
    } else {
        dbg.field("system_segment_kind", &entry.system_segment_kind());
    }

    dbg.field("privilege_level", &entry.privilege_level())
        .field("present", &entry.present())
        .field("available", &entry.available())
        .field("code_segment64", &entry.code_segment64())
        .field("default_operation_size", &entry.default_operation_size())
        .field_with("granularity", |f| {
            if entry.granularity() {
                f.write_str("4KiB")
            } else {
                f.write_str("bytes")
            }
        });

    dbg
}

impl Debug for RawGdtEntry {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let mut dbg = f.debug_struct("RawGdtEntry");

        dbg.field_with("base", |f| {
            f.write_fmt(format_args!("0x{:x}", self.base32()))
        });
        let dbg = debug_commont_gdt_fields(self, &mut dbg);

        dbg.finish()
    }
}

#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DefaultOperationSize {
    Bit16 = 0,
    Bit32 = 1,
}

impl DefaultOperationSize {
    pub fn into_u8(self) -> u8 {
        self as u8
    }
    pub fn from_bool(value: bool) -> Self {
        if value {
            DefaultOperationSize::Bit32
        } else {
            DefaultOperationSize::Bit16
        }
    }
}

impl AsRef<RawGdtEntry> for TssEntry {
    fn as_ref(&self) -> &RawGdtEntry {
        &self.gdt
    }
}

impl RawGdtEntry {
    pub fn new(limit: u32, base: u32, access: u8, flags: u8) -> Self {
        let limit_low = limit as u16;
        let limit_high = (limit >> 16) as u8;
        let base_low = base as u16;
        let base_middle = (base >> 16) as u8;
        let base_high = (base >> 24) as u8;

        RawGdtEntry {
            limit_low,
            base_low,
            base_middle,
            access,
            flags: (flags & 0xf0) | limit_high,
            base_high,
        }
    }
    pub fn empty() -> Self {
        RawGdtEntry {
            limit_low: 0,
            base_low: 0,
            base_middle: 0,
            access: 0,
            flags: 0,
            base_high: 0,
        }
    }

    pub fn with_segment_kind(mut self, kind: CodeDataSegmentKind) -> Self {
        self.set_segment_kind(kind);
        self
    }
    pub fn with_system_segment_kind(mut self, kind: SystemSegmentKind) -> Self {
        self.set_system_segment_kind(kind);
        self
    }
    pub fn with_segment(mut self, segment: bool) -> Self {
        self.set_segment(segment);
        self
    }
    pub fn with_privilege_level(mut self, level: u8) -> Self {
        self.set_privilege_level(level);
        self
    }
    pub fn with_present(mut self, present: bool) -> Self {
        self.set_present(present);
        self
    }

    pub fn system_segment_kind(&self) -> SystemSegmentKind {
        SystemSegmentKind::from_u8_or_reserved(self.access.get_bits(0..4))
    }
    pub fn set_system_segment_kind(&mut self, kind: SystemSegmentKind) {
        self.access.set_bits(0..4, kind.into_u8());
    }
    /// `segment_kind` is valid if `segment` is set, otherwise
    /// `system_segment_kind` is valid. `segment_kind` is the type of segment,
    /// and is only valid for code and data segments.
    pub fn segment_kind(&self) -> CodeDataSegmentKind {
        CodeDataSegmentKind::from_bits_truncate(self.access.get_bits(0..4))
    }
    pub fn set_segment_kind(&mut self, kind: CodeDataSegmentKind) {
        self.access.set_bits(0..4, kind.bits());
    }

    /// `segment` is set if this is a code or data segment, and clear if it is a system segment.
    pub fn segment(&self) -> bool {
        self.access.get_bit(4)
    }
    pub fn set_segment(&mut self, system_segment: bool) {
        self.access.set_bit(4, system_segment);
    }
    pub fn privilege_level(&self) -> u8 {
        self.access.get_bits(5..7)
    }
    pub fn set_privilege_level(&mut self, level: u8) {
        self.access.set_bits(5..7, level);
    }
    pub fn present(&self) -> bool {
        self.access.get_bit(7)
    }
    pub fn set_present(&mut self, present: bool) {
        self.access.set_bit(7, present);
    }

    pub fn with_available(mut self, available: bool) -> Self {
        self.set_available(available);
        self
    }
    pub fn with_code_segment64(mut self, code_segment64: bool) -> Self {
        self.set_code_segment64(code_segment64);
        self
    }
    pub fn with_default_operation_size(mut self, size: DefaultOperationSize) -> Self {
        self.set_default_operation_size(size);
        self
    }
    pub fn with_granularity(mut self, granularity: bool) -> Self {
        self.set_granularity(granularity);
        self
    }
    pub fn with_limit_high(mut self, limit_high: u8) -> Self {
        self.set_limit_high(limit_high);
        self
    }

    fn limit_high(&self) -> u8 {
        self.flags.get_bits(0..4)
    }
    fn set_limit_high(&mut self, limit_high: u8) {
        self.flags.set_bits(0..4, limit_high);
    }
    pub fn available(&self) -> bool {
        self.flags.get_bit(4)
    }
    pub fn set_available(&mut self, available: bool) {
        self.flags.set_bit(4, available);
    }
    pub fn code_segment64(&self) -> bool {
        self.flags.get_bit(5)
    }
    pub fn set_code_segment64(&mut self, code_segment64: bool) {
        self.flags.set_bit(5, code_segment64);
    }

    /// `default_operation_size` is the default size of operands for this segment. If `default_operation_size` is set, the default size is 32 bits, otherwise it is 16 bits.
    /// On 64-bit mode, this should be 1.
    pub fn default_operation_size(&self) -> DefaultOperationSize {
        DefaultOperationSize::from_bool(self.flags.get_bit(6))
    }
    pub fn set_default_operation_size(&mut self, size: DefaultOperationSize) {
        self.flags.set_bit(6, size.into_u8() != 0);
    }

    /// `granularity` is the unit of the limit field. If `granularity` is set,
    /// the limit is in 4KiB blocks, otherwise it is in bytes.
    pub fn granularity(&self) -> bool {
        self.flags.get_bit(7)
    }
    pub fn set_granularity(&mut self, granularity: bool) {
        self.flags.set_bit(7, granularity);
    }

    pub fn limit(&self) -> u32 {
        let limit_low = self.limit_low as u32;
        let limit_high = (self.limit_high() as u32) << 16;
        limit_low | limit_high
    }

    pub fn base32(&self) -> u32 {
        let base_low = self.base_low as u32;
        let base_middle = (self.base_middle as u32) << 16;
        let base_high = (self.base_high as u32) << 24;
        base_low | base_middle | base_high
    }
}

#[repr(C)]
#[derive(Clone, Copy, PartialEq, Eq)]
pub struct TssEntry {
    gdt: RawGdtEntry,
    pub base_ext: u32,
    reserved: u32,
}

impl Debug for TssEntry {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let mut dbg = f.debug_struct("TssEntry");

        dbg.field_with("base", |f| {
            f.write_fmt(format_args!("0x{:x}", self.base64()))
        });
        let dbg = debug_commont_gdt_fields(&self.gdt, &mut dbg);

        dbg.finish()
    }
}

impl TssEntry {
    pub fn base64(&self) -> u64 {
        let base_low = self.gdt.base32() as u64;
        let base_ext = self.base_ext as u64;
        (base_ext << 32) | base_low
    }

    pub fn new(limit: u32, base: u64) -> Self {
        let base_low = base as u32;
        let base_ext = (base >> 32) as u32;

        TssEntry {
            gdt: RawGdtEntry::new(limit, base_low, 0, 0)
                .with_present(true)
                .with_privilege_level(RING0)
                .with_system_segment_kind(SystemSegmentKind::TssAvailable),
            base_ext,
            reserved: 0,
        }
    }
}

impl Deref for TssEntry {
    type Target = RawGdtEntry;

    fn deref(&self) -> &Self::Target {
        &self.gdt
    }
}

impl DerefMut for TssEntry {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.gdt
    }
}

const PRIVILEGE_STACK_TABLE_SIZE: usize = 3;
const INTERRUPT_STACK_TABLE_SIZE: usize = 7;

#[repr(C, packed(4))]
pub struct TaskStateSegment {
    _reserved1: [u8; 4],
    pub privilege_stack_table: [u64; PRIVILEGE_STACK_TABLE_SIZE],
    _reserved2: [u8; 8],
    pub interrupt_stack_table: [u64; INTERRUPT_STACK_TABLE_SIZE],
    _reserved3: [u8; 10],
    pub iomap_base: u16,
}

impl Debug for TaskStateSegment {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let mut dbg = f.debug_struct("TaskStateSegment");

        let rsps = self.privilege_stack_table;
        let ists = self.interrupt_stack_table;

        dbg.field("RSP0", &format_args!("{:#x}", rsps[0]))
            .field("RSP1", &format_args!("{:#x}", rsps[1]))
            .field("RSP2", &format_args!("{:#x}", rsps[2]));

        dbg.field("IST1", &format_args!("{:#x}", ists[0]))
            .field("IST2", &format_args!("{:#x}", ists[1]))
            .field("IST3", &format_args!("{:#x}", ists[2]))
            .field("IST4", &format_args!("{:#x}", ists[3]))
            .field("IST5", &format_args!("{:#x}", ists[4]))
            .field("IST6", &format_args!("{:#x}", ists[5]))
            .field("IST7", &format_args!("{:#x}", ists[6]));

        dbg.field("iomap_base", &format_args!("{:#x}", self.iomap_base));

        dbg.finish()
    }
}

impl TaskStateSegment {
    pub const fn new() -> Self {
        Self {
            privilege_stack_table: [0; 3],
            interrupt_stack_table: [0; 7],
            iomap_base: size_of::<TaskStateSegment>() as u16,
            _reserved1: [0; 4],
            _reserved2: [0; 8],
            _reserved3: [0; 10],
        }
    }
    pub fn set_stack(&mut self, index: u8, stack: &'static Stack) {
        assert!(
            (index as usize) < INTERRUPT_STACK_TABLE_SIZE,
            "Interrupt stack table index out of bounds"
        );
        self.interrupt_stack_table[index as usize] = stack.top().to_address();
    }
}

const impl Default for TaskStateSegment {
    fn default() -> Self {
        Self::new()
    }
}

pub const DF_STACK: u8 = 0;
pub const NMI_STACK: u8 = 1;
pub const MC_STACK: u8 = 2;
pub const PF_STACK: u8 = 3;

const STACK_SIZE: usize = super::PAGE_SIZE * 5;
#[repr(align(4096))]
pub struct Stack {
    _bytes: MaybeUninit<[u8; STACK_SIZE]>,
}

impl Stack {
    const fn new() -> Self {
        Stack {
            _bytes: MaybeUninit::uninit(),
        }
    }
    fn top(&self) -> *const u8 {
        unsafe {
            (&raw const *self)
                .cast::<u8>()
                .byte_add(super::PAGE_SIZE * 5)
        }
    }
}

pub static TSS: LazyLock<TaskStateSegment> = LazyLock::new(|| {
    let mut tss = TaskStateSegment::new();
    static mut STACKS: [Stack; 4] = [const { Stack::new() }; 4];

    tss.set_stack(DF_STACK, unsafe { &STACKS[DF_STACK as usize] });
    tss.set_stack(NMI_STACK, unsafe { &STACKS[NMI_STACK as usize] });
    tss.set_stack(MC_STACK, unsafe { &STACKS[MC_STACK as usize] });
    tss.set_stack(PF_STACK, unsafe { &STACKS[PF_STACK as usize] });

    serial_println!("Initialized TSS {:#?}", tss);

    tss
});