//! task 相关操作

use alloc::{format, string::String, sync::Arc, vec::Vec};
use core::sync::atomic::{AtomicUsize, Ordering};

use moa_bsp_define::smp::NR_CPUS;
use moa_cpumask::CpuId;
use moa_task_define::TaskBase;

use crate::{
    processor::{nr_cpus, this_processor_id},
    sched::task::{
        pid::{Pid, pid_init},
        stack::THREAD_STACK_SIZE,
    },
    space::{
        addr::{Vaddr, init_stack_base},
        kalloc::{GfpFlags, kfree, kmalloc},
        mm::MmStruct,
    },
};

mod pid;
pub mod stack;

/// 任务结构体
#[repr(C)]
#[allow(clippy::struct_field_names)]
pub struct Task {
    task_base: TaskBase,
    name: String,
    cpu: AtomicUsize,
    mm: Option<Arc<MmStruct>>,
    pid: Pid,
}

impl Task {
    const fn new() -> Self {
        Task {
            task_base: TaskBase::new(),
            name: String::new(),
            cpu: AtomicUsize::new(0),
            mm: None,
            pid: Pid::new(),
        }
    }

    fn init(name: &str, cpu: usize, mm: Option<Arc<MmStruct>>, stack: usize) -> Self {
        let tsk = Self {
            task_base: TaskBase::new(),
            name: String::from(name),
            cpu: AtomicUsize::new(cpu),
            mm,
            pid: Pid::create(),
        };
        tsk.task_base.set_stack(stack);
        tsk
    }

    /// 创建一个任务
    ///
    /// # Panics
    /// 内存不足将会 panic.
    pub fn create(name: &str, cpu: usize, mm: Option<Arc<MmStruct>>) -> Self {
        let stack_base = kmalloc(THREAD_STACK_SIZE, GfpFlags::Clean).unwrap().to_value();
        Task::init(name, cpu, mm, stack_base)
    }

    /// 任务名
    pub fn name(&self) -> &str {
        &self.name
    }

    /// 任务所在处理器
    pub fn cpuid(&self) -> CpuId {
        CpuId::from_raw(self.cpu.load(Ordering::Relaxed))
    }

    /// 任务地址空间
    pub fn mm(&self) -> Option<Arc<MmStruct>> {
        self.mm.as_ref().map(Arc::clone)
    }

    /// 设置任务抢占计数
    #[inline(always)]
    pub fn set_preempt(&self, val: u32) {
        self.task_base.set_preempt(val);
    }

    /// 获取任务抢占计数
    #[inline(always)]
    pub fn preempt(&self) -> u32 {
        self.task_base.preempt()
    }

    /// 增加任务抢占计数
    #[inline(always)]
    pub fn preempt_add(&self, val: u32) {
        self.task_base.preempt_add(val);
    }

    /// 减少任务抢占计数
    #[inline(always)]
    pub fn preempt_sub(&self, val: u32) {
        self.task_base.preempt_sub(val);
    }

    /// 增加一次任务抢占计数
    #[inline(always)]
    pub fn preempt_inc(&self) {
        self.preempt_add(1);
    }

    /// 减少一次任务抢占计数
    #[inline(always)]
    pub fn preempt_dec(&self) {
        self.preempt_sub(1);
    }

    /// 任务`pid`
    #[inline(always)]
    pub fn pid(&self) -> usize {
        self.pid.to_value()
    }
}

impl Drop for Task {
    fn drop(&mut self) {
        let stack_base = self.task_base.stack();
        kfree(Vaddr::from(stack_base));
    }
}

/// 当前任务引用
///
/// # Panics
/// 系统故障将会 panic.
#[inline(always)]
pub fn current() -> Arc<Task> {
    let this_task = unsafe { PERCPU_ENTRY_TASK[this_processor_id().raw()].as_ref().unwrap() };
    this_task.clone()
}

/// 当前任务快速访问引用
#[inline(always)]
pub fn current_fast() -> &'static Task {
    cfg_if::cfg_if! {
        if #[cfg(aarch64_moana)] {
            unsafe {
                let task: usize;
                core::arch::asm!(
                    "mrs {0}, sp_el0",
                    out(reg) task
                );
                &*(task as *const Task)
            }
        } else {
            unreachable!()
        }
    }
}

// 仅唤醒从核时获取对应处理器任务信息使用
#[inline]
pub(crate) fn get_idle_task_info(cpu: usize) -> (usize, usize) {
    unsafe {
        let ts = PERCPU_ENTRY_TASK_RAW[cpu];
        let task = &*(ts as *const Task);
        let stack = task.task_base.stack();
        (ts, stack)
    }
}

#[inline(always)]
#[allow(unused_variables)]
fn set_current_fast(ts: usize) {
    cfg_if::cfg_if! {
        if #[cfg(aarch64_moana)] {
            unsafe {
                core::arch::asm!(
                    "msr sp_el0, {0}",
                    in(reg) ts
                );
            }
        } else {
            unreachable!()
        }
    }
}

static DUMMY_TASK: Task = Task::new();

pub(crate) fn early_task_init() {
    set_current_fast(&raw const DUMMY_TASK as usize);
}

fn root_task_init() -> usize {
    let tsk = Arc::new(Task::init(
        "root_server",
        0,
        Some(Arc::new(MmStruct::create().unwrap())),
        init_stack_base().to_value(),
    ));
    let raw_tsk = tsk.as_ref() as *const Task as usize;
    switch_entry_task(0, tsk);
    raw_tsk
}

#[allow(static_mut_refs)]
pub(crate) fn task_init() {
    unsafe {
        for _ in 0..nr_cpus() {
            PERCPU_ENTRY_TASK.push(None);
        }
    }

    pid_init();

    let root_tsk = root_task_init();

    for idx in 1..nr_cpus() {
        let idle = idle_create(idx);
        switch_entry_task(idx, idle);
    }

    let preempt = current_fast().preempt();
    set_current_fast(root_tsk);
    assert_eq!(preempt, current_fast().preempt());

    // ok. 切换用户地址空间
    current().mm().unwrap().switch_mm();
}

// NOTE: 只是创建 idle 任务, 不要创建重复了
pub(crate) fn idle_create(cpu: usize) -> Arc<Task> {
    Arc::new(Task::create(&format!("idle/{cpu}"), cpu, None))
}

#[inline(always)]
pub(crate) fn switch_entry_task(cpu: usize, tsk: Arc<Task>) {
    let raw_tsk = tsk.as_ref() as *const Task as usize;
    unsafe {
        PERCPU_ENTRY_TASK[cpu] = Some(tsk);
        PERCPU_ENTRY_TASK_RAW[cpu] = raw_tsk;
    }
}

// 记录当前任务地址
static mut PERCPU_ENTRY_TASK: Vec<Option<Arc<Task>>> = Vec::new();

// 记录当前任务原始指针
//
// # NOTE
// 仅用于汇编使用, 并且与`PERCPU_ENTRY_TASK`同步更新
#[unsafe(no_mangle)]
static mut PERCPU_ENTRY_TASK_RAW: [usize; NR_CPUS] = [0; NR_CPUS];
