import mm.mm as mm
import paging.paging as paging
import drivers.serial.uart.serial as serial
import drivers.fs.fat32.fat32 as fat32
import drivers.fs.fat32.fat32_types as fat32_types
import drivers.core.cpu.cpu as cpu
import intr.gdt as gdt
import string


define PROC_READY       = 0
define PROC_RUNNING     = 1
define PROC_BLOCKED     = 2
define PROC_DONE        = 3
define MAX_PROCESSES    = 8
define MAX_PROC_THREADS = 8
define PROC_MAX_FDS     = 16
define PROC_MAX_DIRS    = 8
define KERN_STACK_SIZE  = 8192
define USER_STACK_SIZE  = 65536
# 用户栈虚拟区：主线程栈顶 = USER_STACK_VTOP（见 execrunner/elf.vp 的映射），
# 额外用户线程（THREAD_CREATE）从它下方按 USER_THREAD_STACK_GAP 逐个下移。
define USER_STACK_VTOP        = 0x7FFFFFE00000
define USER_STACK_PAGES       = 16
define USER_THREAD_STACK_GAP  = 0x20000
define MAX_USER_THREADS       = 16


class Process:
    pid:          int
    pml4_root:    uint64
    threads:      [8]int
    thread_count: int
    state:        int
    entry:        rawptr
    elf_base:     uint64
    heap_start:   uint64
    heap_size:    uint64
    name:         [32]char
    fd_table:     [PROC_MAX_FDS]rawptr
    fd_used:      [PROC_MAX_FDS]uint8
    dir_table:    [PROC_MAX_DIRS]rawptr
    dir_used:     [PROC_MAX_DIRS]uint8
    dir_pool:     [PROC_MAX_DIRS]fat32_types.dirobj
    kern_stack:   uint64
    user_stack:   uint64

    def addThread(self, tid: int) -> int:
        if self.thread_count >= MAX_PROC_THREADS: return -1
        self.threads[self.thread_count] = tid
        self.thread_count += 1
        return 0

    def exit(self):
        self.state = PROC_DONE
        for i in range(PROC_MAX_FDS):
            if self.fd_used[i]:
                fp: rawptr = self.fd_table[i]
                if fp:
                    fat32.close(fp)
                self.fd_used[i] = 0
                self.fd_table[i] = None
        for i in range(PROC_MAX_DIRS):
            if self.dir_used[i]:
                dp: rawptr = self.dir_table[i]
                if dp:
                    fat32.closedir(dp)
                self.dir_used[i] = 0
                self.dir_table[i] = None
        m: ProcessManager* = _proc_mgr_ptr
        if m.current_pid == self.pid:
            m.current_pid = 0
            m.processes[0].state = PROC_RUNNING
            if self.pml4_root != 0:
                kernel_cr3: uint64 = m.processes[0].pml4_root
                if kernel_cr3 != 0:
                    __asm__(f"mov cr3, {asminp(kernel_cr3, ASM_DESCR.REG_ANY)}",
                        op=[ASM_DESCR.CLOBBER_MEMORY])
                    paging.set_active_pml4(kernel_cr3)

    @property
    def is_running(self) -> bool: return self.state == PROC_RUNNING
    @property
    def is_done(self) -> bool: return self.state == PROC_DONE


