#include <assert.h>
#include <x86.h>
#include <errno.h>
#include <string.h>

#include <kern/fork.h>
#include <kern/syscall.h>
#include <kern/trap.h>
#include <kern/pmap.h>
#include <kern/kmalloc.h>
#include <kern/sche.h>
#include <kern/stdio.h>

// get the flag of a lin addr in page table
static inline u32 get_flag(uintptr_t laddr, u32 cr3) {
	assert(PGOFF(laddr) == 0);
	uintptr_t *pde_ptr = (uintptr_t *)K_PHY2LIN(cr3);
	assert((pde_ptr[PDX(laddr)] & PTE_P) != 0);
	phyaddr_t pte_phy = PTE_ADDR(pde_ptr[PDX(laddr)]);
	uintptr_t *pte_ptr = (uintptr_t *)K_PHY2LIN(pte_phy);
	assert((pte_ptr[PTX(laddr)] & PTE_P) != 0);
	return PTE_FLAG(pte_ptr[PTX(laddr)]);
}

static phyaddr_t
alloc_phy_page(struct page_node **page_list)
{
	phyaddr_t paddr = phy_malloc_4k();
	
	struct page_node *new_node = kmalloc(sizeof(struct page_node));
	new_node->nxt = *page_list;
	new_node->paddr = paddr;
	new_node->laddr = -1;
	*page_list = new_node;

	return paddr;
}

static phyaddr_t
lin_mapping_phy(u32			cr3,
		struct page_node	**page_list,
		uintptr_t		laddr,
		u32			pte_flag)
{
	assert(PGOFF(laddr) == 0);

	uintptr_t *pde_ptr = (uintptr_t *)K_PHY2LIN(cr3);

	if ((pde_ptr[PDX(laddr)] & PTE_P) == 0) {
		phyaddr_t pte_phy = alloc_phy_page(page_list);
		memset((void *)K_PHY2LIN(pte_phy), 0, PGSIZE);
		pde_ptr[PDX(laddr)] = pte_phy | PTE_P | PTE_W | PTE_U;
	}

	phyaddr_t pte_phy = PTE_ADDR(pde_ptr[PDX(laddr)]);
	uintptr_t *pte_ptr = (uintptr_t *)K_PHY2LIN(pte_phy);

	phyaddr_t page_phy;
	if ((pte_ptr[PTX(laddr)] & PTE_P) != 0)
		warn("fork alloc: this page was mapped before, laddr: %x", laddr);
	page_phy = alloc_phy_page(page_list);
	(*page_list)->laddr = laddr;
	pte_ptr[PTX(laddr)] = page_phy | pte_flag;
	return page_phy;
}

// modified from pmap.c::map_elf
// first alloc a new page list, together with a new cr3(with page directory page)
// second map kern
// third iterate over parent pages, 
//     for each parent page(except page table pages, child has its own) , map a page in child proc
//     then copy content of this page to child
// finally, set child->page_list to this brand new one and set cr3
// SET: child's page list and cr3
static inline void copy_parent_pages(PROCESS_0 *child, PROCESS_0 *parent){
	phyaddr_t new_cr3 = phy_malloc_4k();
	memset((void *)K_PHY2LIN(new_cr3), 0, PGSIZE);
	struct page_node *new_page_list = kmalloc(sizeof(struct page_node));
	new_page_list->nxt = NULL;
	new_page_list->paddr = new_cr3;
	new_page_list->laddr = -1;
	// maybe redundant, let it be for now
	map_kern(new_cr3, &new_page_list);
	// iterate over parent pages
	for (struct page_node* pa_page = parent->page_list; pa_page != NULL; pa_page = pa_page->nxt) {
		// no need to take care of non-physical page, we have our own page table and cr3
		if (pa_page->laddr == -1) continue;
		// first alloc a new page for the same laddr, and set flag same as parent
		phyaddr_t new_paddr = lin_mapping_phy(new_cr3, &new_page_list, pa_page->laddr, get_flag(pa_page->laddr, parent->cr3));
		// copy parent page to child page using paddr
		memcpy((void*)K_PHY2LIN(new_paddr), (const void*)K_PHY2LIN(pa_page->paddr), PGSIZE);
		// kprintf("%x -> %x\n", K_PHY2LIN(pa_page->paddr), K_PHY2LIN(new_paddr));
	}
	child->page_list = new_page_list;
	child->cr3 = new_cr3;
}

