module sync

import time
import rand

// how often to try to get data without blocking before to wait for semaphore
const spinloops = u32(750)
const spinloops_sem = u32(4000)

enum BufferElemStat {
    unused = 0
    writing
    written
    reading
}

struct Subscription {
mut:
    sem  &Semaphore     = unsafe { nil }
    prev &&Subscription = unsafe { nil }
    nxt  &Subscription  = unsafe { nil }
}

// append_subscription keeps select waiters in FIFO order, so the oldest waiter
// is woken first when a channel becomes ready.
fn append_subscription(head &&Subscription, sub &Subscription) {
    unsafe {
        mut link := head
        for *link != 0 {
            link = &(*link).nxt
        }
        sub.prev = link
        sub.nxt = nil
        *link = sub
    }
}

pub enum Direction {
    pop
    push
}

@[typedef]
pub struct C.atomic_uintptr_t {}

pub struct Channel {
    ringbuf   &u8 = unsafe { nil } // queue for buffered channels
    statusbuf &u8 = unsafe { nil } // flags to synchronize write/read in ringbuf
    objsize   u32
mut:
    // atomic
    writesem           Semaphore // to wake thread that wanted to write, but buffer was full
    readsem            Semaphore // to wake thread that wanted to read, but buffer was empty
    writesem_im        Semaphore
    readsem_im         Semaphore
    write_adr          C.atomic_uintptr_t // if != NULL the next obj can be written here without wait
    read_adr           C.atomic_uintptr_t // if != NULL an obj can be read from here without wait
    adr_read           C.atomic_uintptr_t // used to identify origin of writesem
    adr_written        C.atomic_uintptr_t // used to identify origin of readsem
    write_free         u32                // for queue state
    read_avail         u32
    buf_elem_write_idx u32
    buf_elem_read_idx  u32
    // for select
    write_subscriber &Subscription = unsafe { nil }
    read_subscriber  &Subscription = unsafe { nil }
    write_sub_mtx    &SpinLock
    read_sub_mtx     &SpinLock
    closed           u16
    close_err        IError = none
pub:
    cap u32 // queue length in #objects
}

pub fn new_channel[T](n u32) &Channel {
    st := if sizeof(T) > 0 { sizeof(T) } else { 1 }
    if isreftype(T) {
        return new_channel_st(n, st)
    } else {
        return new_channel_st_noscan(n, st)
    }
}

fn new_channel_st(n u32, st u32) &Channel {
    wsem := if n > 0 { n } else { 1 }
    rsem := if n > 0 { u32(0) } else { 1 }
    rbuf := if n > 0 { unsafe { malloc(int(n * st)) } } else { &u8(unsafe { nil }) }
    sbuf := if n > 0 { vcalloc_noscan(int(n * 2)) } else { &u8(unsafe { nil }) }
    mut ch := Channel{
        objsize:          st
        cap:              n
        write_free:       n
        read_avail:       0
        ringbuf:          rbuf
        statusbuf:        sbuf
        write_subscriber: unsafe { nil }
        read_subscriber:  unsafe { nil }
        write_sub_mtx:    new_spin_lock()
        read_sub_mtx:     new_spin_lock()
    }
    ch.writesem.init(wsem)
    ch.readsem.init(rsem)
    ch.writesem_im.init(0)
    ch.readsem_im.init(0)
    return &ch
}

fn new_channel_st_noscan(n u32, st u32) &Channel {
    $if gcboehm_opt ? {
        wsem := if n > 0 { n } else { 1 }
        rsem := if n > 0 { u32(0) } else { 1 }
        rbuf := if n > 0 { unsafe { malloc_noscan(int(n * st)) } } else { &u8(unsafe { nil }) }
        sbuf := if n > 0 { vcalloc_noscan(int(n * 2)) } else { &u8(unsafe { nil }) }
        mut ch := Channel{
            objsize:          st
            cap:              n
            write_free:       n
            read_avail:       0
            ringbuf:          rbuf
            statusbuf:        sbuf
            write_subscriber: unsafe { nil }
            read_subscriber:  unsafe { nil }
            write_sub_mtx:    new_spin_lock()
            read_sub_mtx:     new_spin_lock()
        }
        ch.writesem.init(wsem)
        ch.readsem.init(rsem)
        ch.writesem_im.init(0)
        ch.readsem_im.init(0)
        return &ch
    } $else {
        return new_channel_st(n, st)
    }
}