class ProcessManager:
    processes: [MAX_PROCESSES]Process
    current_pid: int
    count: int

    def __init__(self):
        global _proc_mgr_ptr
        _proc_mgr_ptr = &_proc_mgr
        m: ProcessManager* = _proc_mgr_ptr
        m.current_pid = 0
        m.count = 1
        p: Process* = &m.processes[0]
        p.pid          = 0
        p.state        = PROC_RUNNING
        p.thread_count = 0
        p.pml4_root    = 0
        p.entry        = None
        p.elf_base     = 0
        p.heap_start   = 0
        p.heap_size    = 0
        name_str: str = "kernel"
        i: int = 0
        for ch in name_str:
            if i >= 31: break
            p.name[i] = ch
            i += 1
        p.name[i] = 0
        j: int
        for j in range(MAX_PROC_THREADS):
            p.threads[j] = -1
        for j in range(PROC_MAX_FDS):
            p.fd_used[j] = 0
            p.fd_table[j] = None
        for j in range(PROC_MAX_DIRS):
            p.dir_used[j] = 0
            p.dir_table[j] = None
        __asm__(f"""mov {asmout(p.pml4_root, ASM_DESCR.OUTPUT_REG)}, cr3""",
            op=[ASM_DESCR.CLOBBER_MEMORY])

    @staticmethod
    def _switch(pid: int) -> int:
        global _saved_cr3
        m: ProcessManager* = _proc_mgr_ptr
        if pid < 0 or pid >= m.count: return -1
        p: Process* = &m.processes[pid]
        if p.state == PROC_DONE: return -1
        old_pid: int = m.current_pid
        if old_pid == pid: return 0
        if old_pid >= 0 and old_pid < m.count:
            m.processes[old_pid].state = PROC_READY
        p.state = PROC_RUNNING
        m.current_pid = pid
        __asm__(f"""mov {asmout(_saved_cr3, ASM_DESCR.OUTPUT_REG)}, cr3
                mov rax, {asminp(p.pml4_root, ASM_DESCR.REG_ANY)}
                mov cr3, rax""",
            op=[ASM_DESCR.CLOBBER_MEMORY, ASM_DESCR.CLOBBER_RAX])
        return 0

    @staticmethod
    def current() -> Process*:
        m: ProcessManager* = _proc_mgr_ptr
        return &m.processes[m.current_pid]

    @staticmethod
    def get_process(pid: int) -> Process*:
        m: ProcessManager* = _proc_mgr_ptr
        if pid < 0 or pid >= m.count: return None
        return &m.processes[pid]

    @staticmethod
    def create_process(name: str, entry: rawptr) -> Process*:
        m: ProcessManager* = _proc_mgr_ptr
        if m.count >= MAX_PROCESSES: return None
        pid: int = m.count
        m.count += 1
        p: Process* = &m.processes[pid]
        p.pid          = pid
        p.state        = PROC_READY
        p.entry        = entry
        p.elf_base     = 0
        p.thread_count = 0
        p.heap_start   = 0
        p.heap_size    = 0
        string.memset(&p.name, 0, 32)
        i: int = 0
        for ch in name:
            if i >= 31: break
            p.name[i] = ch
            i += 1
        p.name[i] = 0
        for j in range(MAX_PROC_THREADS):
            p.threads[j] = -1
        for j in range(PROC_MAX_FDS):
            p.fd_used[j] = 0
            p.fd_table[j] = None
        for j in range(PROC_MAX_DIRS):
            p.dir_used[j] = 0
            p.dir_table[j] = None
        ks: rawptr = mm.malloc(KERN_STACK_SIZE)
        if ks is None:
            m.count = m.count - 1
            return None
        ks_size: uint64 = KERN_STACK_SIZE
        p.kern_stack = uint64(ks) + ks_size
        us: rawptr = mm.malloc(USER_STACK_SIZE)
        if us is None:
            serial.puts("[process] user stack alloc failed\n")
            mm.free(ks)
            m.count = m.count - 1
            return None
        us_size: uint64 = USER_STACK_SIZE
        p.user_stack = uint64(us) + us_size
        __asm__(f"""mov {asmout(p.pml4_root, ASM_DESCR.OUTPUT_REG)}, cr3""",
            op=[ASM_DESCR.CLOBBER_MEMORY, ASM_DESCR.CLOBBER_RAX])
        return p


_proc_mgr:     ProcessManager = ProcessManager()
_proc_mgr_ptr: ProcessManager* = &_proc_mgr
export _saved_cr3: uint64

def current() -> int:
    m: ProcessManager* = _proc_mgr_ptr
    return m.current_pid

def switch_pid(new_pid: int):
    m: ProcessManager* = _proc_mgr_ptr
    if m.current_pid == new_pid: return
    if m.current_pid >= 0 and m.current_pid < m.count:
        old_p: Process* = &m.processes[m.current_pid]
        if old_p.state != PROC_DONE:
            old_p.state = PROC_READY
    if new_pid >= 0 and new_pid < m.count:
        new_p: Process* = &m.processes[new_pid]
        new_p.state = PROC_RUNNING
        new_cr3: uint64 = new_p.pml4_root
        if new_cr3 != 0:
            __asm__(f"mov cr3, {asminp(new_cr3, ASM_DESCR.REG_ANY)}",
                op=[ASM_DESCR.CLOBBER_MEMORY])
            paging.set_active_pml4(new_cr3)
        # Update TSS RSP0 and GS kernel_rsp0 for the new process
        # This is needed so interrupts from Ring 3 use the correct kernel stack
        if new_p.kern_stack != 0:
            gdt.set_tss_rsp0(new_p.kern_stack)
            cpu.set_kernel_rsp0(new_p.kern_stack)
    m.current_pid = new_pid

