import mm.mm as mm
import asm


define CORO_READY      = 0
define CORO_RUNNING    = 1
define CORO_YIELDED    = 2
define CORO_DONE       = 3
define MAX_COROUTINES  = 32
define CORO_STACK_SIZE = 8192


class Coroutine:
    rsp:        uint64
    stack:      rawptr
    stack_size: uint64
    func:       rawptr
    arg:        rawptr
    state:      int
    cid:        int
    caller_rsp: uint64
    ret_val:    int

    def resume(self) -> int:
        global _coro_old_rsp, _coro_new_rsp
        m: CoroutineManager* = _coro_mgr_ptr
        if self.state == CORO_DONE: return self.ret_val
        if self.state != CORO_READY and self.state != CORO_YIELDED: return -1
        prev_cid: int = m.current_cid
        m.current_cid = self.cid
        self.state = CORO_RUNNING
        _coro_old_rsp = &m.coroutines[prev_cid].rsp
        _coro_new_rsp = &self.rsp
        self.caller_rsp = uint64(_coro_old_rsp)
        _coro_switch()
        if self.state == CORO_DONE:
            return self.ret_val
        return 0

    @staticmethod
    def _yield():
        global _coro_old_rsp, _coro_new_rsp
        m: CoroutineManager* = _coro_mgr_ptr
        cid: int = m.current_cid
        if cid < 0: return
        m.coroutines[cid].state = CORO_YIELDED
        prev_cid: int = -1
        for i in range(m.count):
            if m.coroutines[i].state == CORO_RUNNING and m.coroutines[i].caller_rsp != 0:
                prev_cid = i
                break
        if prev_cid < 0: return
        m.current_cid = prev_cid
        _coro_old_rsp = &m.coroutines[cid].rsp
        _coro_new_rsp = &m.coroutines[prev_cid].rsp
        _coro_switch()

    @property
    def is_done(self) -> bool: return self.state == CORO_DONE
    @property
    def is_yielded(self) -> bool: return self.state == CORO_YIELDED
    @property
    def retval(self) -> int: return self.ret_val


class CoroutineManager:
    coroutines: [MAX_COROUTINES]Coroutine
    current_cid: int
    count: int

    def __init__(self):
        global _coro_mgr_ptr
        _coro_mgr_ptr = &_coro_mgr
        m: CoroutineManager* = _coro_mgr_ptr
        m.current_cid = 0
        m.count = 1
        m.coroutines[0].state = CORO_RUNNING
        m.coroutines[0].cid = 0
        m.coroutines[0].caller_rsp = 0
        m.coroutines[0].ret_val = 0
        m.coroutines[0].rsp = 0
        m.coroutines[0].stack = None
        m.coroutines[0].stack_size = 0
        m.coroutines[0].func = None
        m.coroutines[0].arg = None
        main_rsp_ptr: rawptr = &m.coroutines[0].rsp
        __asm__(f"""mov qword ptr [{asminp(main_rsp_ptr, ASM_DESCR.REG_ANY)}], rsp""",
            op=[ASM_DESCR.CLOBBER_MEMORY])

    @staticmethod
    def current() -> Coroutine*:
        m: CoroutineManager* = _coro_mgr_ptr
        return &m.coroutines[m.current_cid]

    @staticmethod
    def get_coroutine(cid: int) -> Coroutine*:
        m: CoroutineManager* = _coro_mgr_ptr
        if cid < 0 or cid >= m.count: return None
        return &m.coroutines[cid]

    @staticmethod
    def create(func: rawptr, arg: rawptr) -> Coroutine*:
        m: CoroutineManager* = _coro_mgr_ptr
        if m.count >= MAX_COROUTINES: return None
        cid: int = m.count
        m.count += 1
        co: Coroutine* = &m.coroutines[cid]
        co.state = CORO_READY
        co.cid = cid
        co.func = func
        co.arg = arg
        co.caller_rsp = 0
        co.ret_val = 0
        stack_ptr: rawptr = mm.malloc(CORO_STACK_SIZE)
        if not stack_ptr: return None
        co.stack = stack_ptr
        co.stack_size = CORO_STACK_SIZE
        stack_top: uint64 = uint64(stack_ptr) + CORO_STACK_SIZE - 16
        stack_top = stack_top & ~uint64(0xF)
        entry_addr: uint64 = uint64(&_coro_entry)
        sp: rawptr = rawptr(stack_top)
        __asm__(f"""mov rax, {asminp(sp, ASM_DESCR.REG_ANY)}
                mov qword ptr [rax - 56], 0
                mov qword ptr [rax - 48], 0
                mov qword ptr [rax - 40], 0
                mov qword ptr [rax - 32], 0
                mov qword ptr [rax - 24], 0
                mov qword ptr [rax - 16], 0
                mov qword ptr [rax - 8], 0
                mov qword ptr [rax], {asminp(entry_addr, ASM_DESCR.REG_ANY)}""",
            op=[ASM_DESCR.CLOBBER_MEMORY, ASM_DESCR.CLOBBER_RAX])
        stack_top = stack_top - 56
        co.rsp = stack_top
        return co


