//! Process data structures

use core::{
    fmt,
    mem::size_of,
    pin::Pin,
    sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering},
    task::{Context, Poll},
};

use abi::{
    error::Error,
    process::{ExitCode, Signal, ThreadSpawnOptions},
};
use alloc::{
    collections::{BTreeMap, VecDeque},
    string::String,
    sync::Arc,
    vec::Vec,
};
use futures_util::Future;
use kernel_util::util::OneTimeInit;
use vfs::VnodeRef;

use crate::{
    mem::{
        phys,
        pointer::{PhysicalRef, PhysicalRefMut},
        process::ProcessAddressSpace,
        table::MapAttributes,
    },
    proc::{self, io::ProcessIo},
    sync::IrqSafeSpinlock,
    task::context::TaskContextImpl,
};

use super::{
    runtime::QueueWaker,
    sync::UserspaceMutex,
    thread::{Thread, ThreadId, ThreadState},
    TaskContext,
};

#[derive(PartialEq)]
pub enum ProcessState {
    Running,
    Terminated(ExitCode),
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
#[repr(transparent)]
pub struct ProcessId(u64);

// TLS layout (x86-64):
// |     mem_size     | uthread_size |
// |   Data    .......|   self, ???  |
//
// TLS layout (aarch64):
// | uthread_size (0x10?) |     mem_size     |
// |         ???          |  Data       .....|
#[derive(Debug)]
pub struct ProcessTlsInfo {
    pub master_copy_base: usize,
    pub layout: ProcessTlsLayout,
}

#[derive(Debug)]
pub struct ProcessTlsLayout {
    pub data_offset: usize,
    pub uthread_offset: usize,
    pub ptr_offset: usize,

    pub data_size: usize,
    pub mem_size: usize,

    pub full_size: usize,
}

#[cfg(target_arch = "aarch64")]
impl ProcessTlsLayout {
    pub fn new(align: usize, data_size: usize, mem_size: usize) -> Self {
        debug_assert!(align.is_power_of_two());
        let tls_block0_offset = (size_of::<usize>() * 2 + align - 1) & !(align - 1);

        let full_size = (tls_block0_offset + mem_size + align - 1) & !(align - 1);

        Self {
            data_offset: tls_block0_offset,
            uthread_offset: 0,
            ptr_offset: 0,

            data_size,
            mem_size,
            full_size,
        }
    }
}

#[cfg(target_arch = "x86_64")]
impl ProcessTlsLayout {
    pub fn new(align: usize, data_size: usize, mem_size: usize) -> Self {
        // 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);
        // Self-pointer
        let forward_size = size_of::<usize>();

        let full_size = back_size + forward_size;

        Self {
            data_offset: 0,
            uthread_offset: back_size,
            ptr_offset: back_size,

            data_size,
            mem_size,
            full_size,
        }
    }
}

pub struct ProcessImage {
    pub entry: usize,
    pub tls: Option<ProcessTlsInfo>,
}

struct ProcessInner {
    state: ProcessState,

    session_id: ProcessId,
    group_id: ProcessId,

    session_terminal: Option<VnodeRef>,

    threads: Vec<Arc<Thread>>,
    mutexes: BTreeMap<usize, Arc<UserspaceMutex>>,
}

pub struct Process {
    name: String,
    id: ProcessId,

    space: Arc<ProcessAddressSpace>,
    inner: IrqSafeSpinlock<ProcessInner>,
    image: Option<ProcessImage>,

    exit_waker: QueueWaker,
    pub io: IrqSafeSpinlock<ProcessIo>,
}

static PROCESSES: IrqSafeSpinlock<BTreeMap<ProcessId, Arc<Process>>> =
    IrqSafeSpinlock::new(BTreeMap::new());

impl Process {
    pub fn new_with_main<S: Into<String>>(
        name: S,
        space: Arc<ProcessAddressSpace>,
        context: TaskContext,
        image: Option<ProcessImage>,
    ) -> (Arc<Self>, Arc<Thread>) {
        let name = name.into();
        let id = ProcessId::next();

        let process = Arc::new(Self {
            name,
            id,

            image,

            space: space.clone(),
            inner: IrqSafeSpinlock::new(ProcessInner {
                state: ProcessState::Running,
                session_id: id,
                group_id: id,
                session_terminal: None,
                threads: Vec::new(),
                mutexes: BTreeMap::new(),
            }),

            exit_waker: QueueWaker::new(),
            io: IrqSafeSpinlock::new(ProcessIo::new()),
        });

        // Create "main" thread
        let thread = Thread::new_uthread(process.clone(), space, context);
        process.inner.lock().threads.push(thread.clone());

        PROCESSES.lock().insert(id, process.clone());

        (process, thread)
    }

