use alloc::sync::{Arc, Weak};
use alloc::vec::Vec;
use core::cell::RefMut;
use crate::config::{TRAP_CONTEXT};
use crate::mm::address::{PhysPageNum, VirtAddr};
use crate::mm::memory_set::{KERNEL_SPACE, MapPermission, MemorySet};
use crate::println;
use crate::sync::UPSafeCell;
use crate::task::context::{TaskContext};
use crate::task::pid::{kernel_stack_position, KernelStack, pid_alloc, PidHandle};
use crate::trap::{trap_handler, TrapContext};

// TCB的字段, 用来保存任务的状态
#[derive(Copy, Clone, PartialEq)]
pub enum TaskStatus {
    Ready,
    Running,
    Zombie,
}

// 把 ch4 的TaskManagerInner 和 TaskControlBlock 整合为 TaskControlBlock, 现在被 TaskManager 关联和管理
// 在内核中, 它等价与一个进程
pub struct TaskControlBlock {
    // 初始化后不再变化的元数据
    pub pid: PidHandle,
    pub kernel_stack: KernelStack,

    // 初始化后 需要变化的元数据
    inner: UPSafeCell<TaskControlBlockInner>,
}

pub struct TaskControlBlockInner {
    pub user_time: usize,  // 用户态程序用的时间
    pub kernel_time: usize,  // 内核态程序所用的时间
    pub trap_cx_ppn: PhysPageNum,  // tcb访问 trap context所在的真实的物理页, 它对应逻辑页的次高页
    pub task_cx: TaskContext,
    pub task_status: TaskStatus,
    pub memory_set: MemorySet,  // tcb他自己的地址空间
    pub parent: Option<Weak<TaskControlBlock>>,  // 不包含父进程引用计数的弱指针
    pub children: Vec<Arc<TaskControlBlock>>,  // 所有的子进程 arc智能指针, 只有引用计数为0的时候, 绑定的各类资源才能回收
    pub exit_code: i32,  // 该进程退出时的 退出码
    pub base_size: usize,  // 应用地址空间中从0x0开始到用户栈结束一共包含多少字节, 就是用户数据有多大
    pub heap_bottom: usize,  // 堆的起始地址 这个一开始和base_size 是一样的
    pub program_brk: usize,  // 这个一开始和base_size 是一样的, 表示进程当前的堆边界, 即堆的顶部地址, 它指向堆中最后一个已分配内存块的末尾, 下一个内存分配将从该地址开始
    //   base_size(低地址) --- TRAP_CONTEXT(次高页) 的虚拟空间, 可以当做堆来使用
}

impl TaskControlBlock {
    // 得到具有内部可变性的字段
    pub fn inner_exclusive_access(&self) -> RefMut<'_, TaskControlBlockInner> {
        self.inner.exclusive_access()
    }
    pub fn get_pid(&self) -> usize {
        self.pid.0
    }
}

