use alloc::collections::BTreeMap;
use alloc::sync::Arc;
use alloc::vec::Vec;
use core::arch::asm;
use bitflags::bitflags;
use lazy_static::lazy_static;
use riscv::register::satp;
use xmas_elf;
use crate::sync::UPSafeCell;
use crate::config::{MEMORY_END, PAGE_SIZE, TRAMPOLINE, TRAP_CONTEXT, USER_STACK_SIZE};
use crate::mm::address::{PhysAddr, PhysPageNum, VirtAddr, VirtPageNum, VPNRange};
use crate::mm::frame_allocator::{frame_alloc, FrameTracker};
use crate::mm::page_table::{PageTable, PTEFlags};
use crate::println;

// 引入外部符号进来, 下面映射地址空间的逻辑段会用到
extern "C" {
    fn stext();
    fn etext();
    fn srodata();
    fn erodata();
    fn sdata();
    fn edata();
    fn sbss_with_stack();
    fn ebss();
    fn ekernel();
    fn strampoline();  // 这个符号 在linker.ld 中定义
}
lazy_static! {
    /// a memory set instance through lazy_static! managing kernel space
    pub static ref KERNEL_SPACE: Arc<UPSafeCell<MemorySet>> =
        Arc::new(unsafe { UPSafeCell::new(MemorySet::new_kernel()) });
}

// 某逻辑段的映射方式
#[derive(Copy, Clone, PartialEq, Debug)]
pub enum MapType {
    Identical,  // 恒等映射
    Framed,  // 逻辑映射
}


// 这个段的权限, 它是页表项标志位 PTEFlags 的一个子集, 仅仅保留4个标志位就足够了
bitflags! {
    /// map permission corresponding to that in pte: `R W X U`
    pub struct MapPermission: u8 {
        const R = 1 << 1;
        const W = 1 << 2;
        const X = 1 << 3;
        const U = 1 << 4;
    }
}


// 代表一个逻辑段
pub struct MapArea {
    vpn_range: VPNRange,  // 逻辑段的虚拟页起始位置 和 结束位置
    data_frames: BTreeMap<VirtPageNum, FrameTracker>,  // 这个逻辑段   虚拟页 对应的具体的物理帧, 这里使用rust 进行管理其物理帧的释放, 自动drop进行还给物理页帧管理器
    map_type: MapType,
    map_perm: MapPermission,
}

impl MapArea {
    // 从虚拟地址起始和结束, 创建一个逻辑段结构体
    pub fn from(start_va: VirtAddr, end_va: VirtAddr, map_type: MapType, map_perm: MapPermission,) -> Self {
        // 截取抛弃内存碎片
        let start_vpn: VirtPageNum = start_va.floor();
        let end_vpn: VirtPageNum = end_va.ceil();

        Self {
            vpn_range: VPNRange {l:start_vpn, r: end_vpn},
            data_frames: BTreeMap::new(),
            map_type,
            map_perm,
        }
    }
}

impl MapArea {
    // 把一个虚拟页号 创建一个对应的物理页帧 映射到page_table中, 并把创建的物理页帧生命周期 绑定到 self上
    pub fn map_one(&mut self, page_table: &mut PageTable, vpn: VirtPageNum) {
        let ppn: PhysPageNum;
        match self.map_type {
            // 恒等映射
            MapType::Identical => {
                ppn = PhysPageNum(vpn.0);
            }
            MapType::Framed => {
                // 随便创建一个 物理页帧
                let frame = frame_alloc().unwrap();
                ppn = frame.ppn;
                // 把随便创建的物理页帧, 绑定到data_frames 中被 self管理生命周期
                self.data_frames.insert(vpn, frame);
            }
        }

        // 根据上面创建的物理页帧映射pte 到三级页表中
        let pte_flags = PTEFlags::from_bits(self.map_perm.bits).unwrap();
        page_table.map(vpn, ppn, pte_flags);
    }

    // 取消映射一个vpn, 并把 self管理的 物理页帧的生命周期结束
    pub fn unmap_one(&mut self, page_table: &mut PageTable, vpn: VirtPageNum) {
        if self.map_type == MapType::Framed {
            // 删除映射关系, 自动调用 drop 把vpn对应的ppn还给全局内存管理器
            self.data_frames.remove(&vpn);
        }
        // 物理帧不存在了 把vpn 对应的pte设置为空
        page_table.unmap(vpn);
    }