def get_pid(tid: int) -> int:
    m: ProcessManager* = _proc_mgr_ptr
    for i in range(m.count):
        p: Process* = &m.processes[i]
        for j in range(p.thread_count):
            if p.threads[j] == tid:
                return p.pid
    return -1

def drop_to_user_mode(entry: rawptr, user_rsp: uint64, kern_rsp: uint64):
    # NOTE: Do NOT swapgs here. This OS uses reverse GS convention:
    #   GS_BASE = 0 (kernel & user), KERNEL_GS_BASE = &_per_cpu
    #   syscall/ISR entry does swapgs to get GS_BASE = &_per_cpu
    #   syscall/ISR return does swapgs to restore GS_BASE = 0
    # GS state is restored in isrCommonhandler before _yield() if needed.
    # [FIX] Normalize GS to the user convention (GS_BASE = 0) before returning to ring 3.
    # This thread may have been switched in from ANOTHER thread's user-mode timer interrupt:
    # that path never runs the ISR exit swapgs, so GS_BASE is still the kernel view (pointing
    # at the per-cpu struct). Running user code in that state makes the app's first syscall
    # swapgs the wrong way; the kernel-stack lookup then reads garbage from linear address 8,
    # the stack pointer becomes garbage, and even the page-fault frame cannot be pushed
    # => triple fault with zero output (not even the ISR entry probe can fire).
    # Swapgs only when GS_BASE != 0 (idempotent; same idiom as the handler's GS restore).
    __asm__(f"""cli
            mov rsp, {asminp(kern_rsp, ASM_DESCR.REG_ANY)}
            mov ecx, 0xC0000101
            rdmsr
            shl rdx, 32
            or rax, rdx
            test rax, rax
            jz 1f
            swapgs
            1:
            mov rdx, {asminp(user_rsp, ASM_DESCR.REG_ANY)}
            mov rcx, {asminp(entry, ASM_DESCR.REG_ANY)}
            mov r11, 0x202
            push 0x23
            push rdx
            push r11
            push 0x2B
            push rcx
            iretq""",
          op=[ASM_DESCR.CLOBBER_MEMORY, ASM_DESCR.CLOBBER_RAX,
              ASM_DESCR.CLOBBER_RCX, ASM_DESCR.CLOBBER_RDX,
              ASM_DESCR.CLOBBER_R11, ASM_DESCR.CLOBBER_RSI,
              ASM_DESCR.CLOBBER_RDI, ASM_DESCR.CLOBBER_R8,
              ASM_DESCR.CLOBBER_R9, ASM_DESCR.CLOBBER_R10,
              ASM_DESCR.CLOBBER_RBP])


# ============================================================
# 用户态线程（THREAD_CREATE）上下文
# 每个用户线程需要两套栈，否则同进程多线程会互相踩：
#   · 用户栈（Ring 3 执行用）：从 USER_STACK_VTOP 下方按固定间隔分配，映射 PTE_USER；
#   · 内核栈（Ring 3 → Ring 0 的 syscall/中断落栈点 = TSS.RSP0 / gs:[8]）：
#     必须按**线程**写入（由调度器切线程时设置），不能按进程 —— 否则一个线程在
#     syscall 里 yield 时，另一个线程的 syscall 会压掉它的栈帧，返回即跑飞。
# ============================================================
class UserThreadCtx:
    pid:                  int
    tid:                  int
    entry:                rawptr
    arg:                  rawptr
    # 入口 rsp（= 栈顶下一格；该格里放哨兵返回地址 0，见 alloc_user_thread）
    user_stack_entry_rsp: uint64
    # 用户栈映射上界（UnMap 用）
    user_stack_vtop:      uint64
    user_stack_phys:      rawptr
    kern_stack:           uint64

_user_ctx:        [MAX_USER_THREADS]UserThreadCtx
_user_ctx_used:   [MAX_USER_THREADS]int
_user_stack_next: int = 0