impl TaskControlBlock {
    // 根据 elf data 新建一个 TCB, 这个只在初始化 INITPROC 时候用
    pub fn from(elf_data: &[u8]) -> Self {
        // memory_set with elf program headers/trampoline/trap context/user stack
        // 返回 内存空间, 用户栈的栈底和当前栈顶, 和入口地址
        let (memory_set, user_sp, entry_point) = MemorySet::from_elf(elf_data);
        // 在当前的地址空间中, 找到 trap_context 逻辑页 所在的真实物理页, 他逻辑地址在次高页, 真实物理地址是我们在 from_elf 申请映射的未知地址
        // 这个真实的物理地址, 我们后续再陷入的时候会用到, 进行强写数据
        let trap_cx_ppn = memory_set
            .page_table
            .get_pte(VirtAddr::from(TRAP_CONTEXT).into())  // 虚拟地址次高页, 所在对应的物理内存的pte,
            .unwrap()
            .ppn();  // pte 内的ppn, 得到实际 虚拟内存次高页 所在的物理页号

        // 从全局pid分配器, 得到一个pid
        let pid_handle = pid_alloc();

        // 根据pid, 得到内核栈并得到高地址(栈底)
        let kernel_stack = KernelStack::from(&pid_handle);
        let kernel_stack_top = kernel_stack.get_top();

        // 构建 任务控制块
        let task_control_block = Self {
            pid: pid_handle,
            kernel_stack: kernel_stack,
            inner: unsafe {
                UPSafeCell::new(TaskControlBlockInner{
                    user_time: 0,
                    kernel_time: 0,
                    task_status:TaskStatus::Ready,
                    task_cx: TaskContext::goto_trap_return(kernel_stack_top),  // 根据内核栈构造任务切换上下文, 并把 switch任务切换的 ra 设置为 trap_return的函数地址
                    memory_set,  // 新增, 当前任务的地址空间
                    parent: None,
                    children: Vec::new(),
                    trap_cx_ppn,  // 新增  逻辑次高页的trap context, 对应的这个是真实的物理页, 我们这里保存一份, 省的在memory_set 里面查找了
                    base_size: user_sp,  // 用户栈顶以下 是代码中的各种逻辑段+栈段, 应用地址空间中从0x0开始到用户栈结束一共包含多少字节, 所以大小就是截止到user_sp
                    heap_bottom: user_sp,  //
                    program_brk: user_sp,  //
                    exit_code: 0,
                })
            }
        };

        // prepare TrapContext in user space
        // 根据 trap_cx_ppn 构建 陷入 trap context 的结构体
        let trap_cx = task_control_block.inner_exclusive_access().get_trap_cx();  // 根据trap_cx_ppn得到真实的物理页的 trap context的地址
        *trap_cx = TrapContext::from(  // 对 tcb 的 TrapContext 的物理页进行 修改,
           entry_point,  // 在陷入完成后, 准备返回用户态执行的, 用户代码入口地址
           user_sp,  // 用户栈, 这个栈的地址是虚拟地址
           KERNEL_SPACE.exclusive_access().token(),  // 内核 satp 页表的寄存器信息
           kernel_stack_top,  // 用户应用的内核栈顶  在内核空间的位置 (这个是根据appid 计算出来的, 也只能在内核的地址空间中才能根据这个地址 看到相关的栈)
           trap_handler as usize,  // 内核中 trap handler的入口点虚拟地址(恒等映射所以 这里也能找到)
        );
        task_control_block
        // 以上:
        // 任务控制块和 trap控制块, 分别多了一些字段
        // trap控制块(这个每次陷入都会被 sscratch 寄存器临时保存):
        //     多的三个字段 首次写入后, 后续只会读取和恢复,不会被再次写入新值, 这是为了
        //     kernel_satp 内核空间页表所在位置的实际物理地址, 这个值以后会在切换的时候替换到satp寄存器
        //     kernel_sp, 这个保存这个应用 他 根据appid计算出来 所持有的任务内核栈, 在jump context之前 会恢复
        //     trap_handler 这是在内核在陷入的时候, 我们自定义的 trap_handle 的地址, 在陷入保存完寄存器之后, 需要 jmp 到这个函数地址
        // task 控制块 TCB, 这个是在内核进行 switch时的上下文结构:
        //     memory_set 表示当前任务他的地址空间
        //     trap_cx_ppn 用来保存 上方 trap控制块的 实际的物理页帧, 在用户空间中都被应设在了TRAP_CONTEXT相关的次高位置虚拟地址处, 在内核空间我们这保存一份实际的物理地址, 在内核空间读取这个通过恒等映射就得到任务的 trap context
        //     base_size 用户栈顶距离0x0的距离, 即用户应用已知的大小
        //
    }
    