    pub fn spawn_thread(self: &Arc<Self>, options: &ThreadSpawnOptions) -> Result<ThreadId, Error> {
        debugln!(
            "Spawn thread in {} with options: {:#x?}",
            self.id(),
            options
        );

        let tls_address = if let Some(image) = self.image.as_ref() {
            proc::elf::clone_tls(&self.space, image)?
        } else {
            0
        };

        let space = self.space.clone();
        let context = TaskContext::user(
            options.entry as _,
            options.argument as _,
            space.as_address_with_asid(),
            options.stack_top,
            tls_address,
        )?;
        let thread = Thread::new_uthread(self.clone(), space, context);
        let id = thread.id();

        self.inner.lock().threads.push(thread.clone());

        thread.enqueue_somewhere();

        Ok(id)
    }

    pub fn id(&self) -> ProcessId {
        self.id
    }

    pub fn group_id(&self) -> ProcessId {
        self.inner.lock().group_id
    }

    pub fn session_id(&self) -> ProcessId {
        self.inner.lock().session_id
    }

    pub fn name(&self) -> &str {
        self.name.as_ref()
    }

    pub fn set_group_id(&self, id: ProcessId) {
        self.inner.lock().group_id = id;
    }

    pub fn set_session_id(&self, id: ProcessId) {
        self.inner.lock().session_id = id;
    }

    // Resources
    pub fn session_terminal(&self) -> Option<VnodeRef> {
        self.inner.lock().session_terminal.clone()
    }

    pub fn set_session_terminal(&self, node: VnodeRef) {
        self.inner.lock().session_terminal.replace(node);
    }

    pub fn clear_session_terminal(&self) -> Option<VnodeRef> {
        self.inner.lock().session_terminal.take()
    }

    pub fn inherit(&self, parent: &Process) -> Result<(), Error> {
        let mut our_inner = self.inner.lock();
        let their_inner = parent.inner.lock();

        our_inner.session_id = their_inner.session_id;
        our_inner.group_id = their_inner.group_id;
        our_inner.session_terminal = their_inner.session_terminal.clone();

        Ok(())
    }

    // State
    pub fn get_exit_status(&self) -> Option<ExitCode> {
        match self.inner.lock().state {
            ProcessState::Running => None,
            ProcessState::Terminated(x) => Some(x),
        }
    }

    pub fn wait_for_exit(process: Arc<Process>) -> impl Future<Output = ExitCode> {
        struct ProcessExitFuture {
            process: Arc<Process>,
        }

        impl Future for ProcessExitFuture {
            type Output = ExitCode;

            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
                let process = &self.process;

                process.exit_waker.register(cx.waker());

                if let Some(exit_status) = process.get_exit_status() {
                    process.exit_waker.remove(cx.waker());
                    Poll::Ready(exit_status)
                } else {
                    Poll::Pending
                }
            }
        }

        ProcessExitFuture { process }
    }

    pub fn handle_thread_exit(&self, thread: ThreadId, code: ExitCode) {
        debugln!("Thread {} of process {}: {:?}", thread, self.id, code);
        let mut inner = self.inner.lock();

        // TODO make this cleaner
        let old_len = inner.threads.len();
        inner.threads.retain(|t| t.id() != thread);
        assert_ne!(inner.threads.len(), old_len);

        let last_thread = inner.threads.is_empty();

        if last_thread {
            debugln!("Last thread of {} exited", self.id);

            inner.state = ProcessState::Terminated(code);

            self.io.lock().handle_exit();

            drop(inner);

            self.exit_waker.wake_all();
        }
    }

    /// Raises a signal for the specified process
    pub fn raise_signal(self: &Arc<Self>, signal: Signal) {
        let thread = self.inner.lock().threads[0].clone();
        thread.raise_signal(signal);
    }

    /// Raises a signal for the specified process group
    pub fn signal_group(group_id: ProcessId, signal: Signal) {
        let processes = PROCESSES.lock();
        for (_, proc) in processes.iter() {
            let inner = proc.inner.lock();
            if !matches!(inner.state, ProcessState::Terminated(_)) && inner.group_id == group_id {
                debugln!("Deliver group signal to {}: {:?}", proc.id(), signal);
                drop(inner);
                proc.raise_signal(signal);
            }
        }
    }

