//! Process data structures

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

use abi::{
    error::Error,
    process::{ExitCode, Signal, ThreadSpawnOptions},
};
use alloc::{collections::BTreeMap, string::String, sync::Arc, vec::Vec};
use futures_util::Future;
use kernel_util::{runtime::QueueWaker, sync::IrqSafeSpinlock};
use vfs::NodeRef;

use crate::{
    mem::process::ProcessAddressSpace,
    proc::{self, io::ProcessIo},
    task::context::TaskContextImpl,
};

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

/// Represents a process state
#[derive(PartialEq)]
pub enum ProcessState {
    /// Process is running, meaning it still has at least one thread alive
    Running,
    /// Process has finished with some [ExitCode]
    Terminated(ExitCode),
}

/// Unique number assigned to each [Process]
#[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       .....|

/// Describes Thread-Local Storage of a process
#[derive(Debug)]
pub struct ProcessTlsInfo {
    /// Location of the TLS master copy within the process's memory
    pub master_copy_base: usize,
    /// Layout of the TLS
    pub layout: ProcessTlsLayout,
}

/// Describes TLS layout for a program image
#[derive(Debug)]
pub struct ProcessTlsLayout {
    /// Data offset from the TLS base
    pub data_offset: usize,
    /// struct uthread offset from the TLS base
    pub uthread_offset: usize,
    /// Pointer offset from the TLS base. The pointer is passed to the userspace
    pub ptr_offset: usize,

    /// Data size of the TLS segment
    pub data_size: usize,
    /// Memory size of the TLS segment (mem_size >= data_size)
    pub mem_size: usize,

    /// Overall allocation size of the TLS data
    pub full_size: usize,
}

#[cfg(target_arch = "aarch64")]
impl ProcessTlsLayout {
    /// Constructs a new thread-local storage layout info struct
    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 {
    /// Constructs a new thread-local storage layout info struct
    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,
        }
    }
}

/// Describes information about a program's image in memory
pub struct ProcessImage {
    /// Entry point address
    pub entry: usize,
    /// Thread-local storage information
    pub tls: Option<ProcessTlsInfo>,
}

struct ProcessInner {
    state: ProcessState,

    session_id: ProcessId,
    group_id: ProcessId,

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

/// Describes a process within the system
pub struct Process {
    name: String,
    id: ProcessId,

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

    exit_waker: QueueWaker,
    /// Process I/O information
    pub io: IrqSafeSpinlock<ProcessIo>,
}

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

impl Process {
    /// Creates a new process with given main thread
    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)
    }

    /// Spawns a new child thread within the process
    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();

        Ok(id)
    }

    /// Returns the [ProcessId] of this process
    pub fn id(&self) -> ProcessId {
        self.id
    }

    /// Returns the process group ID of the process
    pub fn group_id(&self) -> ProcessId {
        self.inner.lock().group_id
    }

    /// Returns the process session ID of the process
    pub fn session_id(&self) -> ProcessId {
        self.inner.lock().session_id
    }

    /// Returns the process name
    pub fn name(&self) -> &str {
        self.name.as_ref()
    }

    /// Changes the process's group ID
    pub fn set_group_id(&self, id: ProcessId) {
        self.inner.lock().group_id = id;
    }

    /// Changes the process's session ID
    pub fn set_session_id(&self, id: ProcessId) {
        self.inner.lock().session_id = id;
    }

    // Resources

    /// Returns the current session terminal of the process, if set
    pub fn session_terminal(&self) -> Option<NodeRef> {
        self.inner.lock().session_terminal.clone()
    }

    /// Changes the current session terminal of the process
    pub fn set_session_terminal(&self, node: NodeRef) {
        self.inner.lock().session_terminal.replace(node);
    }

    /// Resets the current session terminal of the process
    pub fn clear_session_terminal(&self) -> Option<NodeRef> {
        self.inner.lock().session_terminal.take()
    }

    /// Inherits the process information from the `parent`
    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

    /// Returns the [ExitCode] of the process, if it has exited
    pub fn get_exit_status(&self) -> Option<ExitCode> {
        match self.inner.lock().state {
            ProcessState::Running => None,
            ProcessState::Terminated(x) => Some(x),
        }
    }

    /// Suspends the task until the process exits
    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 }
    }

    /// Handles exit of a single child thread
    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);

            // XXX
            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);
            }
        }
    }

    /// Returns a [UserspaceMutex] associated with the `address`. If one does not exist, will
    /// create it.
    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

    /// Returns the process with given [ProcessId], if it exists
    pub fn get(id: ProcessId) -> Option<Arc<Self>> {
        PROCESSES.lock().get(&id).cloned()
    }

    /// Terminates all children of the process, `except` one
    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 {
        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 {
    /// Generates a new [ProcessId]
    pub fn next() -> Self {
        static COUNTER: AtomicU64 = AtomicU64::new(1);
        let id = COUNTER.fetch_add(1, Ordering::SeqCst);
        Self(id)
    }
}