// close closes the channel and optionally stores an error that will be
// returned by receive operations that use `or {}` or `?` after the
// buffered values have been drained.
pub fn (mut ch Channel) close(errs ...IError) {
    open_val := u16(0)
    if !C.atomic_compare_exchange_strong_u16(&ch.closed, &open_val, 1) {
        return
    }
    if errs.len > 0 {
        ch.close_err = errs[0]
    }
    mut nulladr := unsafe { nil }
    for !C.atomic_compare_exchange_weak_ptr(voidptr(&ch.adr_written), voidptr(&nulladr), isize(-1)) {
        nulladr = unsafe { nil }
    }
    ch.readsem_im.post()
    ch.readsem.post()
    ch.read_sub_mtx.lock()
    if ch.read_subscriber != unsafe { nil } {
        ch.read_subscriber.sem.post()
    }
    ch.read_sub_mtx.unlock()
    ch.write_sub_mtx.lock()
    if ch.write_subscriber != unsafe { nil } {
        ch.write_subscriber.sem.post()
    }
    ch.write_sub_mtx.unlock()
    ch.writesem.post()
    if ch.cap == 0 {
        C.atomic_store_ptr(unsafe { &voidptr(&ch.read_adr) }, unsafe { nil })
    }
    ch.writesem_im.post()

    // Do not destroy `read_sub_mtx` and `write_sub_mtx` here,
    // because we can read from a closed channel later.
}

@[inline]
fn (ch &Channel) closed_error() IError {
    if ch.close_err !is None__ {
        return ch.close_err
    }
    return error('channel closed')
}

@[inline]
pub fn (mut ch Channel) len() int {
    return int(C.atomic_load_u32(&ch.read_avail))
}

@[inline]
pub fn (mut ch Channel) closed() bool {
    return C.atomic_load_u16(&ch.closed) != 0
}

@[inline]
pub fn (mut ch Channel) push(src voidptr) {
    if ch.try_push_priv(src, false) == .closed {
        panic('push on closed channel')
    }
}

@[inline]
pub fn (mut ch Channel) try_push(src voidptr) ChanState {
    return ch.try_push_priv(src, true)
}

