import mm
import drivers.serial.uart.serial as serial
import asm
import t, c

CORO_READY: t.CDefine = 0
CORO_RUNNING: t.CDefine = 1
CORO_YIELDED: t.CDefine = 2
CORO_DONE: t.CDefine = 3
MAX_COROUTINES: t.CDefine = 32
CORO_STACK_SIZE: t.CDefine = 8192


@t.Object
class Coroutine:
    rsp:        t.CUInt64T
    stack:      t.CVoid | t.CPtr
    stack_size: t.CUInt64T
    func:       t.CVoid | t.CPtr
    arg:        t.CVoid | t.CPtr
    state:      t.CInt
    cid:        t.CInt
    caller_rsp: t.CUInt64T
    ret_val:    t.CInt

    def resume(self) -> t.CInt:
        global _coro_old_rsp, _coro_new_rsp
        m: CoroutineManager | t.CPtr = _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: t.CInt = m.current_cid
        m.current_cid = self.cid
        self.state = CORO_RUNNING
        _coro_old_rsp = c.Addr(m.coroutines[prev_cid].rsp)
        _coro_new_rsp = c.Addr(self.rsp)
        self.caller_rsp = t.CUInt64T(_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 | t.CPtr = _coro_mgr_ptr
        cid: t.CInt = m.current_cid
        if cid < 0: return
        m.coroutines[cid].state = CORO_YIELDED
        prev_cid: t.CInt = -1
        i: t.CInt
        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 = c.Addr(m.coroutines[cid].rsp)
        _coro_new_rsp = c.Addr(m.coroutines[prev_cid].rsp)
        _coro_switch()

    @property
    def is_done(self) -> t.CInt:
        if self.state == CORO_DONE:
            return 1
        return 0

    @property
    def is_yielded(self) -> t.CInt:
        if self.state == CORO_YIELDED:
            return 1
        return 0

    @property
    def retval(self) -> t.CInt:
        return self.ret_val


@t.Object
class CoroutineManager:
    coroutines: t.CArray[Coroutine, MAX_COROUTINES]
    current_cid: t.CInt
    count: t.CInt

    def __init__():
        global _coro_mgr_ptr
        _coro_mgr_ptr = c.Addr(_coro_mgr)
        m: CoroutineManager | t.CPtr = _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: t.CVoid | t.CPtr = c.Addr(m.coroutines[0].rsp)
        c.Asm(f"""mov qword ptr [{c.AsmInp(main_rsp_ptr, t.ASM_DESCR.REG_ANY)}], rsp""",
            op=[t.ASM_DESCR.CLOBBER_MEMORY])

    @staticmethod
    def current() -> Coroutine | t.CPtr:
        m: CoroutineManager | t.CPtr = _coro_mgr_ptr
        return c.Addr(m.coroutines[m.current_cid])

    @staticmethod
    def get_coroutine(cid: t.CInt) -> Coroutine | t.CPtr:
        m: CoroutineManager | t.CPtr = _coro_mgr_ptr
        if cid < 0 or cid >= m.count: return None
        return c.Addr(m.coroutines[cid])

    @staticmethod
    def create(func: t.CVoid | t.CPtr, arg: t.CVoid | t.CPtr) -> Coroutine | t.CPtr:
        m: CoroutineManager | t.CPtr = _coro_mgr_ptr
        if m.count >= MAX_COROUTINES: return None
        cid: t.CInt = m.count
        m.count += 1
        co: Coroutine | t.CPtr = c.Addr(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: t.CVoid | t.CPtr = mm.malloc(CORO_STACK_SIZE)
        if not stack_ptr: return None
        co.stack = stack_ptr
        co.stack_size = CORO_STACK_SIZE
        stack_top: t.CUInt64T = t.CUInt64T(stack_ptr) + CORO_STACK_SIZE - 16
        stack_top = stack_top & ~t.CUInt64T(0xF)
        entry_addr: t.CUInt64T = t.CUInt64T(c.Addr(_coro_entry))
        sp: t.CPtr = t.CPtr(stack_top)
        c.Asm(f"""mov rax, {c.AsmInp(sp, t.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], {c.AsmInp(entry_addr, t.ASM_DESCR.REG_ANY)}""",
            op=[t.ASM_DESCR.CLOBBER_MEMORY, t.ASM_DESCR.CLOBBER_RAX])
        stack_top = stack_top - 56
        co.rsp = stack_top
        return co


_coro_mgr: CoroutineManager = CoroutineManager()
_coro_mgr_ptr: CoroutineManager | t.CPtr = c.Addr(_coro_mgr)
_coro_old_rsp: t.CVoid | t.CPtr
_coro_new_rsp: t.CVoid | t.CPtr

@c.Attribute(t.attr.naked)
def _coro_switch():
    c.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=[t.ASM_DESCR.CLOBBER_MEMORY])

@c.Attribute(t.attr.naked)
def _coro_entry():
    c.Asm(f"""sub rsp, 8
            mov rax, qword ptr [rip + _coro_mgr_ptr]
            mov ecx, dword ptr [rax + 1056]
            movsxd rcx, ecx
            shl rcx, 5
            mov rdi, qword ptr [rax + rcx + 24]
            mov rax, qword ptr [rax + rcx + 16]
            call rax
            mov rax, qword ptr [rip + _coro_mgr_ptr]
            mov ecx, dword ptr [rax + 1056]
            movsxd rcx, ecx
            shl rcx, 5
            mov dword ptr [rax + rcx + 32], 3
            mov dword ptr [rax + rcx + 48], eax
            add rsp, 8""",
          op=[t.ASM_DESCR.CLOBBER_MEMORY, t.ASM_DESCR.CLOBBER_RAX,
              t.ASM_DESCR.CLOBBER_RCX, t.ASM_DESCR.CLOBBER_RDX,
              t.ASM_DESCR.CLOBBER_R8, t.ASM_DESCR.CLOBBER_R9,
              t.ASM_DESCR.CLOBBER_R10, t.ASM_DESCR.CLOBBER_R11])
    Coroutine._yield()
    while True:
        asm.sti()
        asm.hlt()