// Copyright (c) 2025 shibedrill
// SPDX-License-Identifier: MIT

use core::arch::asm;
use core::ptr::addr_of_mut;

use intbits::Bits;
use lazy_static::lazy_static;
use x86_64::PrivilegeLevel::{Ring0, Ring3};
use x86_64::structures::gdt;
use x86_64::structures::gdt::*;

use crate::log::*;
use crate::memory::alloc::format;
use crate::{log_info, log_trace};

trait ModifyFlags {
    fn set_flags(self, flags: DescriptorFlags) -> Self;
}

// All this is necessary because there's no fucking method to change the
// flags in a segment descriptor. WHY??????
impl ModifyFlags for Descriptor {
    fn set_flags(self, flags: DescriptorFlags) -> Descriptor {
        match self {
            Descriptor::UserSegment(a) => {
                Descriptor::UserSegment(a.with_bits(52..=55, flags.bits().bits(0..=3)))
            }
            Descriptor::SystemSegment(a, b) => {
                log_trace!("0b{:b}", flags.bits().bits(0..=3));
                Descriptor::SystemSegment(a.with_bits(52..=55, flags.bits().bits(0..=3)), b)
            }
        }
    }
}

lazy_static! {
    pub static ref TSS: x86_64::structures::tss::TaskStateSegment =
        x86_64::structures::tss::TaskStateSegment::new();
    pub static ref GDT: x86_64::structures::gdt::GlobalDescriptorTable<64> = {
        let mut g = GlobalDescriptorTable::empty();
        g.append(Descriptor::kernel_code_segment());
        g.append(
            Descriptor::kernel_data_segment().set_flags(DescriptorFlags::from_bits_retain(0b1010)),
        );
        g.append(Descriptor::user_code_segment());
        g.append(
            Descriptor::user_data_segment().set_flags(DescriptorFlags::from_bits_retain(0b1010)),
        );
        g.append(
            Descriptor::tss_segment(&TSS).set_flags(DescriptorFlags::from_bits_retain(0b1010)),
        );
        g
    };
}
#[allow(unused)]
pub const GDT_SELECTOR_KERNEL_CODE: SegmentSelector = SegmentSelector::new(1, Ring0);
#[allow(unused)]
pub const GDT_SELECTOR_KERNEL_DATA: SegmentSelector = SegmentSelector::new(2, Ring0);
pub const GDT_SELECTOR_USER_CODE: SegmentSelector = SegmentSelector::new(3, Ring3);
pub const GDT_SELECTOR_USER_DATA: SegmentSelector = SegmentSelector::new(4, Ring3);
#[allow(unused)]
pub const GDT_SELECTOR_TSS: SegmentSelector = SegmentSelector::new(5, Ring0);

#[allow(dead_code)]
pub trait GdtEntryRead {
    fn base(&self) -> usize;
    fn limit(&self) -> usize;
    fn flags(&self) -> u8;
    fn access(&self) -> u8;
}

impl GdtEntryRead for Entry {
    fn base(&self) -> usize {
        ((self.raw().bits(56..=63) << 24) & self.raw().bits(16..=39)) as usize
    }
    fn limit(&self) -> usize {
        ((self.raw().bits(48..=51) << 16) & self.raw().bits(16..=31)) as usize
    }
    fn access(&self) -> u8 {
        self.raw().bits(40..=47) as u8
    }
    fn flags(&self) -> u8 {
        self.raw().bits(52..=55) as u8
    }
}

#[allow(dead_code)]
pub fn read_gdt() -> GlobalDescriptorTable {
    let mut addr: u128 = 0;
    unsafe {
        asm! {
            "sgdt [rax]",
            in("rax") addr_of_mut!(addr),
        }
    }
    // x86 gdtr: Lower 2 8-bit bytes (lower 16 bits) is the size of the GDT,
    // in 8-bit bytes, minus one.
    // Bits 16 through 79 are the virtual address of the GDT.
    let limit = addr.bits(0..=15) as usize + 1;
    let base = addr.bits(16..=79) as u64;
    let gdt_raw: &[u64] =
        unsafe { core::slice::from_raw_parts(base as *const u64, (limit / 8) - 1) };
    GlobalDescriptorTable::from_raw_entries(gdt_raw)
}

pub fn dump_gdt() {
    let gdt = crate::arch::x86_64::gdt::read_gdt();
    log_info!("GDT length: {}", gdt.limit() + 1);
    for entry in gdt.entries() {
        let flags = x86_64::structures::gdt::DescriptorFlags::from_bits_retain(entry.raw());
        let seg_type = SegmentType::try_from(entry);
        log_trace!(
            "GDT entry: type: {:?} (ring {}), width: {}, length: 0x{:x}, raw flags: 0b{:04b}, access: 0b{:08b}",
            seg_type.unwrap(),
            entry.access().bits(5..=6),
            if flags.contains(gdt::DescriptorFlags::DEFAULT_SIZE) {
                "32-bit"
            } else if flags.contains(gdt::DescriptorFlags::LONG_MODE) {
                "64-bit"
            } else {
                "16-bit"
            },
            entry.limit(),
            entry.raw().bits(52..=55),
            entry.access(),
        );
    }
}

#[derive(Debug)]
enum SegmentType {
    Null,
    Code,
    Data,
    Ldt,
    Tss16Avail,
    Tss16Busy,
    Tss32Avail,
    Tss32Busy,
    Tss64Avail,
    Tss64Busy,
}

impl TryFrom<&gdt::Entry> for SegmentType {
    type Error = ();
    fn try_from(value: &gdt::Entry) -> Result<Self, ()> {
        let flags = x86_64::structures::gdt::DescriptorFlags::from_bits_retain(value.raw());
        if value.access().bit(4) {
            if value.access().bit(3) {
                Ok(SegmentType::Code)
            } else {
                Ok(SegmentType::Data)
            }
        } else {
            match value.access().bits(0..=3) {
                0b0000 => Ok(SegmentType::Null),
                0b0001 => Ok(SegmentType::Tss16Avail),
                0b0010 => Ok(SegmentType::Ldt),
                0b0011 => Ok(SegmentType::Tss16Busy),
                0b1001 => {
                    if flags.contains(gdt::DescriptorFlags::LONG_MODE) {
                        Ok(SegmentType::Tss64Avail)
                    } else {
                        Ok(SegmentType::Tss32Avail)
                    }
                }
                0b1011 => {
                    if flags.contains(gdt::DescriptorFlags::LONG_MODE) {
                        Ok(SegmentType::Tss64Busy)
                    } else {
                        Ok(SegmentType::Tss32Busy)
                    }
                }
                _ => Err(()),
            }
        }
    }
}