use core::pin::pin;

use alloc::sync::Arc;
use alloc::vec::Vec;
use bitflags::bitflags;
use eonix_sync::{LazyLock, Mutex, MutexGuard, WaitList};
use intrusive_collections::{intrusive_adapter, KeyAdapter, RBTree, RBTreeAtomicLink};

use crate::{
    kernel::{
        constants::{EAGAIN, EINVAL},
        user::UserPointer,
    },
    prelude::KResult,
};

#[derive(PartialEq, Debug, Clone, Copy)]
#[repr(u32)]
#[expect(non_camel_case_types)]
pub enum FutexOp {
    FUTEX_WAIT = 0,
    FUTEX_WAKE = 1,
    FUTEX_FD = 2,
    FUTEX_REQUEUE = 3,
    FUTEX_CMP_REQUEUE = 4,
    FUTEX_WAKE_OP = 5,
    FUTEX_LOCK_PI = 6,
    FUTEX_UNLOCK_PI = 7,
    FUTEX_TRYLOCK_PI = 8,
    FUTEX_WAIT_BITSET = 9,
    FUTEX_WAKE_BITSET = 10,
}

impl TryFrom<u32> for FutexOp {
    type Error = u32;

    fn try_from(value: u32) -> Result<Self, Self::Error> {
        match value {
            0 => Ok(FutexOp::FUTEX_WAIT),
            1 => Ok(FutexOp::FUTEX_WAKE),
            2 => Ok(FutexOp::FUTEX_FD),
            3 => Ok(FutexOp::FUTEX_REQUEUE),
            4 => Ok(FutexOp::FUTEX_CMP_REQUEUE),
            5 => Ok(FutexOp::FUTEX_WAKE_OP),
            6 => Ok(FutexOp::FUTEX_LOCK_PI),
            7 => Ok(FutexOp::FUTEX_UNLOCK_PI),
            8 => Ok(FutexOp::FUTEX_TRYLOCK_PI),
            9 => Ok(FutexOp::FUTEX_WAIT_BITSET),
            10 => Ok(FutexOp::FUTEX_WAKE_BITSET),
            _ => Err(EINVAL),
        }
    }
}

bitflags! {
    pub struct FutexFlags : u32 {
        const FUTEX_PRIVATE         = 128;
        const FUTEX_CLOCK_REALTIME  = 256;
    }
}

const FUTEX_OP_MASK: u32 = 0x0000_000F;
const FUTEX_FLAGS_MASK: u32 = 0xFFFF_FFF0;

pub fn parse_futexop(bits: u32) -> KResult<(FutexOp, FutexFlags)> {
    let op = FutexOp::try_from(bits & FUTEX_OP_MASK)?;

    let flags = {
        let flags_bits = bits & FUTEX_FLAGS_MASK;
        FutexFlags::from_bits(flags_bits).ok_or(EINVAL)
    }?;

    Ok((op, flags))
}

struct FutexTable {
    futex_buckets: Vec<Mutex<FutexBucket>>,
}

struct FutexBucket {
    futex_items: RBTree<FutexItemAdapter>,
}

intrusive_adapter!(
    FutexItemAdapter = Arc<FutexItem>: FutexItem { link: RBTreeAtomicLink }
);

impl<'a> KeyAdapter<'a> for FutexItemAdapter {
    type Key = FutexKey;
    fn get_key(&self, item: &'a FutexItem) -> Self::Key {
        item.key
    }
}

struct FutexItem {
    link: RBTreeAtomicLink,
    key: FutexKey,
    // // A list of waiters that are waiting on this futex.
    wait_list: WaitList,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
struct FutexKey {
    addr: usize,
    pid: Option<u32>,
}

impl FutexKey {
    fn new(addr: usize, pid: Option<u32>) -> Self {
        FutexKey { addr, pid }
    }

    fn hash(&self) -> usize {
        (self.addr >> 2) + self.pid.unwrap_or(0) as usize
    }
}

static FUTEX_TABLE: LazyLock<FutexTable> = LazyLock::new(FutexTable::new);

const FUTEX_BUCKETS_NUM: usize = 100;

impl FutexTable {
    fn new() -> Self {
        let futex_buckets = (0..FUTEX_BUCKETS_NUM)
            .map(|_| Mutex::new(FutexBucket::new()))
            .collect();
        FutexTable { futex_buckets }
    }