fn (mut ch Channel) try_push_priv(src voidptr, no_block bool) ChanState {
    if C.atomic_load_u16(&ch.closed) != 0 {
        return .closed
    }
    spinloops_sem_, spinloops_ := if no_block { u32(1), u32(1) } else { spinloops, spinloops_sem }
    mut have_swapped := false
    for {
        mut got_sem := false
        mut wradr := C.atomic_load_ptr(unsafe { &voidptr(&ch.write_adr) })
        for wradr != C.NULL {
            if C.atomic_compare_exchange_strong_ptr(voidptr(&ch.write_adr), voidptr(&wradr),
                isize(0))
            {
                // there is a reader waiting for us
                unsafe { C.memcpy(wradr, src, ch.objsize) }
                mut nulladr := unsafe { nil }
                for !C.atomic_compare_exchange_weak_ptr(voidptr(&ch.adr_written),
                    voidptr(&nulladr), isize(wradr)) {
                    nulladr = unsafe { nil }
                }
                ch.readsem_im.post()
                return .success
            }
        }
        if no_block && ch.cap == 0 {
            return .not_ready
        }
        // get token to read
        for _ in 0 .. spinloops_sem_ {
            if got_sem {
                break
            }
            got_sem = ch.writesem.try_wait()
        }
        if !got_sem {
            if no_block {
                return .not_ready
            }
            ch.writesem.wait()
        }
        if C.atomic_load_u16(&ch.closed) != 0 {
            ch.writesem.post()
            return .closed
        }
        if ch.cap == 0 {
            // try to advertise current object as readable
            mut read_in_progress := false
            C.atomic_store_ptr(unsafe { &voidptr(&ch.read_adr) }, src)
            wradr = C.atomic_load_ptr(unsafe { &voidptr(&ch.write_adr) })
            if wradr != C.NULL {
                mut src2 := src
                if C.atomic_compare_exchange_strong_ptr(voidptr(&ch.read_adr), voidptr(&src2),
                    isize(0))
                {
                    ch.writesem.post()
                    continue
                } else {
                    read_in_progress = true
                }
            }
            if !read_in_progress {
                ch.read_sub_mtx.lock()
                if ch.read_subscriber != unsafe { nil } {
                    ch.read_subscriber.sem.post()
                }
                ch.read_sub_mtx.unlock()
            }
            mut src2 := src
            for sp := u32(0); sp < spinloops_ || read_in_progress; sp++ {
                if C.atomic_compare_exchange_strong_ptr(voidptr(&ch.adr_read), voidptr(&src2),
                    isize(0))
                {
                    have_swapped = true
                    read_in_progress = true
                    break
                }
                src2 = src
            }
            mut got_im_sem := false
            for sp := u32(0); sp < spinloops_sem_ || read_in_progress; sp++ {
                got_im_sem = ch.writesem_im.try_wait()
                if got_im_sem {
                    break
                }
            }
            for {
                if got_im_sem {
                    got_im_sem = false
                } else {
                    ch.writesem_im.wait()
                }
                if C.atomic_load_u16(&ch.closed) != 0 {
                    if have_swapped
                        || C.atomic_compare_exchange_strong_ptr(voidptr(&ch.adr_read), voidptr(&src2), isize(0)) {
                        ch.writesem.post()
                        return .success
                    } else {
                        return .closed
                    }
                }
                if have_swapped
                    || C.atomic_compare_exchange_strong_ptr(voidptr(&ch.adr_read), voidptr(&src2), isize(0)) {
                    ch.writesem.post()
                    break
                } else {
                    // this semaphore was not for us - repost in
                    ch.writesem_im.post()
                    if src2 == voidptr(-1) {
                        ch.readsem.post()
                        return .closed
                    }
                    src2 = src
                }
            }
            return .success
        } else {
            // buffered channel
            mut space_in_queue := false
            mut wr_free := C.atomic_load_u32(&ch.write_free)
            for wr_free > 0 {
                space_in_queue = C.atomic_compare_exchange_weak_u32(&ch.write_free, &wr_free,
                    wr_free - 1)
                if space_in_queue {
                    break
                }
            }
            if space_in_queue {
                mut wr_idx := C.atomic_load_u32(&ch.buf_elem_write_idx)
                for {
                    mut new_wr_idx := wr_idx + 1
                    for new_wr_idx >= ch.cap {
                        new_wr_idx -= ch.cap
                    }
                    if C.atomic_compare_exchange_strong_u32(&ch.buf_elem_write_idx, &wr_idx,
                        new_wr_idx)
                    {
                        break
                    }
                }
                mut wr_ptr := ch.ringbuf
                mut status_adr := ch.statusbuf
                unsafe {
                    wr_ptr += (wr_idx * ch.objsize)
                    status_adr += wr_idx * sizeof(u16)
                }
                mut expected_status := u16(BufferElemStat.unused)
                for !C.atomic_compare_exchange_weak_u16(status_adr, &expected_status,
                    u16(BufferElemStat.writing)) {
                    expected_status = u16(BufferElemStat.unused)
                }
                unsafe {
                    C.memcpy(wr_ptr, src, ch.objsize)
                }
                C.atomic_store_u16(unsafe { &u16(status_adr) }, u16(BufferElemStat.written))
                C.atomic_fetch_add_u32(voidptr(&ch.read_avail), 1)
                ch.readsem.post()
                ch.read_sub_mtx.lock()
                if ch.read_subscriber != unsafe { nil } {
                    ch.read_subscriber.sem.post()
                }
                ch.read_sub_mtx.unlock()
                return .success
            } else {
                if no_block {
                    return .not_ready
                }
                ch.writesem.post()
            }
        }
    }
    // we should not get here but the V compiler want's to see a return statement
    panic('unknown `try_push_priv` state')
}

