use core::{cell::RefCell, ffi::c_void, ptr::null_mut};

use alloc::{boxed::Box, collections::BTreeMap, vec::Vec};
use yggdrasil_rt::{
    process::{
        thread::{self as rt, ThreadCreateInfo},
        thread_local, ProgramArgumentInner,
    },
    sync::mutex::Mutex,
};

use crate::{
    error::{EResult, OptionExt},
    headers::{
        errno::Errno,
        sys_types::{pthread_attr_t, pthread_t},
    },
};

pub mod tls;

pub struct ThreadDestructor {
    pub destructor: extern "C" fn(*mut c_void),
    pub argument: *mut c_void,
}

static THREADS: Mutex<BTreeMap<pthread_t, Thread>> = Mutex::new(BTreeMap::new());
#[thread_local]
static THREAD_LOCAL_DTORS: RefCell<Vec<ThreadDestructor>> = RefCell::new(Vec::new());

pub const DEFAULT_STACK_SIZE: usize = 4096 * 16;
pub const DEFAULT_SIGNAL_STACK_SIZE: usize = 4096 * 8;

pub const MIN_STACK_SIZE: usize = 4096 * 4;

pub struct Thread {
    handle: rt::ThreadHandle<*mut c_void>,
}

impl Thread {
    pub fn spawn(
        attr: Option<&pthread_attr_t>,
        entry: extern "C" fn(*mut c_void) -> *mut c_void,
        argument: *mut c_void,
    ) -> EResult<pthread_t> {
        let stack = if let Some(attr) = attr {
            assert!(attr.stack_size >= MIN_STACK_SIZE);
            match attr.stack {
                Some(stack) => rt::ThreadStack::Raw(usize::from(stack.addr()) + attr.stack_size),
                None => rt::ThreadStack::Allocate(attr.stack_size),
            }
        } else {
            rt::ThreadStack::Allocate(DEFAULT_STACK_SIZE)
        };
        // TODO honor attr stack
        let info = ThreadCreateInfo {
            signal_stack: rt::ThreadSignalStack::Allocate(DEFAULT_SIGNAL_STACK_SIZE),
            stack,
            #[allow(clippy::redundant_closure)]
            entry: rt::ThreadFunction::Closure(Box::new(move |arg| Self::thread_main(entry, arg))),
            tls_image: tls::TLS_IMAGE.get().as_ref(),
        };
        let handle = rt::Thread::spawn(info, argument, true)?;
        let id = handle.id();
        THREADS.lock().insert(id, Thread { handle });
        EResult::Ok(id)
    }

    pub fn join(id: pthread_t) -> EResult<*mut c_void> {
        let thread = THREADS.lock().remove(&id).e_ok_or(Errno::EINVAL)?;
        // TODO handle EINTR during join
        let result = thread.handle.join_uninterruptible().unwrap_or(null_mut());
        EResult::Ok(result)
    }

    pub fn this() -> EResult<pthread_t> {
        EResult::Ok(unsafe { rt::Thread::<*mut c_void>::current().id() })
    }

    fn thread_main(
        entry: extern "C" fn(*mut c_void) -> *mut c_void,
        arg: *mut c_void,
    ) -> *mut c_void {
        let ret = entry(arg);
        unsafe { Self::call_thread_dtors() };
        ret
    }

    pub fn at_thread_exit(
        destructor: extern "C" fn(*mut c_void),
        argument: *mut c_void,
        _dso_handle: *mut c_void,
    ) {
        THREAD_LOCAL_DTORS.borrow_mut().push(ThreadDestructor {
            destructor,
            argument,
        });
    }

    pub unsafe fn call_thread_dtors() {
        for dtor in THREAD_LOCAL_DTORS.borrow().iter() {
            (dtor.destructor)(dtor.argument);
        }
    }
}

pub fn init_main_thread(arg: &ProgramArgumentInner) {
    // Will create a TLS for the main thread (if not done already).
    // Usually, a dynamic loader will do this for us, but this still needs to be
    // done for statically-linked programs if the kernel transfers control directly to
    // the program.
    let tls_image =
        thread_local::init_tls_from_auxv(arg.auxv(), false).expect("Could not initialize TLS");

    // Store the TLS image, it'll be needed when creating new threads
    tls::TLS_IMAGE.init_with(|| tls_image);

    unsafe { rt::track_main::<*mut c_void>() };
}
