module async

import context
import sync

// PoolConfig configures a fixed-size concurrency pool.
//
// Both values must be positive. `workers` is the maximum number of jobs that
// can execute concurrently; `queue_size` is the bounded backlog accepted by
// try_submit() while all worker slots are busy.
@[params]
pub struct PoolConfig {
pub:
	workers    int
	queue_size int
}

// Pool limits concurrent JobFn execution with fixed worker slots and bounded backlog.
//
// The pool owns one derived context shared by all jobs. Closing the pool stops
// new submissions, waits for every accepted job, and returns the first job error
// if any.
@[heap]
pub struct Pool {
mut:
	ctx      context.Context
	cancel   context.CancelFn = unsafe { nil }
	tokens   chan bool
	wg       &sync.WaitGroup = sync.new_waitgroup()
	max_jobs int
	// Protects lifecycle flags, accepted job count, WaitGroup.add(), and
	// first_err. This lock is never held while a user JobFn runs.
	mutex     &sync.Mutex = sync.new_mutex()
	first_err IError      = none
	closed    bool
	waited    bool
	accepted  int
}

// new_pool creates a Pool with a background parent context.
pub fn new_pool(config PoolConfig) !&Pool {
	return new_pool_with_context(context.background(), config)
}

// new_pool_with_context creates a Pool with a context derived from parent.
//
// The worker limit and queue size are fixed for the pool lifetime. Parent
// cancellation is cooperative: jobs must observe `ctx.done()` and return.
pub fn new_pool_with_context(parent context.Context, config PoolConfig) !&Pool {
	if config.workers <= 0 {
		return error(err_pool_workers_invalid)
	}
	if config.queue_size <= 0 {
		return error(err_pool_queue_size_invalid)
	}
	ctx, cancel := new_cancel_context(parent)
	mut pool := &Pool{
		ctx:      context.Context(ctx)
		cancel:   cancel
		tokens:   chan bool{cap: config.workers}
		wg:       sync.new_waitgroup()
		max_jobs: config.workers + config.queue_size
		mutex:    sync.new_mutex()
	}
	for _ in 0 .. config.workers {
		pool.tokens <- true
	}
	return pool
}

// try_submit accepts f if the pool is open and its bounded backlog has capacity.
//
// It never blocks for queue space. A full backlog returns `async: pool queue is
// full`, making backpressure explicit for callers.
pub fn (mut p Pool) try_submit(f JobFn) ! {
	if f == unsafe { nil } {
		return error(err_nil_job)
	}
	p.mutex.lock()
	if p.closed {
		p.mutex.unlock()
		return error(err_pool_closed)
	}
	if p.accepted >= p.max_jobs {
		p.mutex.unlock()
		return error(err_pool_queue_full)
	}
	p.accepted++
	// add() is protected by the same mutex as wait(), so callers cannot race an
	// accepted job against pool shutdown.
	p.wg.add(1)
	p.mutex.unlock()
	// The JobFn is passed directly to the spawned wrapper instead of being
	// stored in a channel. That keeps closure ownership with V's normal spawn
	// path while the token channel below still enforces the fixed worker limit.
	spawn run_pool_job(mut p, f)
}

// wait closes the pool to new submissions, drains accepted jobs, and waits for completion.
//
// wait is one-shot. It returns the first job error observed by any accepted job.
pub fn (mut p Pool) wait() ! {
	p.mutex.lock()
	if p.waited {
		p.mutex.unlock()
		return error(err_pool_wait_called)
	}
	p.waited = true
	p.closed = true
	p.mutex.unlock()

	p.wg.wait()
	p.cancel()
	err := p.get_first_error()
	if err !is none {
		return err
	}
}

// close is an explicit lifecycle alias for wait().
//
// It rejects later submissions and waits for all accepted jobs before returning.
pub fn (mut p Pool) close() ! {
	p.wait()!
}

fn run_pool_job(mut p Pool, f JobFn) {
	defer {
		p.finish_accepted_job()
		p.wg.done()
	}
	// The token channel is a bounded semaphore. At most `workers` accepted jobs
	// can pass this point and run user code concurrently.
	_ := <-p.tokens
	defer {
		p.tokens <- true
	}
	mut job_ctx := p.ctx
	f(mut job_ctx) or { p.set_first_error(err) }
}

fn (mut p Pool) finish_accepted_job() {
	p.mutex.lock()
	p.accepted--
	p.mutex.unlock()
}

fn (mut p Pool) set_first_error(err IError) {
	p.mutex.lock()
	if p.first_err is none {
		p.first_err = err
	}
	p.mutex.unlock()
}

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