package mysql

import (
	"context"
	"encoding/json"
	"os"
	"sync"
	"testing"
	"time"

	"git.yuntuo.hk/k2safe/orbit/internal/fault"
	"git.yuntuo.hk/k2safe/orbit/platform/auth"
	pulse "git.yuntuo.hk/k2safe/orbit/pulse/model"
	"git.yuntuo.hk/k2safe/orbit/pulse/service"
	"git.yuntuo.hk/k2safe/orbit/pulse/worker"
)

func configureResources(t *testing.T, f *executionFixture, rate, concurrency, unfinished uint32, paused bool) service.QueueView {
	t.Helper()
	q, err := f.tasks.GetQueue(f.ctx, f.token, f.resources.QueueID)
	if err != nil {
		t.Fatal(err)
	}
	out, err := f.tasks.ConfigureQueue(f.ctx, f.token, q.ID, string(executionID(t)), service.QueueConfigInput{ExpectedRevision: q.Revision, Reason: "controlled resource regression", Paused: paused, MaxConcurrency: concurrency, RatePerSecond: rate, MaxUnfinishedTasks: unfinished})
	if err != nil {
		t.Fatal(err)
	}
	return out.Queue
}
func assertQueueCount(t *testing.T, f *executionFixture, expected uint64) {
	t.Helper()
	q, err := f.tasks.GetQueue(f.ctx, f.token, f.resources.QueueID)
	if err != nil || q.UnfinishedTasks != expected {
		t.Fatalf("queue count=%d want=%d err=%v", q.UnfinishedTasks, expected, err)
	}
	actual := f.count("pulse_tasks", "queue_id=? AND state IN ('WAITING','LEASED','RUNNING','NEEDS_REVIEW')", q.ID)
	if uint64(actual) != expected {
		t.Fatalf("counter differs from actual states: %d", actual)
	}
}
func waitStartWindow(t *testing.T, f *executionFixture) {
	t.Helper()
	var until, now time.Time
	if err := f.db.QueryRowContext(f.ctx, `SELECT start_window_until,UTC_TIMESTAMP(6) FROM pulse_queues WHERE namespace_id=? AND id=?`, f.resources.NamespaceID, f.resources.QueueID).Scan(&until, &now); err != nil {
		t.Fatal(err)
	}
	timer := time.NewTimer(max(time.Millisecond, until.Sub(now)+10*time.Millisecond))
	defer timer.Stop()
	select {
	case <-timer.C:
	case <-f.ctx.Done():
		t.Fatal(f.ctx.Err())
	}
}
func TestIntegrationQueueResources(t *testing.T) {
	dsn := os.Getenv("ORBIT_TEST_MYSQL_DSN")
	if dsn == "" {
		t.Skip("requires disposable MySQL")
	}
	ctx, stop := context.WithTimeout(context.Background(), 60*time.Second)
	defer stop()
	db, err := Open(ctx, dsn, 32, 8)
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	if err := Migrate(ctx, db); err != nil {
		t.Fatal(err)
	}
	fixture := func(t *testing.T) *executionFixture { return newExecutionFixture(t, ctx, db) }
	t.Run("concurrent_nodes_share_start_budget", func(t *testing.T) {
		f := fixture(t)
		configureResources(t, f, 4, 4, 10, false)
		hello := f.hello
		hello.InstanceID = executionID(t)
		b, _, err := f.b.Connect(ctx, f.token, hello)
		if err != nil {
			t.Fatal(err)
		}
		type item struct {
			api     *worker.Service
			binding worker.Binding
			ref     worker.AttemptRef
		}
		var items []item
		for i := 0; i < 4; i++ {
			f.submit(1)
			api, binding := f.a, f.binding
			if i%2 == 1 {
				api, binding = f.b, b
			}
			a, err := api.Claim(ctx, f.token, binding)
			if err != nil || a == nil {
				t.Fatal("claim reservation failed", err)
			}
			items = append(items, item{api, binding, worker.AttemptRef{TaskID: a.Task.ID, AttemptID: a.AttemptID, Epoch: a.Epoch, LeaseSecret: a.LeaseSecret}})
		}
		configureResources(t, f, 2, 4, 10, false)
		start := make(chan struct{})
		results := make(chan error, 4)
		var wg sync.WaitGroup
		for _, item := range items {
			message := executionID(t)
			wg.Add(1)
			go func() {
				defer wg.Done()
				<-start
				_, err := item.api.Start(ctx, f.token, item.binding, message, item.ref)
				results <- err
			}()
		}
		close(start)
		wg.Wait()
		close(results)
		allowed := 0
		for err := range results {
			if err == nil {
				allowed++
			} else if fault.Code(err) != "RATE_LIMITED" {
				t.Fatal(err)
			}
		}
		if allowed != 2 {
			t.Fatalf("concurrent starts=%d, want 2", allowed)
		}
		assertQueueCount(t, f, 4)
	})
	t.Run("start_budget_rollback_and_configuration_scope", func(t *testing.T) {
		f := fixture(t)
		q := configureResources(t, f, 1, 4, 10, false)
		f.submit(1)
		_, ref := f.claim()
		message := executionID(t)
		f.exec(`CREATE TRIGGER orbit_start_budget_fail BEFORE INSERT ON orbit_events FOR EACH ROW SIGNAL SQLSTATE '45000' SET MESSAGE_TEXT='controlled rollback'`)
		_, err := f.a.Start(ctx, f.token, f.binding, message, ref)
		f.exec(`DROP TRIGGER orbit_start_budget_fail`)
		if err == nil || f.get(ref.TaskID).State != pulse.TaskLeased {
			t.Fatal("failed start committed")
		}
		var used uint32
		if err := db.QueryRowContext(ctx, `SELECT starts_in_window FROM pulse_queues WHERE namespace_id=? AND id=?`, f.resources.NamespaceID, q.ID).Scan(&used); err != nil || used != 0 {
			t.Fatal("rolled back start spent budget")
		}
		if _, err := f.a.Start(ctx, f.token, f.binding, message, ref); err != nil {
			t.Fatal(err)
		}
		other := fixture(t)
		in := service.QueueConfigInput{ExpectedRevision: q.Revision, Reason: "scope", MaxConcurrency: 4, RatePerSecond: 1, MaxUnfinishedTasks: 10}
		_, err = other.tasks.ConfigureQueue(ctx, other.token, q.ID, "cross-namespace", in)
		expectWorkerCode(t, err, "NOT_FOUND")
		raw, _ := json.Marshal(auth.Permissions{Actions: []string{"queues:manage"}, Queues: []string{string(other.resources.QueueID)}})
		f.exec(`UPDATE orbit_credentials SET permissions=? WHERE namespace_id=? AND id=?`, string(raw), f.resources.NamespaceID, f.resources.CredentialID)
		_, err = f.tasks.ConfigureQueue(ctx, f.token, q.ID, "denied-queue", in)
		expectWorkerCode(t, err, "NOT_FOUND")
		if f.count("orbit_idempotency_keys", "operation='queue.configure' AND scope_key=?", string(q.ID)) != 1 {
			t.Fatal("denied config left identity")
		}
	})
	t.Run("concurrent_admission_duplicate_and_release", func(t *testing.T) {
		f := fixture(t)
		configureResources(t, f, 10, 4, 3, false)
		input := service.SubmitInput{QueueID: f.resources.QueueID, Handler: "object.prepare", HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{}`)}
		type result struct {
			key string
			out service.SubmitResult
			err error
		}
		results := make(chan result, 16)
		var wg sync.WaitGroup
		for i := 0; i < 16; i++ {
			key := string(executionID(t))
			wg.Add(1)
			go func() {
				defer wg.Done()
				out, err := f.tasks.Submit(ctx, f.token, key, input)
				results <- result{key, out, err}
			}()
		}
		wg.Wait()
		close(results)
		var accepted []result
		for r := range results {
			if r.err == nil {
				accepted = append(accepted, r)
			} else {
				expectWorkerCode(t, r.err, "QUEUE_FULL")
			}
		}
		if len(accepted) != 3 || f.count("orbit_idempotency_keys", "operation='task.submit'") != 3 {
			t.Fatal("admission exceeded limit or left rejected identities")
		}
		assertQueueCount(t, f, 3)
		again, err := f.tasks.Submit(ctx, f.token, accepted[0].key, input)
		if err != nil || !again.Deduplicated || again.Task.ID != accepted[0].out.Task.ID {
			t.Fatal("full queue rejected original receipt")
		}
		_, err = f.tasks.CancelTask(ctx, f.token, again.Task.ID, "release", service.CancelInput{ExpectedRevision: again.Task.Revision, Reason: "release slot"})
		if err != nil {
			t.Fatal(err)
		}
		assertQueueCount(t, f, 2)
		if _, err := f.tasks.Submit(ctx, f.token, "new-after-release", input); err != nil {
			t.Fatal(err)
		}
		assertQueueCount(t, f, 3)
	})
	t.Run("configuration_cas_audit_and_pause", func(t *testing.T) {
		f := fixture(t)
		q := configureResources(t, f, 10, 4, 10, true)
		task := f.submit(1)
		if a, err := f.a.Claim(ctx, f.token, f.binding); err != nil || a != nil {
			t.Fatal("paused queue dispatched")
		}
		in := service.QueueConfigInput{ExpectedRevision: q.Revision, Reason: "resume", MaxConcurrency: 4, RatePerSecond: 10, MaxUnfinishedTasks: 10}
		first, err := f.tasks.ConfigureQueue(ctx, f.token, q.ID, "resume", in)
		if err != nil {
			t.Fatal(err)
		}
		again, err := f.tasks.ConfigureQueue(ctx, f.token, q.ID, "resume", in)
		if err != nil || !again.Deduplicated || again.Queue.Revision != first.Queue.Revision {
			t.Fatal("config receipt changed")
		}
		_, err = f.tasks.ConfigureQueue(ctx, f.token, q.ID, "stale", in)
		expectWorkerCode(t, err, "STATE_CONFLICT")
		_, ref := f.claim()
		configureResources(t, f, 10, 4, 10, true)
		_, err = f.a.Start(ctx, f.token, f.binding, executionID(t), ref)
		expectWorkerCode(t, err, "EXECUTION_PAUSED")
		_, err = f.tasks.CancelTask(ctx, f.token, task.ID, "cancel-paused", service.CancelInput{ExpectedRevision: f.get(task.ID).Revision, Reason: "cancel while paused"})
		if err != nil {
			t.Fatal(err)
		}
		assertQueueCount(t, f, 0)
		if f.count("orbit_audit_events", "resource_id=? AND action='queue.configure'", q.ID) != 3 {
			t.Fatal("configuration audit count")
		}
	})
	t.Run("rate_window_survives_reconnect_and_receipt_replay", func(t *testing.T) {
		f := fixture(t)
		configureResources(t, f, 1, 4, 10, false)
		f.submit(1)
		f.submit(1)
		_, first := f.claim()
		if a, err := f.a.Claim(ctx, f.token, f.binding); err != nil || a != nil {
			t.Fatal("unstarted reservation did not throttle claims")
		}
		message := executionID(t)
		if _, err := f.a.Start(ctx, f.token, f.binding, message, first); err != nil {
			t.Fatal(err)
		}
		if out, err := f.a.Start(ctx, f.token, f.binding, message, first); err != nil || !out.Deduplicated {
			t.Fatal("start replay rejected by quota")
		}
		binding, _, err := f.b.Connect(ctx, f.token, f.hello)
		if err != nil {
			t.Fatal(err)
		}
		if _, err := f.b.Resume(ctx, f.token, binding, first); err != nil {
			t.Fatal(err)
		}
		f.a, f.binding = f.b, binding
		if a, err := f.a.Claim(ctx, f.token, f.binding); err != nil || a != nil {
			t.Fatal("node handoff reset rate window")
		}
		waitStartWindow(t, f)
		_, second := f.claim()
		if _, err := f.a.Start(ctx, f.token, f.binding, executionID(t), second); err != nil {
			t.Fatal(err)
		}
		var gap int64
		if err := db.QueryRowContext(ctx, `SELECT TIMESTAMPDIFF(MICROSECOND,a.authorized_at,b.authorized_at) FROM pulse_attempts a JOIN pulse_attempts b ON b.namespace_id=a.namespace_id WHERE a.namespace_id=? AND a.id=? AND b.id=?`, f.resources.NamespaceID, first.AttemptID, second.AttemptID).Scan(&gap); err != nil || gap < 1000000 {
			t.Fatalf("rate gap=%d err=%v", gap, err)
		}
		retryFinish(t, f, first, "SUCCEEDED", false)
		retryFinish(t, f, second, "SUCCEEDED", false)
		assertQueueCount(t, f, 0)
	})
	t.Run("lower_limits_preserve_inflight_and_throttled_start", func(t *testing.T) {
		f := fixture(t)
		configureResources(t, f, 2, 4, 10, false)
		f.submit(1)
		f.submit(1)
		f.submit(1)
		_, a := f.claim()
		_, b := f.claim()
		configureResources(t, f, 1, 1, 1, false)
		if _, err := f.a.Start(ctx, f.token, f.binding, executionID(t), a); err != nil {
			t.Fatal(err)
		}
		message := executionID(t)
		_, err := f.a.Start(ctx, f.token, f.binding, message, b)
		expectWorkerCode(t, err, "RATE_LIMITED")
		if f.get(b.TaskID).State != pulse.TaskLeased || f.get(b.TaskID).AttemptCount != 1 || f.used() != 2 {
			t.Fatal("rate rejection consumed or revoked work")
		}
		retryFinish(t, f, a, "SUCCEEDED", false)
		if got, err := f.a.Claim(ctx, f.token, f.binding); err != nil || got != nil {
			t.Fatal("lowered concurrency admitted excess")
		}
		waitStartWindow(t, f)
		if _, err := f.a.Start(ctx, f.token, f.binding, message, b); err != nil {
			t.Fatal(err)
		}
		retryFinish(t, f, b, "SUCCEEDED", false)
		assertQueueCount(t, f, 1)
		waitStartWindow(t, f)
		f.claim()
	})
	t.Run("rollback_retry_unknown_and_expiry_accounting", func(t *testing.T) {
		f := fixture(t)
		configureResources(t, f, 10, 4, 2, false)
		f.exec(`CREATE TRIGGER orbit_resource_test_fail BEFORE INSERT ON orbit_events FOR EACH ROW SIGNAL SQLSTATE '45000' SET MESSAGE_TEXT='controlled rollback'`)
		_, err := f.tasks.Submit(ctx, f.token, "rollback", service.SubmitInput{QueueID: f.resources.QueueID, Handler: "object.prepare", HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{}`)})
		f.exec(`DROP TRIGGER orbit_resource_test_fail`)
		if err == nil {
			t.Fatal("injected failure ignored")
		}
		assertQueueCount(t, f, 0)
		task := retryTask(t, f, 2, pulse.RecoveryIdempotent)
		ref := retryStart(t, f)
		retryFinish(t, f, ref, "FAILED", true)
		assertQueueCount(t, f, 1)
		makeRetryDue(f, task.ID)
		ref = retryStart(t, f)
		retryFinish(t, f, ref, "UNKNOWN", false)
		assertQueueCount(t, f, 1)
		waiting := f.submit(1)
		f.exec(`UPDATE pulse_tasks SET available_at=DATE_SUB(UTC_TIMESTAMP(6),INTERVAL 2 SECOND),deadline_at=DATE_SUB(UTC_TIMESTAMP(6),INTERVAL 1 SECOND) WHERE namespace_id=? AND id=?`, f.resources.NamespaceID, waiting.ID)
		if _, err := f.storeA.Recover(ctx, 100); err != nil {
			t.Fatal(err)
		}
		assertQueueCount(t, f, 1)
		if _, err := f.storeA.Recover(ctx, 100); err != nil {
			t.Fatal(err)
		}
		assertQueueCount(t, f, 1)
	})
	t.Run("schedule_backpressure_preserves_unaccepted_cursor", func(t *testing.T) {
		f := fixture(t)
		configureResources(t, f, 10, 4, 1, false)
		sf := newScheduleFixture(t, f)
		v := sf.create(sf.definition())
		sf.due(v, 2*time.Second)
		sf.tick()
		assertQueueCount(t, f, 1)
		items := sf.occurrences(v.ID)
		if len(items) != 1 || items[0].TaskID == nil {
			t.Fatal("schedule exceeded backlog limit")
		}
		next := sf.current(v.ID).NextFireAt
		if next == nil || !next.Equal(items[0].ScheduledAt.Add(time.Second)) {
			t.Fatal("backpressure skipped unaccepted point")
		}
		id := *items[0].TaskID
		_, err := f.tasks.CancelTask(ctx, f.token, id, "free-schedule", service.CancelInput{ExpectedRevision: 1, Reason: "free slot"})
		if err != nil {
			t.Fatal(err)
		}
		sf.due(v, 0)
		sf.tick()
		items = sf.occurrences(v.ID)
		if len(items) != 2 {
			t.Fatal("schedule did not resume after release")
		}
		found := false
		for _, item := range items {
			if item.ScheduledAt.Equal(*next) {
				found = true
			}
		}
		if !found {
			t.Fatal("resumed wrong original time")
		}
		assertQueueCount(t, f, 1)
	})
}