    // 从当前进程copy 一个子进程, 他和父进程有相同的地址空间, 但是它拥有自己的内核栈 用来存放 trap_context
    // 子进程后续运行会 条入到 trap_return
    pub fn fork(self: &Arc<Self>) -> Arc<Self> {
        // ---- 当前进程的内容
        let mut parent_inner = self.inner_exclusive_access();

        // 根据当前地址空间, copy 当前进程地址空间的数据到 一个新的地址空间
        let memory_set = MemorySet::from_existed_user(&parent_inner.memory_set);

        // 找到新的进程地址空间中 TRAP_CONTEXT 的物理页
        let trap_cx_ppn = memory_set
            .page_table
            .get_pte(VirtAddr::from(TRAP_CONTEXT).into())
            .unwrap()
            .ppn();

        // 分配一个pid
        let pid_handle = pid_alloc();
        // 根据pid 分配一个 子进程自己的内核栈
        let kernel_stack = KernelStack::from(&pid_handle);

        // 创建tcb
        let kernel_stack_top = kernel_stack.get_top();
        let task_control_block = Arc::new(TaskControlBlock {
            pid: pid_handle,
            kernel_stack,
            inner: unsafe {
                UPSafeCell::new(TaskControlBlockInner {
                    user_time: 0,
                    kernel_time: 0,
                    trap_cx_ppn,
                    base_size: parent_inner.base_size,
                    task_cx: TaskContext::goto_trap_return(kernel_stack_top),
                    task_status: TaskStatus::Ready,
                    memory_set,
                    parent: Some(Arc::downgrade(self)),  // 将父进程的弱引用计数 放到子进程的进程控制块中
                    children: Vec::new(),
                    exit_code: 0,
                    heap_bottom: 0,
                    program_brk: 0,
                })
            },
        });


        // add child
        // 添加子进程信息 到父进程中
        parent_inner.children.push(task_control_block.clone());

        // 得到 tcb 当中 trap context 的可变引用, 上面复制时还是复制的父进程的
        // 现在修改其栈顶为我们新创建的 子进程自己的栈顶, 其他字段则是完全和父进程一样
        let trap_cx = task_control_block.inner_exclusive_access().get_trap_cx();
        trap_cx.kernel_sp = kernel_stack_top;

        // return
        task_control_block
        // ---- release parent PCB automatically
        // **** release children PCB automatically
    }
    // TCB 加载一个新的应用, 从新的elf的 代码和数据替换 原有的旧的 应用地址空间的内容, 并开始执行
    // pid 和 kernel_stack 不变
    // 也不需要修改任务上下文环境, 因为只有被暂停的应用才需要再内核中保留一个任务上下文, 修改trap context 即可因
    pub fn exec(&self, elf_data: &[u8]) {
        // memory_set with elf program headers/trampoline/trap context/user stack
        // 根据 elf data 创建一个地址空间
        let (memory_set, user_sp, entry_point) = MemorySet::from_elf(elf_data);
        let trap_cx_ppn = memory_set
            .page_table
            .get_pte(VirtAddr::from(TRAP_CONTEXT).into())
            .unwrap()
            .ppn();

        // **** access inner exclusively
        let mut inner = self.inner_exclusive_access();
        // substitute memory_set
        // 原有地址空间生命周期结束, 全部被回收
        inner.memory_set = memory_set;
        // 更新trap ppn
        inner.trap_cx_ppn = trap_cx_ppn;
        // 更新应用大小
        inner.base_size = user_sp;
        // 得到trap context 真实内存的指针
        let trap_cx = inner.get_trap_cx();
        *trap_cx = TrapContext::from(
            entry_point,
            user_sp,
            KERNEL_SPACE.exclusive_access().token(),
            self.kernel_stack.get_top(),
            trap_handler as usize,
        );
        // 作用域结束 释放inner
        // **** release inner automatically
    }
}

impl TaskControlBlockInner{
    // 根据 trap context 的实际的物理地址, 强转为 TrapContext 结构体
    pub fn get_trap_cx(&self) -> &'static mut TrapContext {
        self.trap_cx_ppn.get_mut()
    }

    pub fn get_user_token(&self) -> usize {
        self.memory_set.token()
    }

    // 扩张/缩减, 当前应用的地址空间中的堆段, 指定字节的数据
    pub fn change_program_brk(&mut self, size: i32) -> Option<usize> {
        // 堆顶
        let old_break = self.program_brk;
        // 申请后新的堆顶
        let new_brk = self.program_brk as isize + size as isize;

        // 如果新的堆顶 小于 堆底 直接返回(防止过度释放)
        if new_brk < self.heap_bottom as isize {
            return None;
        }

        // 如果是小于0 说明需要释放,缩小堆的大小
        // 传入两个参数 堆底虚拟地址, 和 新的堆顶地址
        let result = if size < 0 {
            self.memory_set
                .shrink_to(VirtAddr(self.heap_bottom), VirtAddr(new_brk as usize))
        } else {
            self.memory_set
                .append_to(VirtAddr(self.heap_bottom), VirtAddr(new_brk as usize))
        };
        println!("size: {}, old_break: {}", size, old_break);
        // 如果成功就把新的堆顶的地址, 复制到结构提上记录, 并把旧的堆顶返回
        if result {
            self.program_brk = new_brk as usize;
            Some(old_break)
        } else {
            None
        }
    }
}