@[inline]
pub fn (mut ch Channel) pop(dest voidptr) bool {
    return ch.try_pop_priv(dest, false) == .success
}

// try_pop returns `.success` if an object is popped without blocking.
// Pass the destination as `mut`: `ch.try_pop(mut value)`, not `ch.try_pop(&value)`.
@[inline]
pub fn (mut ch Channel) try_pop(dest voidptr) ChanState {
    return ch.try_pop_priv(dest, true)
}

// try_pop_select_priv treats already closed channels as unavailable for non-blocking `select ... else`.
@[inline]
fn (mut ch Channel) try_pop_select_priv(dest voidptr) ChanState {
    if C.atomic_load_u16(&ch.closed) != 0 {
        return .closed
    }
    return ch.try_pop_priv(dest, true)
}

fn (mut ch Channel) try_pop_priv(dest voidptr, no_block bool) ChanState {
    spinloops_sem_, spinloops_ := if no_block { u32(1), u32(1) } else { spinloops, spinloops_sem }
    mut have_swapped := false
    mut write_in_progress := false
    for {
        mut got_sem := false
        if ch.cap == 0 {
            // unbuffered channel - first see if a `push()` has adversized
            mut rdadr := C.atomic_load_ptr(unsafe { &voidptr(&ch.read_adr) })
            for rdadr != C.NULL {
                if C.atomic_compare_exchange_strong_ptr(voidptr(&ch.read_adr), voidptr(&rdadr),
                    isize(0))
                {
                    // there is a writer waiting for us
                    unsafe { C.memcpy(dest, rdadr, ch.objsize) }
                    mut nulladr := unsafe { nil }
                    for !C.atomic_compare_exchange_weak_ptr(voidptr(&ch.adr_read),
                        voidptr(&nulladr), isize(rdadr)) {
                        nulladr = unsafe { nil }
                    }
                    ch.writesem_im.post()
                    return .success
                }
            }
            if no_block {
                if C.atomic_load_u16(&ch.closed) == 0 {
                    return .not_ready
                } else {
                    return .closed
                }
            }
        }
        // get token to read
        for _ in 0 .. spinloops_sem_ {
            if got_sem {
                break
            }
            got_sem = ch.readsem.try_wait()
        }
        if !got_sem {
            if no_block {
                if C.atomic_load_u16(&ch.closed) == 0 {
                    return .not_ready
                } else {
                    return .closed
                }
            }
            ch.readsem.wait()
        }
        if ch.cap > 0 {
            // try to get buffer token
            mut obj_in_queue := false
            mut rd_avail := C.atomic_load_u32(&ch.read_avail)
            for rd_avail > 0 {
                obj_in_queue = C.atomic_compare_exchange_weak_u32(&ch.read_avail, &rd_avail,
                    rd_avail - 1)
                if obj_in_queue {
                    break
                }
            }
            if obj_in_queue {
                mut rd_idx := C.atomic_load_u32(&ch.buf_elem_read_idx)
                for {
                    mut new_rd_idx := rd_idx + 1
                    for new_rd_idx >= ch.cap {
                        new_rd_idx -= ch.cap
                    }
                    if C.atomic_compare_exchange_weak_u32(&ch.buf_elem_read_idx, &rd_idx,
                        new_rd_idx)
                    {
                        break
                    }
                }
                mut rd_ptr := ch.ringbuf
                mut status_adr := ch.statusbuf
                unsafe {
                    rd_ptr += rd_idx * ch.objsize
                    status_adr += rd_idx * sizeof(u16)
                }
                mut expected_status := u16(BufferElemStat.written)
                for !C.atomic_compare_exchange_weak_u16(status_adr, &expected_status,
                    u16(BufferElemStat.reading)) {
                    expected_status = u16(BufferElemStat.written)
                }
                unsafe {
                    C.memcpy(dest, rd_ptr, ch.objsize)
                }
                C.atomic_store_u16(unsafe { &u16(status_adr) }, u16(BufferElemStat.unused))
                C.atomic_fetch_add_u32(voidptr(&ch.write_free), 1)
                ch.writesem.post()
                ch.write_sub_mtx.lock()
                if ch.write_subscriber != unsafe { nil } {
                    ch.write_subscriber.sem.post()
                }
                ch.write_sub_mtx.unlock()
                return .success
            }
        }
        // try to advertise `dest` as writable
        C.atomic_store_ptr(unsafe { &voidptr(&ch.write_adr) }, dest)
        if ch.cap == 0 {
            mut rdadr := C.atomic_load_ptr(unsafe { &voidptr(&ch.read_adr) })
            if rdadr != C.NULL {
                mut dest2 := dest
                if C.atomic_compare_exchange_strong_ptr(voidptr(&ch.write_adr), voidptr(&dest2),
                    isize(0))
                {
                    ch.readsem.post()
                    continue
                } else {
                    write_in_progress = true
                }
            }
        }
        if ch.cap == 0 && !write_in_progress {
            ch.write_sub_mtx.lock()
            if ch.write_subscriber != unsafe { nil } {
                ch.write_subscriber.sem.post()
            }
            ch.write_sub_mtx.unlock()
        }
        mut dest2 := dest
        for sp := u32(0); sp < spinloops_ || write_in_progress; sp++ {
            if C.atomic_compare_exchange_strong_ptr(voidptr(&ch.adr_written), voidptr(&dest2),
                isize(0))
            {
                have_swapped = true
                break
            } else if dest2 == voidptr(-1) {
                ch.readsem.post()
                return .closed
            }
            dest2 = dest
        }
        mut got_im_sem := false
        for sp := u32(0); sp < spinloops_sem_ || write_in_progress; sp++ {
            got_im_sem = ch.readsem_im.try_wait()
            if got_im_sem {
                break
            }
        }
        for {
            if got_im_sem {
                got_im_sem = false
            } else {
                ch.readsem_im.wait()
            }
            if have_swapped
                || C.atomic_compare_exchange_strong_ptr(voidptr(&ch.adr_written), voidptr(&dest2), isize(0)) {
                ch.readsem.post()
                break
            } else {
                // this semaphore was not for us - repost in
                ch.readsem_im.post()
                if dest2 == voidptr(-1) {
                    ch.readsem.post()
                    return .closed
                }
                dest2 = dest
            }
        }
        break
    }
    return .success
}