    pub fn get_or_insert_mutex(&self, address: usize) -> Arc<UserspaceMutex> {
        let mut inner = self.inner.lock();
        inner
            .mutexes
            .entry(address)
            .or_insert_with(|| Arc::new(UserspaceMutex::new(address)))
            .clone()
    }

    // Process list
    pub fn get(id: ProcessId) -> Option<Arc<Self>> {
        PROCESSES.lock().get(&id).cloned()
    }

    pub async fn terminate_others(&self, except: ThreadId) {
        let mut inner = self.inner.lock();

        for thread in inner.threads.iter() {
            if thread.id() == except {
                continue;
            }

            infoln!("Terminate thread {}", thread.id());
            thread.terminate().await;
        }

        inner.threads.retain(|t| t.id() == except);
    }
}

impl fmt::Display for ProcessId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        use fmt::Write;

        write!(f, "<Process {}>", self.0)
    }
}

// XXX TODO Remove this
impl From<ProcessId> for u32 {
    fn from(value: ProcessId) -> Self {
        value.0 as _
    }
}

impl From<u32> for ProcessId {
    fn from(value: u32) -> Self {
        Self(value as _)
    }
}

impl ProcessId {
    pub fn next() -> Self {
        static COUNTER: AtomicU64 = AtomicU64::new(1);
        let id = COUNTER.fetch_add(1, Ordering::SeqCst);
        Self(id)
    }
}

