package task_test

import (
	"context"
	"errors"
	"git.yuntuo.hk/k2safe/go-omni/lifecycle"
	"git.yuntuo.hk/k2safe/go-omni/task"
	"sync/atomic"
	"testing"
	"time"
)

func TestRunnerPhasedShutdown(t *testing.T) {
	runtime := lifecycle.NewRuntime(context.Background())
	runner, err := task.New(task.Config{Workers: 1, QueueSize: 1})
	if err != nil {
		t.Fatal(err)
	}
	var resourceClosed atomic.Bool
	runtime.OnClose(func() error { resourceClosed.Store(true); return nil })
	started, release, draining := make(chan struct{}), make(chan struct{}), make(chan struct{})
	runtime.OnStop("jobs", runner.StopAccepting)
	runtime.OnDrain("jobs", func(ctx context.Context) error { close(draining); return runner.Wait(ctx) })
	future, err := runner.Submit(runtime.WorkContext(), "write", func(ctx context.Context) error {
		close(started)
		<-release
		if resourceClosed.Load() {
			return errors.New("resource already closed")
		}
		return ctx.Err()
	})
	if err != nil {
		t.Fatal(err)
	}
	<-started
	runtime.RequestShutdown(nil)
	done := make(chan error, 1)
	go func() { done <- runtime.Close() }()
	select {
	case <-draining:
	case <-time.After(time.Second):
		t.Fatal("did not drain")
	}
	if _, err := runner.Submit(context.Background(), "late", func(context.Context) error { return nil }); !errors.Is(err, task.ErrClosed) {
		t.Fatal(err)
	}
	if resourceClosed.Load() {
		t.Fatal("closed resource before task")
	}
	close(release)
	if err := <-done; err != nil {
		t.Fatal(err)
	}
	if err := future.Wait(context.Background()); err != nil {
		t.Fatal(err)
	}
	if !resourceClosed.Load() {
		t.Fatal("resource leaked")
	}
}