import drivers.serial.uart.serial as serial
import services.desktop as desktop
import drivers.fs.fat32.fat32 as fat32
import drivers.fs.fat32.fat32_types as fat32_types
import sched.sched as sched
import sched.process as process
import viperlib
import string
import mm.mm as mm
import paging.paging as paging
import execrunner.elf as elf
import t, c

CREATE_WINDOW: t.CDefine = 1
DRAW_TEXT:     t.CDefine = 2
EXIT:         t.CDefine = 3
OPEN:         t.CDefine = 4
CLOSE:        t.CDefine = 5
READ:         t.CDefine = 6
WRITE:        t.CDefine = 7
SEEK:         t.CDefine = 8
TELL:         t.CDefine = 9
SIZE:         t.CDefine = 10
MKDIR:        t.CDefine = 11
REMOVE:       t.CDefine = 12
OPENDIR:      t.CDefine = 13
READDIR:      t.CDefine = 14
CLOSEDIR:     t.CDefine = 15
STAT:         t.CDefine = 16
WAIT_WINDOW:  t.CDefine = 17
SERIAL_PUTS: t.CDefine = 18
LOAD_SO:     t.CDefine = 19
SO_CALL1:    t.CDefine = 20
GET_WINBUF:  t.CDefine = 21
WIN_FLUSH:   t.CDefine = 22
POLL_EVENT:  t.CDefine = 23
SET_CURSOR: t.CDefine = 24
YIELD:      t.CDefine = 25
SET_TITLE:  t.CDefine = 26
SET_GEOMETRY: t.CDefine = 27
SPAWN:       t.CDefine = 28
MMAP:        t.CDefine = 29
MUNMAP:      t.CDefine = 30
SET_RESIZABLE: t.CDefine = 31
GET_WIN_SIZE: t.CDefine = 32
THREAD_CREATE: t.CDefine = 33

IA32_EFER:  t.CDefine = 0xC0000080
IA32_STAR:  t.CDefine = 0xC0000081
IA32_LSTAR: t.CDefine = 0xC0000082
IA32_FMASK: t.CDefine = 0xC0000084

def wrmsr(msr: t.CUInt32T, val: t.CUInt64T):
    lo: t.CUInt32T = t.CUInt32T(val & 0xFFFFFFFF)
    hi: t.CUInt32T = t.CUInt32T(val >> 32)
    c.Asm(f"""mov ecx, {c.AsmInp(msr, t.ASM_DESCR.REG_ANY)}
              mov eax, {c.AsmInp(lo, t.ASM_DESCR.REG_ANY)}
              mov edx, {c.AsmInp(hi, t.ASM_DESCR.REG_ANY)}
              wrmsr""",
          op=[t.ASM_DESCR.CLOBBER_RCX, t.ASM_DESCR.CLOBBER_RAX, t.ASM_DESCR.CLOBBER_RDX])

def rdmsr(msr: t.CUInt32T) -> t.CUInt64T:
    lo: t.CUInt32T
    hi: t.CUInt32T
    c.Asm(f"""mov ecx, {c.AsmInp(msr, t.ASM_DESCR.REG_ANY)}
              rdmsr""",
          out=[c.AsmOut(lo, t.ASM_DESCR.OUTPUT_REG), c.AsmOut(hi, t.ASM_DESCR.OUTPUT_REG)],
          op=[t.ASM_DESCR.CLOBBER_RCX, t.ASM_DESCR.CLOBBER_RAX, t.ASM_DESCR.CLOBBER_RDX])
    return t.CUInt64T(lo) | (t.CUInt64T(hi) << 32)

def _current_proc() -> process.Process | t.CPtr:
    return process.ProcessManager.current()

def _fd_alloc(p: process.Process | t.CPtr) -> t.CInt:
    for i in range(process.PROC_MAX_FDS):
        if not p.fd_used[i]:
            p.fd_used[i] = 1
            return i
    return -1

def _fd_free(p: process.Process | t.CPtr, fd: t.CInt):
    if fd >= 0 and fd < process.PROC_MAX_FDS:
        p.fd_used[fd] = 0
        p.fd_table[fd] = t.CVoid(0, t.CPtr)

