module async

import context
import sync

// Group coordinates a set of related concurrent jobs.
//
// Jobs share a derived context. The first job error is returned by wait() and
// cancels the shared context so sibling jobs can stop cooperatively.
@[heap]
pub struct Group {
mut:
    ctx    context.Context
    cancel context.CancelFn = unsafe { nil }
    wg     &sync.WaitGroup  = sync.new_waitgroup()
    // Protects both WaitGroup lifecycle state and first_err. WaitGroup.add()
    // must not race with WaitGroup.wait(), so go() and wait() share this lock.
    mutex &sync.Mutex = sync.new_mutex()
    // Stored once. Later job errors never replace the first failure observed by
    // the group, even when several jobs fail concurrently after cancellation.
    first_err IError = none
    // Set before wait() calls wg.wait(). Once true, no further jobs are accepted.
    waiting bool
}

// new_group creates a Group with a shared cancellable context derived from parent.
//
// The parent is accepted by value to keep the public call site simple. The
// derived context is owned by the group and canceled on first job error or when
// wait() completes.
pub fn new_group(parent context.Context) &Group {
    ctx, cancel := new_cancel_context(parent)
    return &Group{
        ctx:    context.Context(ctx)
        cancel: cancel
        wg:     sync.new_waitgroup()
        mutex:  sync.new_mutex()
    }
}

// go starts f in a new concurrent task.
//
// Calling go after wait has started returns an error. The task should not
// panic; panics in spawned work are not recovered by x.async.
pub fn (mut g Group) go(f JobFn) ! {
    if f == unsafe { nil } {
        return error(err_nil_job)
    }
    g.mutex.lock()
    if g.waiting {
        g.mutex.unlock()
        return error(err_group_go_after_wait)
    }
    // add() happens while holding the same mutex that guards wait(), preventing
    // callers from triggering sync.WaitGroup's add-while-waiting misuse panic.
    g.wg.add(1)
    g.mutex.unlock()
    spawn run_group_job(mut g, f)
}

// wait blocks until all accepted group jobs finish.
//
// It returns the first job error, if any. wait may be called once; after it
// starts, the group no longer accepts new jobs.
pub fn (mut g Group) wait() ! {
    g.mutex.lock()
    if g.waiting {
        g.mutex.unlock()
        return error(err_group_wait_called)
    }
    g.waiting = true
    g.mutex.unlock()

    g.wg.wait()
    // Always cancel after all jobs finish to release the derived context and to
    // make the lifecycle symmetric with context.with_cancel().
    g.cancel()
    err := g.get_first_error()
    if err !is none {
        return err
    }
}

fn run_group_job(mut g Group, f JobFn) {
    defer {
        g.wg.done()
    }
    // Each job gets its own local mutable interface value. The underlying
    // context is shared and synchronized by the context module.
    mut job_ctx := g.ctx
    f(mut job_ctx) or { g.set_first_error(err) }
}

fn (mut g Group) set_first_error(err IError) {
    mut should_cancel := false
    g.mutex.lock()
    if g.first_err is none {
        g.first_err = err
        should_cancel = true
    }
    g.mutex.unlock()
    if should_cancel {
        // Cancel outside the group mutex. context cancellation can notify child
        // contexts, so keeping our own lock out of that path avoids lock nesting.
        g.cancel()
    }
}

fn (mut g Group) get_first_error() IError {
    g.mutex.lock()
    err := g.first_err
    g.mutex.unlock()
    return err
}