
header "sys" as sys

use mem
use gc
use type
use time
use thread

global coro_run_first : ?Coro
global coro_run_last : ?Coro
global current_coro : ?Coro

global coro_poll : ?CoroPoll
global g_coros : Array[?Coro] (.new(16))
global g_coro_sleepers: Array[Coro] (.new(16))
global g_coro_sleepers_next: Array[Coro] (.new(16))
global g_coro_indexes : type:Pool[uint] (.new(16))
global g_coro_poll_count : uint
global g_coro_count : uint

class Coro {
    args: ptr
    gc_args: Array[gc:GcPtr] (.{})

    stack: ?ptr
    stack_mem_size: uint (0)
    stack_size: uint (0)

    start_func: fnRef(Coro)()
    handler: fnRef()()
    follow_up: ?Coro
    next_resume: ?Coro
    g_list_index: uint (0)

    s_stack: ?ptr
    s_top: ptr (0.@cast(ptr))
    s_bottom: ptr (0.@cast(ptr))
    s_cont: ptr (0.@cast(ptr))
    s_size: uint (0)
    s_back_buf: ptr (0.@cast(ptr))

    error_code: u32 (0)
    error_msg: String ("")

    poll_event: PollEvent (0)
    sleep_until: uint (0)
    test: uint (10)
    error: u32 (0)
    started: bool (false)
    done: bool (false)

    static fn new(handler: ptr, start_func: ptr) SELF {

        let coro = Coro {
            args: mem:alloc(128)
            //
            start_func: start_func.@cast(fnRef(Coro)())
            handler: handler.@cast(fnRef()())
        }

        let indexes = g_coro_indexes
        if indexes.count == 0 {
            let index = g_coros.length
            g_coros.append(coro)
            coro.g_list_index = index
        } else {
            let index = indexes.get() !? 0
            g_coros.set(index, coro) _
            coro.g_list_index = index
        }

        let last = coro_run_last
        if isset(last) {
            last.next_resume = coro
            coro_run_last = coro
        } else {
            coro_run_first = coro
            coro_run_last = coro
        }

        g_coro_count++

        return coro
    }

    static fn get_poll() CoroPoll {
        let poll = coro_poll
        if poll == null {
            let p = CoroPoll.new()
            coro_poll = p
            return p
        }
        return poll.@cast(CoroPoll)
    }

    //////////////
    // AWAIT
    //////////////

    fn await_coro(coro: Coro) {
        if coro.done : return
        let ob = coro
        let ref = ob.follow_up
        while isset(ref) {
            ob = ref
            ref = ref.follow_up
        }
        ob.follow_up = this
        this.block()
    }

    fn await_fd(fd: FD, read: bool, write: bool) PollEvent {
        this.poll_fd(fd, read, write)
        this.block()
        return this.poll_event
    }

    fn poll_fd(fd: FD, read: bool, write: bool) {
        let poll = Coro.get_poll()
        poll.set(fd, this, read, write)
        g_coro_poll_count++
    }

    fn await_last() {
        let last = coro_run_last
        if isset(last) {
            last.next_resume = this
            coro_run_last = this
        } else {
            coro_run_first = this
            coro_run_last = this
        }
        this.block()
    }

    //////////////
    // FLOW
    //////////////

    fn block() {
        // let here = @stack(inline [ptr, 5])
        current_coro = null
        let here = @stack(<sys:libc_jmp_buf>)
        if @setjmp(here) == 0 {
            this.s_cont = here.@cast(ptr)
            this.s_top = this.s_cont - 32

            #if OS == win
            @ptrv(here, ptr, 0) += 0x80
            #end

            // Save stack
            let size = (this.s_bottom - this.s_top).@cast(uint)
            let stack = this.s_stack
            let has_stack = isset(stack)
            let create_stack = !has_stack || size > this.s_size
            if create_stack {
                if has_stack : mem:free(stack.@cast(ptr))
                stack = mem:alloc(size)
                this.s_size = size
                this.s_stack = stack
            }

            mem:copy(this.s_top, stack.@cast(ptr), size)

            // Save GC stack
            let gc_stack = this.stack
            let gc_size : uint = gc:stack_top - (&gc:stack_root).@cast(ptr)
            if !isset(gc_stack) || gc_size > this.stack_mem_size {
                if isset(gc_stack) : mem:free(gc_stack)
                gc_stack = mem:alloc(gc_size)
                this.stack_mem_size = gc_size
                this.stack = gc_stack
            }
            this.stack_size = gc_size

            mem:copy(gc:stack_root, gc_stack, gc_size)

            // Jump out
            @longjmp(this.s_back_buf)
        }
    }