ssize_t
kern_fork(PROCESS_0 *p_fa)
{
	// 这可能是你第一次实现一个比较完整的功能，你可能会比较畏惧
	// 但是放心，别怕，先别想自己要实现一个这么大的东西而毫无思路
	// 这样你在焦虑的同时也在浪费时间，就跟你在实验五中被页表折磨一样
	// 人在触碰到未知的时候总是害怕的，这是天性，所以请你先冷静下来
	// fork系统调用会一步步引导你写出来，不会让你本科造火箭的
	// panic("Unimplement! CALM DOWN!");
	
	// 推荐是边写边想，而不是想一车然后写，这样非常容易计划赶不上变化
	//? before all, lock father
	while (xchg(&p_fa->lock, 1) == 1)
		schedule();
	// fork的第一步你需要找一个空闲（IDLE）的进程作为你要fork的子进程
	PROCESS *proc_child = NULL;
	int i = 0;
	//? maybe cli will be better when preserving a proc resource
	DISABLE_INT();
	for (i = 1; i < PCB_SIZE; ++ i) {
		if (proc_table[i].pcb.statu == IDLE) {
			proc_table[i].pcb.statu = INITING;
			proc_child = &proc_table[i];
			break;
		}
	}
	ENABLE_INT();
	if (proc_child == NULL) {
		// NO hell. no free proc_table found. 
		xchg(&p_fa->lock, 0);
		return -EAGAIN;
	}
	assert(proc_child->pcb.lock == 0);
	while (xchg(&proc_child->pcb.lock, 1) == 1)
		schedule();
	PROCESS_0 *p_child = &proc_child->pcb;
	// 再之后你需要做的是好好阅读一下pcb的数据结构，搞明白结构体中每个成员的语义
	// 别光扫一遍，要搞明白这个成员到底在哪里被用到了，具体是怎么用的
	// 可能exec和exit系统调用的代码能够帮助你对pcb的理解，不先理解好pcb你fork是无从下手的
	// panic("Unimplement! read pcb");

	// 在阅读完pcb之后终于可以开始fork工作了
	// 本质上相当于将父进程的pcb内容复制到子进程pcb中
	// 但是你需要想清楚，哪些应该复制到子进程，哪些不应该复制，哪些应该子进程自己初始化
	// 其中有三个难点
	// 1. 子进程"fork"的返回值怎么处理？（需要你对系统调用整个过程都掌握比较清楚，如果非常清晰这个问题不会很大）
	// 2. 子进程内存如何复制？（别傻乎乎地复制父进程的cr3，本质上相当于与父进程共享同一块内存，
	// 而共享内存肯定不符合fork的语义，这样一个进程写内存某块地方会影响到另一个进程，这个东西需要你自己思考如何复制父进程的内存）
	// 3. 在fork结束后，肯定会调度到子进程，那么你怎么保证子进程能够正常进入用户态？
	// （你肯定会觉得这个问题问的莫名其妙的，只能说你如果遇到那一块问题了就会体会到这个问题的重要性，
	// 这需要你对调度整个过程都掌握比较清楚）
	//? Start to COPY
	//?   1. COPY PCB
	// p_fa shares the same base addr as its padding
	// anyways, copy the whole proc_table item(8K?)
	DISABLE_INT(); // make sure this process is never interrupted
	// memset(proc_child, 0, sizeof(PROCESS)); // clear child's kernel stack //? seem to be useless
	// the commented fields below will be set later
	// p_child->cr3
	p_child->exit_code = p_fa->exit_code;
	// p_child->fork_tree
	// p_child->kern_regs
	// p_child->lock
	// p_child->page_list
	p_child->pid = i;
	p_child->priority = p_fa->priority;
	p_child->statu = INITING; //! important!!! if you copied this from parent, haha, waste one hour
	p_child->ticks = p_fa->ticks;
	p_child->user_regs = p_fa->user_regs;
	ENABLE_INT(); // TODO: may be a useless atomic block, try to remove and test again
	//?   2. ALLOC PAGES AND COPY PHYSICAL MEMORY
	copy_parent_pages(p_child, p_fa);
	// panic("Unimplement! soul torture1");
	//?   3. SET restart point
	// //! maybe, kern stack different, use offset relevant to stack base addr to set child's kern esp
	// u32 esp_off = p_fa->kern_regs.esp - (u32)p_fa;
	// p_child->kern_regs.esp = (u32)p_child + esp_off;
	// u32 ebp_off = p_fa->kern_regs.ebp - (u32)p_fa;
	// p_child->kern_regs.ebp = (u32)p_child + ebp_off;
	// kprintf("%x %x\n", p_child->kern_regs.esp, p_fa->kern_regs.esp);
	// kprintf("%x %x\n", p_child->user_regs.kernel_esp, p_fa->user_regs.kernel_esp);
	//! kern_ctx is not part of our FSM model of user process, never copy it from parent
	//! or 1 another hour will be wasted
	// just set it as a new process from scratch, directly jump to restart
	p_child->kern_regs.esp = (u32)(proc_child + 1) - 8;
	*(u32 *)(p_child->kern_regs.esp + 0) = (u32)restart;
	*(u32 *)(p_child->kern_regs.esp + 4) = (u32)p_child;
	p_child->user_regs.eax = 0; // eax as the retval

	// 别忘了维护进程树，将这对父子进程关系添加进去
	// ? maintain process tree
	p_child->fork_tree.p_fa = p_fa;
	// malloc a son_node
	struct son_node* p_son = (struct son_node*)kmalloc(sizeof(struct son_node));
	p_son->p_son = p_child;
	// head insert into parent's son list
	p_son->pre = NULL;
	p_son->nxt = p_fa->fork_tree.sons;
	if (p_fa->fork_tree.sons != NULL) {
		p_son->nxt->pre = p_son;
	}
	p_fa->fork_tree.sons = p_son;
	// 最后你需要将子进程的状态置为READY，说明fork已经好了，子进程准备就绪了
	p_child->statu = READY;
	xchg(&p_child->lock, 0);
	xchg(&p_fa->lock, 0);
	// 在你写完fork代码时先别急着运行跑，先要对自己来个灵魂拷问
	// 1. 上锁上了吗？所有临界情况都考虑到了吗？（永远要相信有各种奇奇怪怪的并发问题）
	// 2. 所有错误情况都判断到了吗？错误情况怎么处理？（RTFM->`man 2 fork`）
	// 3. 你写的代码真的符合fork语义吗？
	// panic("Unimplement! soul torture");
	// ? return to parent proc by the normal syscall return machanism
	return p_child->pid;
}

ssize_t
do_fork(void)
{
	return kern_fork(&p_proc_ready->pcb);
}