// channel_select waits `timeout` on any of `channels[i]` until one of them can
// push (`dir[i] == .push`) or pop (`dir[i] == .pop`) the object referenced by
// `objrefs[i]`. `timeout = time.infinite` means wait unlimited time.
// `timeout <= 0` means return immediately if no transaction can be performed
// without waiting. It returns the selected channel index, `-1` on timeout, and
// `-2` when all channels are closed.
pub fn channel_select(mut channels []&Channel, dir []Direction, mut objrefs []voidptr, timeout time.Duration) int {
    skip_closed_pop := timeout < 0
    actual_timeout := if skip_closed_pop { time.Duration(0) } else { timeout }
    closed_pop_mode := if skip_closed_pop {
        SelectClosedPopMode.skip
    } else {
        SelectClosedPopMode.closed
    }
    return channel_select_priv(mut channels, dir, mut objrefs, actual_timeout, closed_pop_mode)
}

enum SelectClosedPopMode {
    closed
    ready
    skip
}

// channel_select_lang is used by the language `select` implementation.
// Closed receive cases stay selectable for blocking/timed selects to match
// plain `<-ch` semantics, while `select ... else` still skips them.
fn channel_select_lang(mut channels []&Channel, dir []Direction, mut objrefs []voidptr, timeout time.Duration) int {
    skip_closed_pop := timeout < 0
    actual_timeout := if skip_closed_pop { time.Duration(0) } else { timeout }
    closed_pop_mode := if skip_closed_pop {
        SelectClosedPopMode.skip
    } else {
        SelectClosedPopMode.ready
    }
    return channel_select_priv(mut channels, dir, mut objrefs, actual_timeout, closed_pop_mode)
}