//     /// Returns the state of the process.
//     ///
//     /// # Note
//     ///
//     /// Maybe I should remove this and make ALL state changes atomic.
//     pub fn state(&self) -> ProcessState {
//         self.state.load(Ordering::Acquire)
//     }
//
//     /// Atomically updates the state of the process and returns the previous one.
//     pub fn set_state(&self, state: ProcessState) -> ProcessState {
//         self.state.swap(state, Ordering::SeqCst)
//     }
//
//     /// Marks the task as running on the specified CPU.
//     ///
//     /// # Safety
//     ///
//     /// Only meant to be called from scheduler routines.
//     pub unsafe fn set_running(&self, cpu: u32) {
//         self.cpu_id.store(cpu, Ordering::Release);
//         self.state.store(ProcessState::Running, Ordering::Release);
//     }
//
//     /// Returns the address space of the task
//     pub fn address_space(&self) -> &ProcessAddressSpace {
//         self.space.as_ref().unwrap()
//     }
//
//     /// Returns the address space of the task, if one is set
//     pub fn get_address_space(&self) -> Option<&ProcessAddressSpace> {
//         self.space.as_ref()
//     }
//
//     /// Replaces the task's session terminal device with another one
//     pub fn set_session_terminal(&self, terminal: VnodeRef) {
//         self.inner.lock().session_terminal.replace(terminal);
//     }
//
//     /// Removes the task's current terminal
//     pub fn clear_session_terminal(&self) -> Option<VnodeRef> {
//     }
//
//     /// Returns the current terminal of the task
//     pub fn session_terminal(&self) -> Option<VnodeRef> {
//         self.inner.lock().session_terminal.clone()
//     }
//
//     /// Sets the session ID of the task
//     pub fn set_session_id(&self, sid: ProcessId) {
//         self.inner.lock().session_id = sid;
//     }
//
//     /// Sets the process group ID of the task
//     pub fn set_group_id(&self, mut gid: ProcessId) {
//         if gid == 0 {
//             gid = self.id();
//         }
//         self.inner.lock().group_id = gid;
//     }
//
//     /// Returns the process group ID of the task
//     pub fn group_id(&self) -> ProcessId {
//         self.inner.lock().group_id
//     }
//
//     /// Returns the CPU number this task in running on (or the last one)
//     pub fn cpu_id(&self) -> u32 {
//         self.cpu_id.load(Ordering::Acquire)
//     }
//
//     /// Selects a suitable CPU queue and submits the process for execution.
//     ///
//     /// # Panics
//     ///
//     /// Currently, the code will panic if the process is queued/executing on any queue.
//     pub fn enqueue_somewhere(self: Arc<Self>) -> usize {
//         // Doesn't have to be precise, so even if something changes, we can still be rebalanced
//         // to another CPU
//         let (index, queue) = CpuQueue::least_loaded().unwrap();
//
//         self.enqueue_to(queue);
//
//         index
//     }
//
//     /// Submits the process to a specific queue.
//     ///
//     /// # Panics
//     ///
//     /// Currently, the code will panic if the process is queued/executing on any queue.
//     pub fn enqueue_to(self: Arc<Self>, queue: &'static CpuQueue) {
//         let _irq = IrqGuard::acquire();
//
//         {
//             let mut inner = self.inner.lock();
//             let old_queue = inner.queue.replace(queue);
//             if old_queue.is_some() {
//                 // Already in some queue
//                 return;
//             }
//         }
//         match self.state.compare_exchange(
//             ProcessState::Suspended,
//             ProcessState::Ready,
//             Ordering::SeqCst,
//             Ordering::Relaxed,
//         ) {
//             Err(ProcessState::Terminated) => {
//                 // Process might've been killed while `await`ing in a `block!`
//                 debugln!(
//                     "Process {} {:?} already terminated, dropping",
//                     self.id(),
//                     self.name()
//                 );
//             }
//             Err(state) => {
//                 todo!("Unexpected process state when enqueueing: {:?}", state)
//             }
//             Ok(_) => unsafe {
//                 queue.enqueue(self);
//             },
//         }
//     }
//
//     /// Returns an exit code if the process exited, [None] if it didn't
//     pub fn get_exit_status(&self) -> Option<ExitCode> {
//         if self.state() == ProcessState::Terminated {
//             Some(ExitCode::from(self.inner.lock().exit_status))
//         } else {
//             None
//         }
//     }
//
//     /// Returns the [Process] currently executing on local CPU, None if idling.
//     pub fn get_current() -> Option<CurrentProcess> {
//         let queue = Cpu::local().queue();
//         queue.current_process()
//     }
//
//     /// Returns a process by its ID
//     pub fn get(pid: ProcessId) -> Option<Arc<Self>> {
//         PROCESSES.lock().get(pid).cloned()
//     }
//
//     /// Wraps [Process::get_current()] for cases when the caller is absolutely sure there is a
//     /// running process (e.g. the call itself comes from a process).
//     pub fn current() -> CurrentProcess {
//         Self::get_current().unwrap()
//     }
//
//     /// Handles the cleanup of an exited process
//     pub fn handle_exit(&self) {
//     }
//
//     /// Inherits the data from a parent process. Meant to be called from SpawnProcess handler.
//     pub fn inherit(&self, parent: &Arc<Process>) -> Result<(), Error> {
//     }
//
//
//     pub fn wait_for_exit(process: Arc<Self>) -> impl Future<Output = ExitCode> {
//     }
// }
//
// impl ArcWake for Process {
//     fn wake_by_ref(arc_self: &Arc<Self>) {
//         arc_self.clone().enqueue_somewhere();
//     }
// }
//
// impl Drop for Process {
//     fn drop(&mut self) {
//         infoln!("Drop process!");
//     }
// }
//
// impl CurrentProcess {
//     /// Wraps a process in this structure.
//     ///
//     /// # Safety
//     ///
//     /// Only meant to be called from [Process::current] or [CpuQueue::current_process].
//     pub unsafe fn new(inner: Arc<Process>, guard: IrqGuard) -> Self {
//         Self(inner, guard)
//     }
//
//     /// Configures signal entry information for the process
//     pub fn set_signal_entry(&self, entry: usize, stack: usize) {
//         let mut inner = self.inner.lock();
//         inner.signal_entry.replace(SignalEntry { entry, stack });
//     }
//
//     pub fn suspend(&self) -> Result<(), Error> {
//         self.dequeue(ProcessState::Suspended);
//
//         let inner = self.inner.lock();
//         if !inner.signal_stack.is_empty() {
//             return Err(Error::Interrupted);
//         }
//
//         Ok(())
//     }
//
//     /// Terminate the current process
//     pub fn exit(&self, status: ExitCode) {
//         self.inner.lock().exit_status = status.into();
//         debugln!("Process {} exited with code {:?}", self.id(), status);
//
//         self.handle_exit();
//         self.dequeue(ProcessState::Terminated);
//     }
//
// }
//
// impl Deref for CurrentProcess {
//     type Target = Arc<Process>;
//
//     fn deref(&self) -> &Self::Target {
//         &self.0
//     }
// }