    // 映射 当前逻辑段中的每一个虚拟页, 到页表中
    pub fn map(&mut self, page_table: &mut PageTable) {

        for vpn in self.vpn_range {
            self.map_one(page_table, vpn);
        }
    }

    // 在页表中取消映射每一个虚拟页
    pub fn unmap(&mut self, page_table: &mut PageTable) {
        for vpn in self.vpn_range {
            self.unmap_one(page_table, vpn);
        }
    }

    // 拷贝外部的数据二进制 到当前的逻辑段中 按照页 为单位
    pub fn copy_data(&mut self, page_table: &mut PageTable, data: &[u8]) {
        assert_eq!(self.map_type, MapType::Framed);

        let mut start: usize = 0;

        // 当前逻辑段的起始页位置
        let mut current_vpn = self.vpn_range.l;
        // 需要拷贝的字节总大小
        let len = data.len();
        loop {
            // 拷贝 每次最大拷贝1个页
            let src = &data[start..len.min(start + PAGE_SIZE)];

            // 得到逻辑段每个页的起始位置
            let dst = &mut page_table
                .get_pte(current_vpn)  // 在页表中, 根据逻辑段中的vpn 得到 54bit的pte (44+10)
                .unwrap()
                .ppn()  // 根据54bit的pte, 得到44bit 的物理页帧的位置
                .get_bytes_array()[..src.len()];  // 得到物理页帧那一页的起始的地址 44bit<<12
            dst.copy_from_slice(src);
            start += PAGE_SIZE;
            if start >= len {
                break;
            }
            current_vpn = (current_vpn.0 + 1).into();
        }
    }
    // 当前段, 缩小到新的结束位置, 即把new_end后面的页, 都还给物理页帧管理器, 并取消映射(通过unmap_one方法即可)
    pub fn shrink_to(&mut self, page_table: &mut PageTable, new_end: VirtPageNum) {
        let tmp_vpn_range = VPNRange{
            l: new_end,
            r: self.vpn_range.r,
        };
        for vpn in tmp_vpn_range {
            self.unmap_one(page_table, vpn)
        }

        // 注意 当前段的区间 也要调整哦new(self.vpn_range.get_start(), new_end);
        self.vpn_range.r = new_end;
    }


    // 当前段, 把当前结束的位置, 扩展到新的结束的位置
    pub fn append_to(&mut self, page_table: &mut PageTable, new_end: VirtPageNum) {
        let tmp_vpn_range = VPNRange{
            l: self.vpn_range.r,
            r: new_end,
        };

        for vpn in tmp_vpn_range {
            self.map_one(page_table, vpn)
        }

        self.vpn_range.r = new_end;
    }
}


// 表示一个地址空间, 它是一个页表和 包含所有段信息的页表组成
pub struct MemorySet {
    pub page_table: PageTable,
    pub areas: Vec<MapArea>,
}

impl MemorySet {
    pub fn new() -> Self {
        Self {
            page_table: PageTable::new(),
            areas: Vec::new(),
        }
    }

