use core::{
    fmt,
    sync::atomic::{AtomicU32, AtomicU64, Ordering},
};

use atomic_enum::atomic_enum;
use yggdrasil_abi::process::ProcessId;

pub trait AllocateProcessId {
    fn new() -> Self;
}

/// Represents the states a thread can be at some point in time
#[atomic_enum]
#[derive(PartialEq)]
pub enum ThreadState {
    /// Thread is ready for execution and is present in some CPU's queue
    Ready,
    /// Thread is currently being executed by some CPU
    Running,
    /// Thread is present in a global list, but is not queued for execution until it is resumed
    Suspended,
    /// Thread is terminated and waits to be reaped
    Terminated,
}

/// Mask describing CPUs a thread is allowed to be scheduled to
#[derive(Debug)]
#[repr(transparent)]
pub struct ThreadAffinity(AtomicU64);

/// Unique number describing a single kernel or userspace thread
#[derive(Debug, PartialEq, Eq, Clone, Copy, Ord, PartialOrd)]
pub enum ThreadId {
    /// Describes a kernel-space thread
    Kernel(u64),
    /// Describes an user-space thread
    User(u64),
}

/// Describes a TLS master copy
#[derive(Debug, Clone)]
pub struct TlsImage {
    /// Where the kernel placed the segment in the process' memory.
    /// NOTE: is zero if there's no actual data in the master copy.
    pub master_copy_base: usize,
    /// Segment alignment
    pub align: usize,
    /// Part of the segment occupied by actual data
    pub data_size: usize,
    /// Full size of the segment. `full_size - data_size` is zeroed on load.
    pub full_size: usize,
}

// // TLS layout (x86-64):
// // |     mem_size     | uthread_size |
// // |   Data    .......|   self, ???  |
// //
// // TLS layout (aarch64):
// // | uthread_size (0x10?) |     mem_size     |
// // |         ???          |  Data       .....|
// //
// // TLS layout (i686):
// // | self | ... |
//
// /// Describes Thread-Local Storage of a process
// #[derive(Clone, Debug)]
// pub struct ProcessTlsInfo {
//     /// Master copy of the TLS
//     pub master_copy: Option<Arc<PageBox<[u8]>>>,
//     /// Layout of the TLS
//     pub layout: ProcessTlsLayout,
// }
//
// /// Describes TLS layout for a program image
// #[derive(Clone, Debug)]
// pub struct ProcessTlsLayout {
//     /// Pointer offset from the TLS base. The pointer is passed to the userspace
//     pub ptr_offset: usize,
//     pub data_offset: usize,
//     pub data_size: usize,
//     pub mem_size: usize,
//     pub slot_offset: usize,
//     /// Overall allocation size of the TLS data
//     pub full_size: usize,
// }

impl ThreadAffinity {
    /// Mask value for a thread to be scheduled onto any CPU
    pub const ANY_CPU: u64 = u64::MAX;

    /// Constructs an affinity mask allowing a thread to be scheduled onto any CPU
    pub const fn any_cpu() -> Self {
        Self(AtomicU64::new(Self::ANY_CPU))
    }

    /// Constructs an affinity mask targeting a single CPU
    pub const fn only_cpu(index: usize) -> Self {
        Self(AtomicU64::new(1 << index))
    }

    /// Returns the current value of the thread's CPU affinity mask
    #[inline]
    pub fn get(&self) -> u64 {
        self.0.load(Ordering::Relaxed)
    }

    /// Updates the thread's CPU affinity mask
    #[inline]
    pub fn set(&self, value: u64) {
        debug_assert_ne!(value, 0);
        self.0.store(value, Ordering::Relaxed);
    }
}

impl ThreadId {
    /// Generates a new kernel-space thread ID
    pub fn next_kernel() -> Self {
        static COUNT: AtomicU64 = AtomicU64::new(1);
        let id = COUNT.fetch_add(1, Ordering::SeqCst);
        Self::Kernel(id)
    }

    /// Generates a new user-space thread ID
    pub fn next_user() -> Self {
        static COUNT: AtomicU64 = AtomicU64::new(1);
        let id = COUNT.fetch_add(1, Ordering::SeqCst);
        Self::User(id)
    }

    /// Returns the number of the userspace thread represented by this ID.
    ///
    /// # Panics
    ///
    /// Will panic if this [ThreadId] does not represent a user-space thread.
    pub fn as_user(&self) -> u64 {
        match self {
            Self::Kernel(_) => panic!(),
            &Self::User(id) => id,
        }
    }

    /// Returns `true` if the [ThreadId] represents a user-space thread
    pub fn is_user(&self) -> bool {
        matches!(self, ThreadId::User(_))
    }
}

impl fmt::Display for ThreadId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Kernel(id) => write!(f, "#[{id}]"),
            Self::User(id) => write!(f, "#{id}"),
        }
    }
}

impl AllocateProcessId for ProcessId {
    fn new() -> Self {
        static COUNTER: AtomicU32 = AtomicU32::new(1);
        let id = COUNTER.fetch_add(1, Ordering::SeqCst);
        unsafe { Self::from_raw(id) }
    }
}

// #[cfg(any(target_arch = "aarch64", rust_analyzer))]
// impl ProcessTlsLayout {
//     /// Constructs a new thread-local storage layout info struct
//     pub fn new(align: usize, data_size: usize, mem_size: usize) -> Self {
//         use core::mem::size_of;
//
//         debug_assert!(align.is_power_of_two());
//         let tls_block0_offset = (size_of::<usize>() * 2 + align - 1) & !(align - 1);
//         // Allocate two words in the back for extra storage
//         let back_size = size_of::<usize>() * 2;
//
//         let full_size = back_size + (tls_block0_offset + mem_size + align - 1) & !(align - 1);
//
//         Self {
//             data_offset: tls_block0_offset + back_size,
//             ptr_offset: back_size,
//             slot_offset: 0,
//
//             data_size,
//             mem_size,
//             full_size,
//         }
//     }
// }
//
// #[cfg(any(target_arch = "x86_64", rust_analyzer))]
// impl ProcessTlsLayout {
//     /// Constructs a new thread-local storage layout info struct
//     pub fn new(align: usize, data_size: usize, mem_size: usize) -> Self {
//         use core::mem::size_of;
//
//         // The static TLS blocks are placed below TP
//         // TP points to the TCB
//         debug_assert!(align.is_power_of_two());
//         let back_size = (mem_size + align - 1) & !(align - 1);
//         // 0: self-pointer
//         // 1: -
//         // 2: segment pointer
//         // 3: segment size
//         let forward_size = size_of::<usize>() * 4;
//
//         let full_size = back_size + forward_size;
//
//         Self {
//             data_offset: 0,
//             ptr_offset: back_size,
//             slot_offset: back_size + 2 * size_of::<usize>(),
//
//             data_size,
//             mem_size,
//             full_size,
//         }
//     }
// }
//
// #[cfg(any(target_arch = "x86", rust_analyzer))]
// impl ProcessTlsLayout {
//     pub fn new(align: usize, data_size: usize, mem_size: usize) -> Self {
//         use core::mem::size_of;
//
//         debug_assert!(align.is_power_of_two());
//         let back_size = (mem_size + align - 1) & !(align - 1);
//         let forward_size = size_of::<usize>() * 4;
//         let full_size = back_size + forward_size;
//
//         Self {
//             data_offset: 0,
//             ptr_offset: back_size,
//             slot_offset: back_size + 2 * size_of::<usize>(),
//
//             data_size,
//             mem_size,
//             full_size,
//         }
//     }
// }