def alloc_user_thread(pid: int, entry: rawptr, arg: rawptr) -> int:
    """登记一个用户线程：分配内核栈 + 用户栈，并把用户栈页映射进该进程地址空间。

    必须在调用者进程上下文里调用（MapPage 作用于当前 pml4 = THREAD_CREATE 调用者的）。
    返回槽位号；失败返回 -1（不兜底，调用方显式处理失败）。
    """
    global _user_stack_next
    slot: int = -1
    i: int
    for i in range(MAX_USER_THREADS):
        if _user_ctx_used[i] == 0:
            slot = i
            break
    if slot < 0:
        serial.puts("[proc] no free user-thread slot\n")
        return -1
    ks: rawptr = mm.malloc(KERN_STACK_SIZE)
    if ks is None:
        serial.puts("[proc] user-thread kern stack alloc failed\n")
        return -1
    us: rawptr = mm.malloc(USER_STACK_SIZE)
    if us is None:
        serial.puts("[proc] user-thread user stack alloc failed\n")
        mm.free(ks)
        return -1
    us_size: uint64 = uint64(USER_STACK_SIZE)
    phys_bottom: uint64 = uint64(us) & ~uint64(0xFFF)
    vtop: uint64 = USER_STACK_VTOP - uint64(_user_stack_next + 1) * uint64(USER_THREAD_STACK_GAP)
    _user_stack_next = _user_stack_next + 1
    si: int
    for si in range(USER_STACK_PAGES):
        vaddr: uint64 = (vtop - us_size) + uint64(si) * uint64(0x1000)
        paddr: uint64 = phys_bottom + uint64(si) * uint64(0x1000)
        paging.MapPage(vaddr, paddr, paging.PTE_PRESENT | paging.PTE_WRITABLE | paging.PTE_USER)
    # 入口 rsp 指向栈顶下一格，并在该格写哨兵返回地址 0：用户线程函数若 return，
    # ret 会跳到 0 ⇒ Ring3 取指缺页（cr2=rip=0，日志一眼可辨）⇒ ISR 终止并回收资源。
    entry_rsp: uint64 = vtop - uint64(8)
    ret_slot: uint64* = uint64*(entry_rsp)
    *ret_slot = 0
    _user_ctx_used[slot] = 1
    _user_ctx[slot].pid = pid
    _user_ctx[slot].tid = -1
    _user_ctx[slot].entry = entry
    _user_ctx[slot].arg = arg
    _user_ctx[slot].user_stack_entry_rsp = entry_rsp
    _user_ctx[slot].user_stack_vtop = vtop
    _user_ctx[slot].user_stack_phys = us
    _user_ctx[slot].kern_stack = uint64(ks) + uint64(KERN_STACK_SIZE)
    return slot

def free_user_thread(slot: int):
    """回收用户线程资源（线程已不在用户态执行）。

    · 解映射用户栈 VA 并交还物理页；
    · 清槽位标记（允许复用）。
    内核栈**故意不在这里 free**：终止路径（异常 ISR / THREAD_EXIT）此刻正跑在该线程
    自己的内核栈上，free 后它会被别的线程分配走并覆盖，直到 _do_switch 切走为止。
    """
    if slot < 0 or slot >= MAX_USER_THREADS: return
    if _user_ctx_used[slot] == 0: return
    vtop: uint64 = _user_ctx[slot].user_stack_vtop
    si: int
    for si in range(USER_STACK_PAGES):
        vaddr: uint64 = (vtop - uint64(USER_STACK_SIZE)) + uint64(si) * uint64(0x1000)
        paging.UnMapPage(vaddr)
    us: rawptr = _user_ctx[slot].user_stack_phys
    if us is not None:
        mm.free(us)
    _user_ctx[slot].user_stack_phys = None
    _user_ctx_used[slot] = 0

def find_user_slot_by_tid(tid: int) -> int:
    i: int
    for i in range(MAX_USER_THREADS):
        if _user_ctx_used[i] != 0 and _user_ctx[i].tid == tid:
            return i
    return -1

def user_thread_ctx(slot: int) -> UserThreadCtx*:
    if slot < 0 or slot >= MAX_USER_THREADS: return None
    if _user_ctx_used[slot] == 0: return None
    return &_user_ctx[slot]