def _dir_alloc(p: process.Process | t.CPtr) -> t.CInt:
    for i in range(process.PROC_MAX_DIRS):
        if not p.dir_used[i]:
            p.dir_used[i] = 1
            return i
    return -1

def _dir_free(p: process.Process | t.CPtr, dd: t.CInt):
    if dd >= 0 and dd < process.PROC_MAX_DIRS:
        p.dir_used[dd] = 0
        p.dir_table[dd] = t.CVoid(0, t.CPtr)

def CREATE_WINDOW(x: t.CInt, y: t.CInt, w: t.CInt, h: t.CInt, title: t.CConst | t.CChar | t.CPtr, style: t.CInt) -> t.CInt:
    desktop.set_app_pid(process.current())
    return desktop.create(x, y, w, h, title, style)

def DRAW_TEXT(win_id: t.CInt, x: t.CInt, y: t.CInt, text: t.CConst | t.CChar | t.CPtr, color: t.CUInt32T) -> t.CInt:
    return desktop.draw_text(win_id, x, y, text, color)

def EXIT(code: t.CInt) -> t.CInt:
    serial.puts("[syscall] exit called\n")
    desktop.destroy_by_pid(process.current())
    return 0

def WAIT_WINDOW(win_id: t.CInt) -> t.CInt:
    while desktop.is_active(win_id):
        sched.Scheduler.try_reschedule()
    return 0

def SERIAL_PUTS(str_ptr: t.CVoid | t.CPtr) -> t.CInt:
    s: t.CConst | t.CChar | t.CPtr = t.CVoid(str_ptr, t.CPtr)
    serial.puts(s)
    return 0

_loaded_so_base: t.CVoid | t.CPtr = None
_loaded_so_path: t.CArray[t.CChar, 64]

def LOAD_SO(path_ptr: t.CVoid | t.CPtr) -> t.CUInt64T:
    global _loaded_so_base
    path: t.CConst | t.CChar | t.CPtr = t.CVoid(path_ptr, t.CPtr)
    serial.puts("[so] loading: ")
    serial.puts(path)
    serial.puts("\n")
    base: t.CVoid | t.CPtr = elf.load_so(path)
    if base is not None:
        _loaded_so_base = base
        for i in range(63):
            ch: t.CChar = c.Deref(path + i)
            _loaded_so_path[i] = ch
            if ch == 0: break
        _loaded_so_path[63] = 0
        return t.CUInt64T(base)
    return 0

def SO_CALL1(sym_name_ptr: t.CVoid | t.CPtr, arg1: t.CVoid | t.CPtr) -> t.CInt:
    if _loaded_so_base is None: return -1
    sym_name: t.CConst | t.CChar | t.CPtr = t.CVoid(sym_name_ptr, t.CPtr)
    fn: t.CVoid | t.CPtr = elf.get_so_symbol(_loaded_so_base, sym_name, c.Addr(_loaded_so_path[0]))
    if fn is None:
        serial.puts("[so] symbol not found\n")
        return -2
    c.Asm(f"""mov rdi, {c.AsmInp(arg1, t.ASM_DESCR.REG_ANY)}
            call {c.AsmInp(fn, t.ASM_DESCR.REG_ANY)}""",
          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_RDI, t.ASM_DESCR.CLOBBER_RSI,
              t.ASM_DESCR.CLOBBER_R8, t.ASM_DESCR.CLOBBER_R9,
              t.ASM_DESCR.CLOBBER_R10, t.ASM_DESCR.CLOBBER_R11])
    serial.puts("[so] call returned\n")
    return 0

def GET_WINBUF(win_id: t.CInt) -> t.CUInt64T:
    buf: t.CVoid | t.CPtr = desktop.get_fb(win_id)
    if buf is None:
        return 0
    return t.CUInt64T(buf)

def WIN_FLUSH(win_id: t.CInt) -> t.CInt:
    desktop.flush(win_id)
    return 0