    pub fn from_elf(elf_data: &[u8]) -> (Self, usize, usize) {
        let mut memory_set = Self::new();
        // map trampoline
        // 创建一个地址空间
        memory_set.map_trampoline();
        // map program headers of elf, with U flag
        // 分析外部传进来的 elf文件数据
        let elf = xmas_elf::ElfFile::new(elf_data).unwrap();
        let elf_header = elf.header;
        let magic = elf_header.pt1.magic;
        assert_eq!(magic, [0x7f, 0x45, 0x4c, 0x46], "invalid elf!");
        let ph_count = elf_header.pt2.ph_count();
        // 得到所有的段数量
        let mut max_end_vpn = VirtPageNum(0);
        // 循环所有的段
        for i in 0..ph_count {
            // 得到当前段
            let ph = elf.program_header(i).unwrap();
            // 确定有加载的必要
            if ph.get_type().unwrap() == xmas_elf::program::Type::Load {
                // 当前段开始位置
                let start_va: VirtAddr = (ph.virtual_addr() as usize).into();
                // 当前段结束位置 (开始位置+大小)
                let end_va: VirtAddr = ((ph.virtual_addr() + ph.mem_size()) as usize).into();

                // 当前段权限
                let mut map_perm = MapPermission::U;
                let ph_flags = ph.flags();
                if ph_flags.is_read() {
                    map_perm |= MapPermission::R;
                }
                if ph_flags.is_write() {
                    map_perm |= MapPermission::W;
                }
                if ph_flags.is_execute() {
                    map_perm |= MapPermission::X;
                }
                // 创建一个 逻辑段
                let map_area = MapArea::from(start_va, end_va, MapType::Framed, map_perm);
                // 每次都更新最后的逻辑段的最后结束最大的逻辑地址
                max_end_vpn = map_area.vpn_range.r;

                // 当前逻辑段, 添加到地址空间, 并copy数据进去
                memory_set.push(
                    map_area,
                    Some(&elf.input[ph.offset() as usize..(ph.offset() + ph.file_size()) as usize]),
                );

            }
        }

        // map user stack with U flags
        // 将最后一个最大的逻辑段的结束的页, 转为逻辑地址  左移12位 得到39位的VirtAddr
        let max_end_va: VirtAddr = max_end_vpn.into();
        // 根据riscv 的规则, 转成一个 合法的地址 (todo 为什么这里需要转成合法的地址:将高位全部设置为1, 上面的start_va 不转就传?, 而且不直接可以使用end_va 吗? 这里还创建一个max_end_vpn?)
        let mut user_stack_bottom: usize = max_end_va.into();

        // guard page  放置 一页 保护页隔离用户栈
        user_stack_bottom += PAGE_SIZE;

        // 保护页之后 再增加 用户栈大小,得到栈底和当前的栈顶
        // 此时 user_stack_top 距离 max_end_vpn 三个页
        let user_stack_top = user_stack_bottom + USER_STACK_SIZE;

        // 添加 用户逻辑栈段并映射物理地址页, 到当前的地址空间
        memory_set.push(
            MapArea::from(
                user_stack_bottom.into(),
                user_stack_top.into(),
                MapType::Framed,
                MapPermission::R | MapPermission::W | MapPermission::U,
            ),
            None,
        );

        // 用户的堆空间段 只有一个起始位置即可, 因为 user_stack_top上面全是堆的空间(除了最后两页)
        memory_set.push(
            MapArea::from(
                user_stack_top.into(),
                user_stack_top.into(),
                MapType::Framed,
                MapPermission::R | MapPermission::W | MapPermission::U,
            ),
            None,
        );

        // 设置 TrapContext为统一的虚拟地址
        // 将 虚拟地址的 次高页添加到用户空间, 后面会被内核强行找到物理地址并修改为 trap context
        memory_set.push(
            MapArea::from(
                TRAP_CONTEXT.into(),
                TRAMPOLINE.into(),
                MapType::Framed,
                MapPermission::R | MapPermission::W,
            ),
            None,
        );

        (
            memory_set,  // 地址空间
            user_stack_top,  // 用户栈底和当前的栈顶, 他在地址空间各个逻辑段的上面, 同时
            elf.header.pt2.entry_point() as usize,  // elf文件的入口地址
        )
    }
}

impl MemorySet {
    // 得到当前地址空间的 页表token
    pub fn token(&self) -> usize {
        self.page_table.token()
    }

    // 激活当前地址空间
    pub fn activate(&self) {
        // let satp = self.page_table.token();
        // 得到当前地址空间的 页表 的规定格式
        let satp = self.token();;
        unsafe {
            // 设置satp CSR寄存器为我们
            satp::write(satp);

            // 刷新快表缓存TLB, 这个指令最简单的作用就是清空旧的快表
            // 更深层次其实是一个内存屏障
            asm!("sfence.vma");
        }
    }

    // 根据 start va 和end va 创建一个 逻辑段类型为逻辑映射, 插入到 当前地址空间
    pub fn insert_framed_area(&mut self, start_va: VirtAddr, end_va: VirtAddr, permission: MapPermission) {
        self.push(MapArea::from(start_va, end_va, MapType::Framed, permission), None);
    }

    // 把一个 逻辑段, 映射到当前页表中
    fn push(&mut self, mut map_area: MapArea, data: Option<&[u8]>) {
        map_area.map(&mut self.page_table);
        if let Some(data) = data {
            map_area.copy_data(&mut self.page_table, data);
        }
        self.areas.push(map_area);
    }