_coro_mgr: CoroutineManager = CoroutineManager()
export _coro_mgr_ptr: CoroutineManager* = &_coro_mgr
export _coro_old_rsp: rawptr
export _coro_new_rsp: rawptr

@__attribute__(naked)
def _coro_switch():
    __asm__(f"""push rax
            push rbx
            push rbp
            push r12
            push r13
            push r14
            push r15
            mov rax, qword ptr [rip + _coro_old_rsp]
            mov qword ptr [rax], rsp
            mov rax, qword ptr [rip + _coro_new_rsp]
            mov rsp, qword ptr [rax]
            pop r15
            pop r14
            pop r13
            pop r12
            pop rbp
            pop rbx
            pop rax
            ret""",
          op=[ASM_DESCR.CLOBBER_MEMORY])

@__attribute__(naked)
def _coro_entry():
    # 结构体布局（以 stub.ll 为权威）：
    #   Coroutine       = {rsp@0, stack@8, stack_size@16, func@24, arg@32,
    #                      state@40, cid@44, caller_rsp@48, ret_val@56}， sizeof=64
    #   CoroutineManager= {coroutines[32]@0 (32*64=2048), current_cid@2048, count@2052}
    # 取 func/arg 入口，写 state=CORO_DONE(3)、ret_val。
    __asm__(f"""sub rsp, 8
            mov rax, qword ptr [rip + _coro_mgr_ptr]
            mov ecx, dword ptr [rax + 2048]
            movsxd rcx, ecx
            shl rcx, 6
            mov rdi, qword ptr [rax + rcx + 32]
            mov rax, qword ptr [rax + rcx + 24]
            call rax
            mov rax, qword ptr [rip + _coro_mgr_ptr]
            mov ecx, dword ptr [rax + 2048]
            movsxd rcx, ecx
            shl rcx, 6
            mov dword ptr [rax + rcx + 40], 3
            mov dword ptr [rax + rcx + 56], eax
            add rsp, 8""",
          # `call rax` 破坏全部 caller-saved 寄存器（x64 SysV）：除 RAX/RCX/RDX/
          # R8-R11 外，RSI 与 RDI 也必须声明，否则编译器可能假设其值存活（N9）。
          # RSP 由 `sub rsp,8` / `add rsp,8` 配对抵消，净变化为 0，无需声明。
          op=[ASM_DESCR.CLOBBER_MEMORY, ASM_DESCR.CLOBBER_RAX,
              ASM_DESCR.CLOBBER_RCX, ASM_DESCR.CLOBBER_RDX,
              ASM_DESCR.CLOBBER_RSI, ASM_DESCR.CLOBBER_RDI,
              ASM_DESCR.CLOBBER_R8, ASM_DESCR.CLOBBER_R9,
              ASM_DESCR.CLOBBER_R10, ASM_DESCR.CLOBBER_R11])
    Coroutine._yield()
    while True:
        asm.sti()
        asm.hlt()