module onecontext

import context
import sync
import time

// canceled is the error returned when the cancel function is called on a merged context
pub const canceled = error('canceled context')

struct OneContext {
mut:
	ctx        context.Context
	ctxs       []context.Context
	done       chan int
	err        IError = none
	err_mutex  sync.Mutex
	cancel_fn  context.CancelFn
	cancel_ctx context.Context
}

// merge allows to merge multiple contexts
// it returns the merged context
pub fn merge(ctx context.Context, ctxs ...context.Context) (context.Context, context.CancelFn) {
	mut background := context.background()
	cancel_ctx, cancel := context.with_cancel(mut &background)
	mut octx := &OneContext{
		done: chan int{cap: 3}
		ctx: ctx
		ctxs: ctxs
		cancel_fn: cancel
		cancel_ctx: cancel_ctx
	}
	spawn octx.run()
	return context.Context(octx), context.CancelFn(cancel)
}

pub fn (octx OneContext) deadline() ?time.Time {
	mut min := time.Time{}

	if deadline := octx.ctx.deadline() {
		min = deadline
	}

	for ctx in octx.ctxs {
		if deadline := ctx.deadline() {
			if min.unix_time() == 0 || deadline < min {
				min = deadline
			}
		}
	}

	if min.unix_time() == 0 {
		return none
	}

	return min
}

pub fn (octx OneContext) done() chan int {
	return octx.done
}

pub fn (mut octx OneContext) err() IError {
	octx.err_mutex.@lock()
	defer {
		octx.err_mutex.unlock()
	}
	return octx.err
}

pub fn (octx OneContext) value(key context.Key) ?context.Any {
	if value := octx.ctx.value(key) {
		return value
	}

	for ctx in octx.ctxs {
		if value := ctx.value(key) {
			return value
		}
	}

	return none
}

pub fn (mut octx OneContext) run() {
	mut wrapped_ctx := &octx.ctx
	if octx.ctxs.len == 1 {
		mut first_ctx := &octx.ctxs[0]
		octx.run_two_contexts(mut wrapped_ctx, mut first_ctx)
		return
	}

	octx.run_multiple_contexts(mut wrapped_ctx)
	for mut ctx in octx.ctxs {
		octx.run_multiple_contexts(mut &ctx)
	}
}

pub fn (octx OneContext) str() string {
	return ''
}

pub fn (mut octx OneContext) cancel(err IError) {
	octx.cancel_fn()
	octx.err_mutex.@lock()
	octx.err = err
	octx.err_mutex.unlock()
	if !octx.done.closed {
		octx.done <- 0
		octx.done.close()
	}
}

pub fn (mut octx OneContext) run_two_contexts(mut ctx1 context.Context, mut ctx2 context.Context) {
	spawn fn (mut octx OneContext, mut ctx1 context.Context, mut ctx2 context.Context) {
		octx_cancel_done := octx.cancel_ctx.done()
		c1done := ctx1.done()
		c2done := ctx2.done()
		select {
			_ := <-octx_cancel_done {
				octx.cancel(onecontext.canceled)
			}
			_ := <-c1done {
				octx.cancel(ctx1.err())
			}
			_ := <-c2done {
				octx.cancel(ctx1.err())
			}
		}
	}(mut &octx, mut &ctx1, mut &ctx2)
}

pub fn (mut octx OneContext) run_multiple_contexts(mut ctx context.Context) {
	spawn fn (mut octx OneContext, mut ctx context.Context) {
		octx_cancel_done := octx.cancel_ctx.done()
		cdone := ctx.done()
		select {
			_ := <-octx_cancel_done {
				octx.cancel(onecontext.canceled)
			}
			_ := <-cdone {
				octx.cancel(ctx.err())
			}
		}
	}(mut &octx, mut &ctx)
}