// 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::structures::gdt;
use x86_64::structures::gdt::*;
use x86_64::structures::tss::TaskStateSegment;

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) => {
                Descriptor::SystemSegment(a.with_bits(52..=55, flags.bits().bits(0..=3)), b)
            }
        }
    }
}

trait ModifyAccess {
    fn set_access(self, access: u8) -> Self;
}

impl ModifyAccess for Descriptor {
    fn set_access(self, access: u8) -> Self {
        match self {
            Descriptor::UserSegment(_) => todo!(),
            Descriptor::SystemSegment(a, b) => {
                Descriptor::SystemSegment(a.with_bits(40..=47, access.bits(0..=7).into()), b)
            },
        }
    }
}

pub struct Selectors {
    pub sel_kernel_code: SegmentSelector,
    pub sel_kernel_stack: SegmentSelector,
    pub sel_user_code: SegmentSelector,
    pub sel_user_stack: SegmentSelector,
    #[allow(unused)]
    pub sel_tss: SegmentSelector,
    pub gdt: GlobalDescriptorTable<48>,
}

pub const TSS: TaskStateSegment = TaskStateSegment::new();

lazy_static! {
    #[derive(Debug)]
    pub static ref SELECTORS: Selectors = {
        let mut gdt = GlobalDescriptorTable::empty();
        // DO NOT REORDER
        let sel_k_cs = gdt.append(Descriptor::kernel_code_segment());
        // DO NOT REORDER
        let sel_k_ss = gdt.append(
            Descriptor::kernel_data_segment().set_flags(DescriptorFlags::from_bits_retain(0b1010)),
        );
        // DO NOT REORDER
        let sel_u_ss = gdt.append(
            Descriptor::user_data_segment().set_flags(DescriptorFlags::from_bits_retain(0b1010)),
        );
        // DO NOT REORDER
        let sel_u_cs = gdt.append(Descriptor::user_code_segment());
        // DO NOT REORDER
        let sel_tss = gdt.append(
            Descriptor::tss_segment(&TSS).set_flags(DescriptorFlags::from_bits_retain(0b1010)),
        );
        Selectors {
            sel_kernel_code: sel_k_cs,
            sel_kernel_stack: sel_k_ss,
            sel_user_code: sel_u_cs,
            sel_user_stack: sel_u_ss,
            sel_tss,
            gdt,
        }
    };
}

#[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)]
pub 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(()),
            }
        }
    }
}