def SET_CURSOR(win_id: t.CInt, cursor_t: t.CInt) -> t.CInt:
    desktop.set_cursor_type(win_id, cursor_t)
    return 0

def OPEN(path_ptr: t.CVoid | t.CPtr, mode: t.CUInt32T) -> t.CInt:
    p: process.Process | t.CPtr = _current_proc()
    path: t.CConst | t.CChar | t.CPtr = t.CVoid(path_ptr, t.CPtr)
    fat32.lock()
    fp: t.CVoid | t.CPtr = fat32.open(path, t.CUInt8T(mode))
    if not fp:
        fat32.unlock()
        return -1
    fd: t.CInt = _fd_alloc(p)
    if fd < 0:
        fat32.close(fp)
        fat32.unlock()
        return -1
    p.fd_table[fd] = fp
    fat32.unlock()
    return fd

def CLOSE(fd: t.CInt) -> t.CInt:
    p: process.Process | t.CPtr = _current_proc()
    if fd < 0 or fd >= process.PROC_MAX_FDS:
        return t.CInt(fat32_types.FRESULT.FR_INVALID_OBJECT)
    fp: t.CVoid | t.CPtr = p.fd_table[fd]
    if not fp:
        return t.CInt(fat32_types.FRESULT.FR_INVALID_OBJECT)
    fat32.lock()
    res: t.CInt = t.CInt(fat32.close(fp))
    fat32.unlock()
    _fd_free(p, fd)
    return res

def READ(fd: t.CInt, buf_ptr: t.CVoid | t.CPtr, count: t.CUInt32T) -> t.CInt:
    p: process.Process | t.CPtr = _current_proc()
    if fd < 0 or fd >= process.PROC_MAX_FDS:
        return -1
    fp: t.CVoid | t.CPtr = p.fd_table[fd]
    if not fp:
        return -1
    fat32.lock()
    br: t.CUInt32T = fat32.read(fp, buf_ptr, count)
    fat32.unlock()
    return t.CInt(br)

def WRITE(fd: t.CInt, buf_ptr: t.CVoid | t.CPtr, count: t.CUInt32T) -> t.CInt:
    p: process.Process | t.CPtr = _current_proc()
    if fd < 0 or fd >= process.PROC_MAX_FDS:
        return -1
    fp: t.CVoid | t.CPtr = p.fd_table[fd]
    if not fp:
        return -1
    fat32.lock()
    bw: t.CUInt32T = fat32.write(fp, buf_ptr, count)
    fat32.unlock()
    return t.CInt(bw)

def SEEK(fd: t.CInt, offset: t.CUInt32T) -> t.CInt:
    p: process.Process | t.CPtr = _current_proc()
    if fd < 0 or fd >= process.PROC_MAX_FDS:
        return t.CInt(fat32_types.FRESULT.FR_INVALID_OBJECT)
    fp: t.CVoid | t.CPtr = p.fd_table[fd]
    if not fp:
        return t.CInt(fat32_types.FRESULT.FR_INVALID_OBJECT)
    fat32.lock()
    r: t.CInt = t.CInt(fat32.seek(fp, offset))
    fat32.unlock()
    return r

def TELL(fd: t.CInt) -> t.CUInt32T:
    p: process.Process | t.CPtr = _current_proc()
    if fd < 0 or fd >= process.PROC_MAX_FDS:
        return 0
    fp: t.CVoid | t.CPtr = p.fd_table[fd]
    if not fp:
        return 0
    fat32.lock()
    pos: t.CUInt32T = fat32.tell(fp)
    fat32.unlock()
    return pos

def fsize(fd: t.CInt) -> t.CUInt32T:
    p: process.Process | t.CPtr = _current_proc()
    if fd < 0 or fd >= process.PROC_MAX_FDS:
        return 0
    fp: t.CVoid | t.CPtr = p.fd_table[fd]
    if not fp:
        return 0
    fat32.lock()
    sz: t.CUInt32T = fat32.size(fp)
    fat32.unlock()
    return sz

