package orbit

import (
	"context"
	"encoding/json"
	"errors"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func TestIntegrationWorkerRuntimeKeyedConcurrency(t *testing.T) {
	api, _, token := integrationWorkerClient(t)
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	name := "keyed." + workerTestID(t)
	gateway := newWorkerFaultGateway(t, api.baseURL)
	client, err := NewClient(gateway.server.URL, Options{Token: token, HTTPClient: gateway.server.Client()})
	if err != nil {
		t.Fatal(err)
	}
	var active, calls atomic.Int32
	var mu sync.Mutex
	tokens := map[string]uint64{}
	var last uint64
	handler := runtimeHandler(name, func(ctx context.Context, e WorkerExecution) (json.RawMessage, error) {
		calls.Add(1)
		if active.Add(1) != 1 {
			t.Error("same resource handlers overlapped")
		}
		defer active.Add(-1)
		mu.Lock()
		if e.Task.ConcurrencyKey != name || e.ConcurrencyToken <= last {
			t.Error("key token was lost, duplicated or reset")
		}
		last = e.ConcurrencyToken
		tokens[e.Task.ID] = e.ConcurrencyToken
		mu.Unlock()
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-time.After(100 * time.Millisecond):
		}
		return e.Task.Payload, nil
	})
	var tasks []Task
	var workers []*Worker
	for i := 0; i < 2; i++ {
		q, err := api.CreateQueue(ctx, workerTestID(t), CreateQueueRequest{Name: name + workerTestID(t), MaxConcurrency: 4, RatePerSecond: 100})
		if err != nil {
			t.Fatal(err)
		}
		for j := 0; j < 2; j++ {
			out, err := api.SubmitTask(ctx, workerTestID(t), SubmitTaskRequest{QueueID: q.Queue.ID, Handler: name, HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{"n":9007199254740993}`), ConcurrencyKey: name})
			if err != nil {
				t.Fatal(err)
			}
			tasks = append(tasks, out.Task)
		}
		w, err := NewWorker(client, WorkerOptions{QueueID: q.Queue.ID, Capacity: 4, Handlers: []WorkerHandler{handler}})
		if err != nil {
			t.Fatal(err)
		}
		workers = append(workers, w)
	}
	// 在真实 Start 提交后丢 ACK；恢复必须沿用原 key/token，不能重跑处理器。
	gateway.drop.Store("task.start")
	run, stop := context.WithCancel(ctx)
	done := make(chan error, 2)
	for _, w := range workers {
		go func() { done <- w.Run(run) }()
	}
	defer func() {
		stop()
		for range workers {
			if err := <-done; err != nil {
				t.Error(err)
			}
		}
	}()
	for _, task := range tasks {
		finished := waitRuntimeState(t, ctx, api, task.ID, TaskSucceeded)
		if finished.AttemptCount != 1 || finished.ConcurrencyKey != name {
			t.Fatal("keyed task lost identity")
		}
		history, err := api.ListTaskAttempts(ctx, task.ID, HistoryPageRequest{})
		if err != nil || len(history.Items) != 1 {
			t.Fatal("missing keyed history", err)
		}
		mu.Lock()
		want := tokens[task.ID]
		mu.Unlock()
		if want == 0 || history.Items[0].ConcurrencyToken != want {
			t.Fatal("history token differs from handler")
		}
		detail, err := api.GetTaskAttempt(ctx, task.ID, history.Items[0].ID)
		if err != nil || detail.ConcurrencyToken != want {
			t.Fatal("detail token differs from handler", err)
		}
	}
	gateway.expectDrop(t, "task.start")
	if calls.Load() != 4 || active.Load() != 0 {
		t.Fatal("incorrect handler count")
	}
}

func TestIntegrationWorkerProcessKeyedConcurrency(t *testing.T) {
	a := startOrbitProcess(t)
	_, _, token := integrationWorkerClient(t)
	api := processClient(t, a, token)
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	name := "key-crash." + workerTestID(t)
	var registrations []WorkerRegistration
	var tasks []Task
	for i := 0; i < 2; i++ {
		q, err := api.CreateQueue(ctx, workerTestID(t), CreateQueueRequest{Name: name + workerTestID(t), MaxConcurrency: 2, RatePerSecond: 100})
		if err != nil {
			t.Fatal(err)
		}
		r, err := NewWorkerRegistration(q.Queue.ID, 1, []WorkerCapability{{Handler: name, HandlerVersion: "v1", PayloadVersion: "v1"}})
		if err != nil {
			t.Fatal(err)
		}
		r.ConcurrencyKeys = true
		registrations = append(registrations, r)
		out, err := api.SubmitTask(ctx, workerTestID(t), SubmitTaskRequest{QueueID: q.Queue.ID, Handler: name, HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{}`), ConcurrencyKey: name})
		if err != nil {
			t.Fatal(err)
		}
		tasks = append(tasks, out.Task)
	}
	dial := func(api *Client, r WorkerRegistration) *WorkerConnection {
		t.Helper()
		w, err := api.DialWorker(ctx, r)
		if err != nil {
			t.Fatal(err)
		}
		t.Cleanup(func() { w.Close() })
		return w
	}
	capacity := func(w *WorkerConnection) {
		t.Helper()
		if _, err := w.Capacity(ctx, workerTestID(t), WorkerCapacity{Sequence: 1, Capacity: 1}); err != nil {
			t.Fatal(err)
		}
	}
	w1 := dial(api, registrations[0])
	capacity(w1)
	first, err := w1.NextAssignment(ctx)
	if err != nil || first.ConcurrencyToken == 0 {
		t.Fatal("first key not granted", err)
	}
	started, err := w1.Start(ctx, workerTestID(t), first.Attempt())
	if err != nil {
		t.Fatal(err)
	}
	if _, active := started.Lease.ExecutionDeadline(); !active {
		t.Fatal("no current execution authorization")
	}
	a.kill(t)
	b := startOrbitProcess(t)
	api = processClient(t, b, token)
	if a.nodeID == b.nodeID {
		t.Fatal("node id reused")
	}
	w1 = dial(api, registrations[0])
	lease, err := w1.Resume(ctx, workerTestID(t), first.Attempt())
	if err != nil {
		t.Fatal(err)
	}
	if _, active := lease.ExecutionDeadline(); !active {
		t.Fatal("original attempt not resumed")
	}
	w2 := dial(api, registrations[1])
	capacity(w2)
	brief, stop := context.WithTimeout(ctx, 300*time.Millisecond)
	_, err = w2.NextAssignment(brief)
	stop()
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatal("restart lost key exclusion", err)
	}
	waiting, err := api.GetTask(ctx, tasks[1].ID)
	if err != nil || waiting.AttemptCount != 0 {
		t.Fatal("blocked key consumed attempt", err)
	}
	history, err := api.GetTaskAttempt(ctx, tasks[0].ID, first.AttemptID)
	if err != nil || history.ConcurrencyToken != first.ConcurrencyToken {
		t.Fatal("restart changed key token", err)
	}
	finishID := workerTestID(t)
	finish := WorkerFinish{Ref: first.Attempt(), Outcome: WorkerSucceeded, Result: json.RawMessage(`{}`)}
	if _, err = w1.Finish(ctx, finishID, finish); err != nil {
		t.Fatal(err)
	}
	second, err := w2.NextAssignment(ctx)
	if err != nil || second.ConcurrencyToken <= first.ConcurrencyToken {
		t.Fatal("restart reset sequence", err)
	}
	if _, err = w2.Start(ctx, workerTestID(t), second.Attempt()); err != nil {
		t.Fatal(err)
	}
	// 第二次许可已经开始，旧完成回执的重放不能释放它。
	if out, err := w1.Finish(ctx, finishID, finish); err != nil || !out.Deduplicated {
		t.Fatal("duplicate finish changed ownership", err)
	}
	if _, err = w2.Renew(ctx, workerTestID(t), second.Attempt()); err != nil {
		t.Fatal(err)
	}
	if _, err = w2.Finish(ctx, workerTestID(t), WorkerFinish{Ref: second.Attempt(), Outcome: WorkerSucceeded, Result: json.RawMessage(`{}`)}); err != nil {
		t.Fatal(err)
	}
	for _, task := range tasks {
		waitRuntimeState(t, ctx, api, task.ID, TaskSucceeded)
	}
}