    fn start() {
        let pad = @stack([u8 x 2048]).@cast(ptr)
        this.started = true
        this.s_bottom = pad
        let start = this.start_func
        start(this)
    }
    fn continue() {
        // Restore stack
        let size = (this.s_bottom - this.s_top).@cast(uint)
        let stack = this.s_stack.@cast(ptr)
        mem:copy(stack, this.s_top, size)
        // Restore GC stack
        let gc_stack = this.stack
        let gc_size = this.stack_size
        if isset(gc_stack) : mem:copy(gc_stack, gc:stack_root, gc_size)
        gc:stack_top = (&gc:stack_root).@cast(ptr) + gc_size
        // Jump
        @longjmp(this.s_cont)
    }

    static fn setjmp(buf: ptr) i32 {
        if @setjmp(buf) == 0 {
            return 0
        }
        return 1
    }
    fn resume() {
        if this.done : return
        current_coro = this

        gc:backup_gc_stack()

        // Start/continue coro
        let jmpbuf = @stack(<sys:libc_jmp_buf>)
        if @setjmp(jmpbuf) == 0 {

            #if OS == win
            @ptrv(jmpbuf, ptr, 0) += 0x80
            #end

            this.s_back_buf = jmpbuf
            if this.started == false {
                this.start()
            } else {
                this.continue()
            }
        }

        gc:restore_gc_stack()

        if this.done {
            let next = this.follow_up
            if isset(next) {
                this.follow_up = null
                next.resume()
            }
        }
    }

    fn complete() {
        this.done = true
        current_coro = null
        g_coro_count--
        let index = this.g_list_index

        let gc_stack = this.stack
        if isset(gc_stack) {
            mem:free(gc_stack)
            this.stack = null
            this.stack_size = 0
        }
        let s = this.s_stack
        if isset(s) : mem:free(s)

        g_coros.set(index, null) _
        g_coro_indexes.add(index)
        @longjmp(this.s_back_buf)
    }

    //////////////
    // LOOP
    //////////////

    static fn loop(until: ?Coro (null)) {

        let poll = Coro.get_poll()
        let coros = Array[Coro].new(16)
        let has_run = true

        while g_coro_count > 0 {
            if isset(until) && until.done : break

            // Run coro chain
            let coro = coro_run_first
            while isset(coro) {
                coro.resume()
                let next = coro.next_resume
                coro.next_resume = null
                coro = next
            }
            coro_run_first = null
            coro_run_last = null

            // Done check
            if isset(until) {
                if until.done : break
            }

            // Sleep info
            let timeout : uint = 1000
            let sleepers = g_coro_sleepers
            let sleepers_next = g_coro_sleepers_next
            let utime : uint = 0

            // Sleep coros
            if sleepers.length > 0 {
                utime = time:microtime()
                each sleepers as coro {
                    if utime >= coro.sleep_until {
                        coro.resume()
                    } else {
                        sleepers_next.append(coro)
                        let tm = (coro.sleep_until - utime) / 1000 + 1
                        if tm < timeout : timeout = tm
                    }
                }
                g_coro_sleepers = sleepers_next
                g_coro_sleepers_next = sleepers
                sleepers.clear()
            }

            // Done check
            if isset(until) {
                if until.done : break
            }

            // Wait for poll coros
            if g_coro_poll_count > 0 {
                poll.wait(timeout.@cast(i32), coros) ! {
                    panic("Too many file descriptors")
                }
                //
                if coros.length == 0 {
                    if has_run {
                        // has_run = false
                        gc:collect()
                    }
                } else {
                    has_run = true
                }
                // Resume coroutines
                g_coro_poll_count -= coros.length
                each coros as coro {
                    coro.resume()
                }

            } else if utime > 0 && timeout > 0 && g_coro_sleepers.length > 0 {
                gc:unlock()
                thread:sleep_ms(timeout)
                gc:lock()
            }
        }
    }

    ///////////////////
    // GC
    ///////////////////

    fn gc_free() {
        mem:free(this.args)
    }
}
