// priority 演示近同时到期任务的优先级顺序，真实受理与执行由 Orbit 完成。
package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"time"

	orbit "git.yuntuo.hk/k2safe/orbit/sdk/go"
)

func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}
func run() error {
	token := os.Getenv("ORBIT_TOKEN")
	if path := os.Getenv("ORBIT_CREDENTIALS_FILE"); path != "" && token == "" {
		raw, err := os.ReadFile(path)
		if err != nil {
			return err
		}
		var c struct {
			Token string `json:"token"`
		}
		if err = json.Unmarshal(raw, &c); err != nil {
			return err
		}
		token = c.Token
	}
	api, err := orbit.NewClient(os.Getenv("ORBIT_URL"), orbit.Options{Token: token})
	if err != nil {
		return err
	}
	ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second)
	defer cancel()
	op, err := orbit.NewWorkerMessageID()
	if err != nil {
		return err
	}
	name := "priority." + op
	created, err := api.CreateQueue(ctx, "queue-"+op, orbit.CreateQueueRequest{Name: name, MaxConcurrency: 1, RatePerSecond: 100, MaxUnfinishedTasks: 3})
	if err != nil {
		return err
	}
	var tasks []orbit.Task
	for _, priority := range []int16{0, 9, 4} {
		in := orbit.SubmitTaskRequest{QueueID: created.Queue.ID, Handler: name, HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{"object_ref":"objects/example","n":9007199254740993}`), Priority: priority, TaskTimeoutMS: 20000, Retry: orbit.RetryPolicy{MaxAttempts: 1, AttemptTimeoutMS: 10000}}
		out, err := api.SubmitTask(ctx, fmt.Sprintf("task-%s-%d", op, priority), in)
		if err != nil {
			return err
		}
		tasks = append(tasks, out.Task)
	}
	order := make(chan int16, 3)
	w, err := orbit.NewWorker(api, orbit.WorkerOptions{QueueID: created.Queue.ID, Capacity: 1, ShutdownTimeout: 5 * time.Second, Handlers: []orbit.WorkerHandler{{Capability: orbit.WorkerCapability{Handler: name, HandlerVersion: "v1", PayloadVersion: "v1"}, Handle: func(ctx context.Context, e orbit.WorkerExecution) (json.RawMessage, error) {
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		order <- e.Task.Priority
		return e.Task.Payload, nil
	}}}})
	if err != nil {
		return err
	}
	run, stop := context.WithCancel(ctx)
	done := make(chan error, 1)
	go func() { done <- w.Run(run) }()
	waited := false
	defer func() {
		stop()
		if !waited {
			<-done
		}
	}()
	for _, task := range tasks {
		for {
			v, err := api.GetTask(ctx, task.ID)
			if err != nil {
				return err
			}
			if v.State == orbit.TaskSucceeded {
				if v.Priority != task.Priority || v.AttemptCount != 1 {
					return errors.New("frozen priority or execution count changed")
				}
				break
			}
			if v.State == orbit.TaskFailed || v.State == orbit.TaskNeedsReview || v.State == orbit.TaskCancelled {
				return fmt.Errorf("unexpected task state: %s", v.State)
			}
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-time.After(80 * time.Millisecond):
			}
		}
	}
	got := []int16{<-order, <-order, <-order}
	if got[0] != 9 || got[1] != 4 || got[2] != 0 {
		return fmt.Errorf("unexpected dispatch order: %v", got)
	}
	stop()
	err = <-done
	waited = true
	if err != nil {
		return err
	}
	fmt.Println("submitted_priorities=0,9,4 executed_priorities=9,4,0")
	fmt.Println("completed=3 attempts_per_task=1 worker_drained=true")
	return nil
}