fn channel_select_priv(mut channels []&Channel, dir []Direction, mut objrefs []voidptr, timeout time.Duration, closed_pop_mode SelectClosedPopMode) int {
    $if debug_channels ? {
        assert channels.len == dir.len
        assert dir.len == objrefs.len
    }
    mut subscr := []Subscription{len: channels.len}
    mut sem := unsafe { Semaphore{} }
    sem.init(0)
    for i, ch in channels {
        subscr[i].sem = unsafe { &sem }
        sub_mtx, subscriber := if dir[i] == .push {
            ch.write_sub_mtx, &ch.write_subscriber
        } else {
            ch.read_sub_mtx, &ch.read_subscriber
        }
        sub_mtx.lock()
        unsafe {
            append_subscription(subscriber, &subscr[i])
        }
        sub_mtx.unlock()
    }
    stopwatch := if timeout == time.infinite || timeout <= 0 {
        time.StopWatch{}
    } else {
        time.new_stopwatch()
    }
    mut event_idx := -1 // negative index means `timed out`

    outer: for {
        rnd := rand.intn(channels.len) or { 0 }
        mut num_closed := 0
        mut ready_closed_idx := -1
        for j, _ in channels {
            mut i := j + rnd
            if i >= channels.len {
                i -= channels.len
            }
            stat := if dir[i] == .push {
                channels[i].try_push_priv(objrefs[i], true)
            } else if closed_pop_mode == .skip {
                channels[i].try_pop_select_priv(objrefs[i])
            } else {
                channels[i].try_pop_priv(objrefs[i], true)
            }
            if stat == .success {
                event_idx = i
                break outer
            } else if stat == .closed && dir[i] == .pop && closed_pop_mode == .ready {
                unsafe {
                    C.memset(objrefs[i], 0, channels[i].objsize)
                }
                if ready_closed_idx == -1 {
                    ready_closed_idx = i
                }
                num_closed++
            } else if stat == .closed {
                num_closed++
            }
        }
        if ready_closed_idx >= 0 {
            event_idx = ready_closed_idx
            break outer
        }
        if num_closed == channels.len {
            event_idx = -2
            break outer
        }
        if timeout <= 0 {
            break outer
        }
        if timeout != time.infinite {
            remaining := timeout - stopwatch.elapsed()
            if !sem.timed_wait(remaining) {
                break outer
            }
        } else {
            sem.wait()
        }
    }
    // reset subscribers
    for i, ch in channels {
        sub_mtx := if dir[i] == .push {
            ch.write_sub_mtx
        } else {
            ch.read_sub_mtx
        }
        sub_mtx.lock()
        unsafe {
            *subscr[i].prev = subscr[i].nxt
        }
        if unsafe { subscr[i].nxt != 0 } {
            subscr[i].nxt.prev = subscr[i].prev
            subscr[i].nxt.sem.post()
        }
        sub_mtx.unlock()
    }
    sem.destroy()
    return event_idx
}