def MKDIR(path_ptr: t.CVoid | t.CPtr) -> t.CInt:
    path: t.CConst | t.CChar | t.CPtr = t.CVoid(path_ptr, t.CPtr)
    fat32.lock()
    r: t.CInt = t.CInt(fat32.mkdir(path))
    fat32.unlock()
    return r

def REMOVE(path_ptr: t.CVoid | t.CPtr) -> t.CInt:
    path: t.CConst | t.CChar | t.CPtr = t.CVoid(path_ptr, t.CPtr)
    fat32.lock()
    r: t.CInt = t.CInt(fat32.remove(path))
    fat32.unlock()
    return r

def OPENdir(path_ptr: t.CVoid | t.CPtr) -> t.CInt:
    p: process.Process | t.CPtr = _current_proc()
    path: t.CConst | t.CChar | t.CPtr = t.CVoid(path_ptr, t.CPtr)
    dd: t.CInt = _dir_alloc(p)
    if dd < 0:
        return -1
    dp: fat32_types.dirobj | t.CPtr = c.Addr(p.dir_pool[dd])
    fat32.lock()
    res: t.CInt = t.CInt(fat32.opendir(path, dp))
    fat32.unlock()
    if res != 0:
        _dir_free(p, dd)
        return -1
    p.dir_table[dd] = t.CVoid(t.CUInt64T(dp), t.CPtr)
    return dd

def READdir(dd: t.CInt, info_ptr: t.CVoid | t.CPtr) -> t.CInt:
    p: process.Process | t.CPtr = _current_proc()
    if dd < 0 or dd >= process.PROC_MAX_DIRS:
        return t.CInt(fat32_types.FRESULT.FR_INVALID_OBJECT)
    dp: t.CVoid | t.CPtr = p.dir_table[dd]
    if not dp:
        return t.CInt(fat32_types.FRESULT.FR_INVALID_OBJECT)
    info: fat32_types.fileinfo | t.CPtr = info_ptr
    fat32.lock()
    r: t.CInt = t.CInt(fat32.readdir(dp, info))
    fat32.unlock()
    return r

def CLOSEdir(dd: t.CInt) -> t.CInt:
    p: process.Process | t.CPtr = _current_proc()
    if dd < 0 or dd >= process.PROC_MAX_DIRS:
        return t.CInt(fat32_types.FRESULT.FR_INVALID_OBJECT)
    dp: t.CVoid | t.CPtr = p.dir_table[dd]
    if not dp:
        return t.CInt(fat32_types.FRESULT.FR_INVALID_OBJECT)
    fat32.lock()
    fat32.closedir(dp)
    fat32.unlock()
    _dir_free(p, dd)
    return 0

def STAT(path_ptr: t.CVoid | t.CPtr, info_ptr: t.CVoid | t.CPtr) -> t.CInt:
    path: t.CConst | t.CChar | t.CPtr = t.CVoid(path_ptr, t.CPtr)
    info: fat32_types.fileinfo | t.CPtr = info_ptr
    fat32.lock()
    r: t.CInt = t.CInt(fat32.stat(path, info))
    fat32.unlock()
    return r

def SPAWN(path_ptr: t.CVoid | t.CPtr, blocking: t.CInt) -> t.CInt:
    path: t.CConst | t.CChar | t.CPtr = t.CVoid(path_ptr, t.CPtr)
    pid: t.CInt = elf.spawn_elf(path)
    if pid < 0:
        return pid
    if blocking:
        p: process.Process | t.CPtr = process.ProcessManager.get_process(pid)
        if p:
            while p.state != process.PROC_DONE:
                sched.Scheduler.try_reschedule()
    return pid

USER_HEAP_BASE: t.CDefine = 0x20000000

def MMAP(size: t.CUInt64T) -> t.CUInt64T:
    p: process.Process | t.CPtr = _current_proc()
    if size == 0: return 0
    aligned_size: t.CUInt64T = (size + 4095) & ~t.CUInt64T(4095)
    if p.heap_start == 0:
        p.heap_start = USER_HEAP_BASE
    heap_end: t.CUInt64T = p.heap_start + p.heap_size
    new_end: t.CUInt64T = heap_end + aligned_size
    for va in range(heap_end, new_end, 4096):
        phys_page: t.CVoid | t.CPtr = mm.malloc(4096)
        if phys_page is None: return 0
        paging.MapPage(va, t.CUInt64T(phys_page), paging.PTE_PRESENT | paging.PTE_WRITABLE | paging.PTE_USER)
    p.heap_size = p.heap_size + aligned_size
    return heap_end