    fn get_bucket(&self, key: &FutexKey) -> (usize, &Mutex<FutexBucket>) {
        let idx = key.hash() % self.futex_buckets.len();
        (idx, &self.futex_buckets[idx])
    }
}

impl FutexBucket {
    fn new() -> Self {
        FutexBucket {
            futex_items: RBTree::new(FutexItemAdapter::new()),
        }
    }

    fn find(&mut self, key: FutexKey) -> Option<Arc<FutexItem>> {
        let cursor = self.futex_items.find(&key);
        cursor.clone_pointer()
    }

    fn find_or_insert(&mut self, key: FutexKey) -> Arc<FutexItem> {
        if let Some(item) = self.find(key) {
            return item;
        }

        let item = Arc::new(FutexItem {
            link: RBTreeAtomicLink::new(),
            key,
            wait_list: WaitList::new(),
        });

        self.futex_items.insert(item.clone());

        item
    }
}

pub async fn futex_wait(
    uaddr: usize,
    pid: Option<u32>,
    expected_val: u32,
    _timeout: Option<u32>,
) -> KResult<()> {
    let futex_key = FutexKey::new(uaddr, pid);

    let futex_item = {
        let (_, futex_bucket_ref) = FUTEX_TABLE.get_bucket(&futex_key);
        let mut futex_bucket = futex_bucket_ref.lock().await;

        let val = UserPointer::new(uaddr as *const u32)?.read()?;

        if val != expected_val {
            return Err(EAGAIN);
        }

        futex_bucket.find_or_insert(futex_key)
    };

    let mut wait = pin!(futex_item.wait_list.prepare_to_wait());
    wait.as_mut().add_to_wait_list();
    wait.await;
    Ok(())
}

pub async fn futex_wake(uaddr: usize, pid: Option<u32>, max_wake_count: u32) -> KResult<usize> {
    let futex_key = FutexKey::new(uaddr, pid);

    let (_, futex_bucket_ref) = FUTEX_TABLE.get_bucket(&futex_key);
    let mut futex_bucket = futex_bucket_ref.lock().await;

    if let Some(futex_item) = futex_bucket.find(futex_key) {
        let mut count = 0;
        loop {
            let not_empty = futex_item.wait_list.notify_one();
            if not_empty == false || count == max_wake_count {
                break;
            }
            count += 1;
        }
        return Ok(count as usize);
    }

    Ok(0)
}

// should make sure that two keys in different buckets
// make the lock order to avoid deadlocks
async fn double_lock_bucket(
    futex_key0: FutexKey,
    futex_key1: FutexKey,
) -> (
    MutexGuard<'static, FutexBucket>,
    MutexGuard<'static, FutexBucket>,
) {
    let (bucket_idx0, bucket_ref0) = FUTEX_TABLE.get_bucket(&futex_key0);
    let (bucket_idx1, bucket_ref1) = FUTEX_TABLE.get_bucket(&futex_key1);

    if bucket_idx0 < bucket_idx1 {
        let bucket0 = bucket_ref0.lock().await;
        let bucket1 = bucket_ref1.lock().await;
        (bucket0, bucket1)
    } else {
        let bucket1 = bucket_ref1.lock().await;
        let bucket0 = bucket_ref0.lock().await;
        (bucket0, bucket1)
    }
}

async fn futex_requeue(
    uaddr: usize,
    pid: Option<u32>,
    wake_count: u32,
    requeue_uaddr: usize,
    requeue_count: u32,
) -> KResult<usize> {
    let futex_key = FutexKey::new(uaddr, pid);
    let futex_requeue_key = FutexKey::new(requeue_uaddr, pid);

    let (bucket_idx0, bucket_ref0) = FUTEX_TABLE.get_bucket(&futex_key);
    let (bucket_idx1, bucket_ref1) = FUTEX_TABLE.get_bucket(&futex_requeue_key);

    if bucket_idx0 == bucket_idx1 {
        // If the keys are the same, we can just wake up the waiters.
        return futex_wake(uaddr, pid, wake_count).await;
    }

    let (futex_bucket, futex_requeue_bucket) =
        double_lock_bucket(futex_key, futex_requeue_key).await;

    todo!()
}