    /// 创建跳板, 跳板就是 trap.S中的代码段中.text.trampoline, 由 linker.ld手动布局在代码段
    /// 内核空间和用户空间,他们都会运行这个函数, 插入同样的代码, 把虚拟地址映射为同样的物理地址, 这里内核空间和用户空间执行这段代码的时候, 每个  虚拟地址对应的物理地址是相同的
    /// 所以 (内核空间的 TRAMPOLINE+指令offset) == (用户空间的 TRAMPOLINE+offset)  == (strampoline + offset)
    /// (内核的 TRAMPOLINE) 和 (用户的 TRAMPOLINE) 和 strampoline 他们都是相等的
    /// 在 trap.S 切换用户空间之后, 取值执行, 切换空间之前和切换空间之后, 他们对应的虚拟地址空间是相同的, 且都映射到了同一个物理页帧
    fn map_trampoline(&mut self) {
        // 直接插入到 把 linker.ld 中的符号strampoline, 映射到最高的虚拟页, 他的物理页帧是 链接脚本中我们手动内存布局的符号
        // 如果是用户应用, 这个符号的最下面就是 次高页 即 trap context
        self.page_table.map(
            VirtAddr::from(TRAMPOLINE).into(),
            PhysAddr::from(strampoline as usize).into(),
            PTEFlags::R | PTEFlags::X,
        );
    }

    // 创建一个内核用的地址空间
    pub fn new_kernel() -> Self {
        let mut memory_set = Self::new();
        // 跳板创建
        memory_set.map_trampoline();
        // map kernel sections
        println!(".text [{:#x}, {:#x})", stext as usize, etext as usize);
        println!(".rodata [{:#x}, {:#x})", srodata as usize, erodata as usize);
        println!(".data [{:#x}, {:#x})", sdata as usize, edata as usize);
        println!(
            ".bss [{:#x}, {:#x})",
            sbss_with_stack as usize, ebss as usize
        );
        println!("mapping .text section");
        memory_set.push(
            MapArea::from(
                (stext as usize).into(),
                (etext as usize).into(),
                MapType::Identical,
                MapPermission::R | MapPermission::X,
            ),
            None,
        );
        println!("mapping .rodata section");
        memory_set.push(
            MapArea::from(
                (srodata as usize).into(),
                (erodata as usize).into(),
                MapType::Identical,
                MapPermission::R,
            ),
            None,
        );
        println!("mapping .data section");
        memory_set.push(
            MapArea::from(
                (sdata as usize).into(),
                (edata as usize).into(),
                MapType::Identical,
                MapPermission::R | MapPermission::W,
            ),
            None,
        );
        println!("mapping .bss section");
        memory_set.push(
            MapArea::from(
                (sbss_with_stack as usize).into(),
                (ebss as usize).into(),
                MapType::Identical,
                MapPermission::R | MapPermission::W,
            ),
            None,
        );
        println!("mapping physical memory");
        memory_set.push(
            MapArea::from(
                (ekernel as usize).into(),
                MEMORY_END.into(),
                MapType::Identical,
                MapPermission::R | MapPermission::W,
            ),
            None,
        );
        memory_set
    }


    // 当前地址空间, 缩小 虚拟地址start位置到虚拟地址 new end的位置
    pub fn shrink_to(&mut self, start: VirtAddr, new_end: VirtAddr) -> bool {
        // 起始虚拟地址所在的页
        let start_vpn = start.floor();

        // 找到起始地址所在的段, 传入页表, 把当前段缩小到指定结束位置new_end
        if let Some(area) = self.areas.iter_mut().find(|area| area.vpn_range.l == start_vpn) {
            area.shrink_to(&mut self.page_table, new_end.ceil());
            true
        } else {
            false
        }
    }

    // 增加当前的地址空间 扩展为 start到 新的结束的位置
    pub fn append_to(&mut self, start: VirtAddr, new_end: VirtAddr) -> bool {
        // 同上, 找到逻辑段之后, 只不过这里是 扩展
        if let Some(area) = self.areas.iter_mut().find(|area| area.vpn_range.l == start.floor()) {
            area.append_to(&mut self.page_table, new_end.ceil());
            true
        } else {
            false
        }
    }
}