def MUNMAP(ptr: t.CUInt64T, size: t.CUInt64T) -> t.CInt:
    return 0

def syscall_entry() -> t.CExtern | t.CVoid | t.State: pass

def syscall_handler(num: t.CUInt64T, arg1: t.CUInt64T, arg2: t.CUInt64T, arg3: t.CUInt64T, arg4: t.CUInt64T, arg5: t.CUInt64T, arg6: t.CUInt64T) -> t.CExport | t.CUInt64T:
    n: t.CInt = t.CInt(num)
    if n == CREATE_WINDOW:
        title_ptr: t.CVoid | t.CPtr = t.CVoid(arg5, t.CPtr)
        return t.CUInt64T(CREATE_WINDOW(t.CInt(arg1), t.CInt(arg2), t.CInt(arg3), t.CInt(arg4), title_ptr, t.CInt(arg6)))
    elif n == DRAW_TEXT:
        text_ptr: t.CVoid | t.CPtr = t.CVoid(arg4, t.CPtr)
        return t.CUInt64T(DRAW_TEXT(t.CInt(arg1), t.CInt(arg2), t.CInt(arg3), text_ptr, t.CUInt32T(arg5)))
    elif n == EXIT:
        return t.CUInt64T(EXIT(t.CInt(arg1)))
    elif n == OPEN:
        path_ptr: t.CVoid | t.CPtr = t.CVoid(arg1, t.CPtr)
        return t.CUInt64T(OPEN(path_ptr, t.CUInt32T(arg2)))
    elif n == CLOSE:
        return t.CUInt64T(CLOSE(t.CInt(arg1)))
    elif n == READ:
        buf_ptr: t.CVoid | t.CPtr = t.CVoid(arg2, t.CPtr)
        return t.CUInt64T(READ(t.CInt(arg1), buf_ptr, t.CUInt32T(arg3)))
    elif n == WRITE:
        buf_ptr: t.CVoid | t.CPtr = t.CVoid(arg2, t.CPtr)
        return t.CUInt64T(WRITE(t.CInt(arg1), buf_ptr, t.CUInt32T(arg3)))
    elif n == SEEK:
        return t.CUInt64T(SEEK(t.CInt(arg1), t.CUInt32T(arg2)))
    elif n == TELL:
        return t.CUInt64T(TELL(t.CInt(arg1)))
    elif n == SIZE:
        return t.CUInt64T(fsize(t.CInt(arg1)))
    elif n == MKDIR:
        path_ptr: t.CVoid | t.CPtr = t.CVoid(arg1, t.CPtr)
        return t.CUInt64T(MKDIR(path_ptr))
    elif n == REMOVE:
        path_ptr: t.CVoid | t.CPtr = t.CVoid(arg1, t.CPtr)
        return t.CUInt64T(REMOVE(path_ptr))
    elif n == OPENDIR:
        path_ptr: t.CVoid | t.CPtr = t.CVoid(arg1, t.CPtr)
        return t.CUInt64T(OPENdir(path_ptr))
    elif n == READDIR:
        info_ptr: t.CVoid | t.CPtr = t.CVoid(arg2, t.CPtr)
        return t.CUInt64T(READdir(t.CInt(arg1), info_ptr))
    elif n == CLOSEDIR:
        return t.CUInt64T(CLOSEdir(t.CInt(arg1)))
    elif n == STAT:
        path_ptr: t.CVoid | t.CPtr = t.CVoid(arg1, t.CPtr)
        info_ptr: t.CVoid | t.CPtr = t.CVoid(arg2, t.CPtr)
        return t.CUInt64T(STAT(path_ptr, info_ptr))
    elif n == WAIT_WINDOW:
        if t.CInt(arg2) == 1:
            return t.CUInt64T(desktop.is_active(t.CInt(arg1)))
        elif t.CInt(arg2) == 2:
            sched.Scheduler.try_reschedule()
            return 0
        return t.CUInt64T(WAIT_WINDOW(t.CInt(arg1)))
    elif n == SERIAL_PUTS:
        str_ptr_sp: t.CVoid | t.CPtr = t.CVoid(arg1, t.CPtr)
        return t.CUInt64T(SERIAL_PUTS(str_ptr_sp))
    elif n == LOAD_SO:
        path_ptr_so: t.CVoid | t.CPtr = t.CVoid(arg1, t.CPtr)
        return LOAD_SO(path_ptr_so)
    elif n == SO_CALL1:
        sym_ptr: t.CVoid | t.CPtr = t.CVoid(arg1, t.CPtr)
        return t.CUInt64T(SO_CALL1(sym_ptr, arg2))
    elif n == GET_WINBUF:
        return GET_WINBUF(t.CInt(arg1))
    elif n == WIN_FLUSH:
        return t.CUInt64T(WIN_FLUSH(t.CInt(arg1)))
    elif n == POLL_EVENT:
        buf_ptr_ev: t.CVoid | t.CPtr = t.CVoid(arg1, t.CPtr)
        caller_pid_ev: t.CInt = process.current()
        return t.CUInt64T(desktop.poll_event(buf_ptr_ev, caller_pid_ev))
    elif n == SET_CURSOR:
        desktop.set_cursor_type(t.CInt(arg1), t.CInt(arg2))
        return 0
    elif n == YIELD:
        sched.Scheduler.try_reschedule()
        return 0
    elif n == SET_TITLE:
        title_ptr_st: t.CVoid | t.CPtr = t.CVoid(arg2, t.CPtr)
        return t.CUInt64T(desktop.set_title(t.CInt(arg1), title_ptr_st))
    elif n == SET_GEOMETRY:
        return t.CUInt64T(desktop.set_geometry(t.CInt(arg1), t.CInt(arg2), t.CInt(arg3), t.CInt(arg4), t.CInt(arg5)))
    elif n == SET_RESIZABLE:
        return t.CUInt64T(desktop.set_resizable(t.CInt(arg1), t.CInt(arg2), t.CInt(arg3)))
    elif n == GET_WIN_SIZE:
        return t.CUInt64T(desktop.get_size(t.CInt(arg1)))
    elif n == SPAWN:
        path_ptr_sp: t.CVoid | t.CPtr = t.CVoid(arg1, t.CPtr)
        return t.CUInt64T(SPAWN(path_ptr_sp, t.CInt(arg2)))
    elif n == MMAP:
        return MMAP(arg1)
    elif n == MUNMAP:
        return t.CUInt64T(MUNMAP(arg1, arg2))
    elif n == THREAD_CREATE:
        func_ptr: t.CVoid | t.CPtr = t.CVoid(arg1, t.CPtr)
        arg_ptr: t.CVoid | t.CPtr = t.CVoid(arg2, t.CPtr)
        cur_pid: t.CInt = process.current()
        th: sched.Thread | t.CPtr = sched.Scheduler.create_thread(func_ptr, arg_ptr, cur_pid)
        if th is None: return t.CUInt64T(0xFFFFFFFFFFFFFFFF)
        p: process.Process | t.CPtr = process.ProcessManager.get_process(cur_pid)
        if p:
            p.addThread(th.tid)
        return t.CUInt64T(th.tid)
    return 0

def init():
    efer: t.CUInt64T = rdmsr(IA32_EFER)
    efer = efer | 1
    wrmsr(IA32_EFER, efer)
    star_val: t.CUInt64T = (t.CUInt64T(0x08) << 32) | (t.CUInt64T(0x1B) << 48)
    wrmsr(IA32_STAR, star_val)
    wrmsr(IA32_LSTAR, t.CUInt64T(syscall_entry))
    wrmsr(IA32_FMASK, 0x200)
    serial.puts("[syscall] MSR init done\n")