package orbit

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/coder/websocket"
)

// 由 scripts/test-mysql.sh 启动真实服务与隔离 MySQL；SDK 不导入服务端内部包。
func integrationWorkerClient(t *testing.T) (*Client, string, string) {
	t.Helper()
	path := os.Getenv("ORBIT_TEST_WORKER_CREDENTIALS_FILE")
	base := os.Getenv("ORBIT_TEST_WORKER_URL")
	if path == "" || base == "" {
		t.Skip("requires isolated Orbit server and Worker credentials")
	}
	if !strings.HasPrefix(base, "http://127.0.0.1:") {
		t.Fatal("Worker integration requires an isolated loopback server")
	}
	raw, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	var credentials struct {
		Token   string `json:"token"`
		QueueID string `json:"queue_id"`
	}
	if json.Unmarshal(raw, &credentials) != nil || credentials.Token == "" || !validID(credentials.QueueID) {
		t.Fatal("invalid integration credentials")
	}
	c, err := NewClient(base, Options{Token: credentials.Token})
	if err != nil {
		t.Fatal(err)
	}
	return c, credentials.QueueID, credentials.Token
}

// workerFaultGateway 在实际服务端确认后丢弃指定 ACK；SDK 到网关使用真实 TLS。
// 只丢响应，不伪造任务状态，测试凭据不进入 URL 或输出。
type workerFaultGateway struct {
	drop    atomic.Value
	offline atomic.Bool
	mu      sync.Mutex
	active  map[*websocket.Conn]struct{}
	dropped chan string
	server  *httptest.Server
}

func newWorkerFaultGateway(t *testing.T, upstream string) *workerFaultGateway {
	t.Helper()
	g := &workerFaultGateway{dropped: make(chan string, 4), active: make(map[*websocket.Conn]struct{})}
	g.drop.Store("")
	ctx, cancel := context.WithCancel(context.Background())
	var handlers sync.WaitGroup
	g.server = httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if g.offline.Load() {
			http.Error(w, "gateway offline", 503)
			return
		}
		handlers.Add(1)
		defer handlers.Done()
		bridge, stop := context.WithCancel(ctx)
		defer stop()
		header := http.Header{}
		header.Set("Authorization", r.Header.Get("Authorization"))
		up, _, err := websocket.Dial(bridge, upstream+"/v1/worker/connect", &websocket.DialOptions{HTTPHeader: header, Subprotocols: []string{workerSubprotocol}})
		if err != nil {
			http.Error(w, "upstream unavailable", 503)
			return
		}
		defer up.CloseNow()
		up.SetReadLimit(512 << 10)
		down, err := websocket.Accept(w, r, &websocket.AcceptOptions{Subprotocols: []string{workerSubprotocol}})
		if err != nil {
			return
		}
		defer down.CloseNow()
		g.mu.Lock()
		g.active[down] = struct{}{}
		g.mu.Unlock()
		defer func() { g.mu.Lock(); delete(g.active, down); g.mu.Unlock() }()
		if g.offline.Load() {
			return
		}
		down.SetReadLimit(512 << 10)
		var mu sync.Mutex
		var forwarding sync.WaitGroup
		forwarding.Add(2)
		operations := map[string]string{}
		finished := make(chan struct{}, 1)
		go func() {
			defer forwarding.Done()
			defer func() { finished <- struct{}{} }()
			for {
				kind, raw, err := down.Read(bridge)
				if err != nil {
					return
				}
				var f workerFrame
				if json.Unmarshal(raw, &f) != nil {
					return
				}
				mu.Lock()
				operations[f.MessageID] = f.Type
				mu.Unlock()
				if up.Write(bridge, kind, raw) != nil {
					return
				}
			}
		}()
		go func() {
			defer forwarding.Done()
			defer stop()
			for {
				kind, raw, err := up.Read(bridge)
				if err != nil {
					return
				}
				var f workerFrame
				if json.Unmarshal(raw, &f) != nil {
					return
				}
				mu.Lock()
				operation := operations[f.ReplyTo]
				delete(operations, f.ReplyTo)
				mu.Unlock()
				if f.Type == "ack" && operation != "" && g.drop.CompareAndSwap(operation, "") {
					g.dropped <- operation
					return
				}
				if down.Write(bridge, kind, raw) != nil {
					return
				}
			}
		}()
		select {
		case <-bridge.Done():
		case <-finished:
		}
		stop()
		down.CloseNow()
		up.CloseNow()
		forwarding.Wait()
	}))
	t.Cleanup(func() { cancel(); g.server.Close(); handlers.Wait() })
	return g
}
func (g *workerFaultGateway) disconnectActive() {
	g.mu.Lock()
	active := make([]*websocket.Conn, 0, len(g.active))
	for c := range g.active {
		active = append(active, c)
	}
	g.mu.Unlock()
	for _, c := range active {
		c.CloseNow()
	}
}
func (g *workerFaultGateway) expectDrop(t *testing.T, operation string) {
	t.Helper()
	select {
	case got := <-g.dropped:
		if got != operation {
			t.Fatal("wrong dropped operation")
		}
	case <-time.After(2 * time.Second):
		t.Fatal("ACK was not dropped after upstream confirmation")
	}
}

func TestIntegrationWorkerConnectionACKLoss(t *testing.T) {
	api, queue, token := integrationWorkerClient(t)
	ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second)
	defer cancel()
	r, err := NewWorkerRegistration(queue, 1, []WorkerCapability{{Handler: "sdk.ack-loss", HandlerVersion: "v1", PayloadVersion: "v1"}})
	if err != nil {
		t.Fatal(err)
	}
	g := newWorkerFaultGateway(t, api.baseURL)
	client, err := NewClient(g.server.URL, Options{Token: token, HTTPClient: g.server.Client()})
	if err != nil {
		t.Fatal(err)
	}
	dial := func() *WorkerConnection {
		t.Helper()
		w, err := client.DialWorker(ctx, r)
		if err != nil {
			t.Fatal(err)
		}
		t.Cleanup(func() { w.Close() })
		return w
	}
	accepted, err := api.SubmitTask(ctx, "sdk-ack-"+workerTestID(t), SubmitTaskRequest{QueueID: queue, Handler: r.Capabilities[0].Handler, HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{"value":9007199254740993}`), TaskTimeoutMS: 60000, Retry: RetryPolicy{MaxAttempts: 1, AttemptTimeoutMS: 30000}})
	if err != nil {
		t.Fatal(err)
	}
	w := dial()
	if _, err := w.Capacity(ctx, workerTestID(t), WorkerCapacity{Sequence: 1, Capacity: 1}); err != nil {
		t.Fatal(err)
	}
	a, err := w.NextAssignment(ctx)
	if err != nil {
		t.Fatal(err)
	}
	if a.Task.ID != accepted.Task.ID {
		t.Fatal("wrong assigned task")
	}
	startID := workerTestID(t)
	g.drop.Store("task.start")
	if _, err := w.Start(ctx, startID, a.Attempt()); err == nil {
		t.Fatal("lost start ACK reported success")
	}
	g.expectDrop(t, "task.start")
	w.Close()
	w = dial()
	lease, err := w.Resume(ctx, workerTestID(t), a.Attempt())
	if err != nil {
		t.Fatal(err)
	}
	if _, active := lease.ExecutionDeadline(); !active {
		t.Fatal("committed start did not resume")
	}
	started, err := w.Start(ctx, startID, a.Attempt())
	if err != nil {
		t.Fatal(err)
	}
	if !started.Deduplicated {
		t.Fatal("start receipt was recreated")
	}
	if _, active := started.Lease.ExecutionDeadline(); !active {
		t.Fatal("resume missing current authority")
	}
	// 处理器仅在首次确认当前授权后运行。之后连接恢复只核对，不再次运行它。
	var handlerCalls int
	handlerCalls++
	if _, err := w.Capacity(ctx, workerTestID(t), WorkerCapacity{Sequence: 1, Capacity: 1}); err != nil {
		t.Fatal(err)
	}
	renewID := workerTestID(t)
	first, err := w.Renew(ctx, renewID, a.Attempt())
	if err != nil {
		t.Fatal(err)
	}
	second, err := w.Renew(ctx, renewID, a.Attempt())
	if err != nil {
		t.Fatal(err)
	}
	if first.Deduplicated || !second.Deduplicated || first.Receipt.TaskRevision != second.Receipt.TaskRevision || !first.Receipt.LeaseUntil.Equal(second.Receipt.LeaseUntil) {
		t.Fatal("duplicate renew changed original receipt")
	}
	finishID := workerTestID(t)
	result := json.RawMessage(`{"normalized":"hello","sequence":9007199254740993,"ratio":0.1234567890123456789012345}`)
	finish := WorkerFinish{Ref: a.Attempt(), Outcome: WorkerSucceeded, Result: result}
	g.drop.Store("task.finish")
	if _, err := w.Finish(ctx, finishID, finish); err == nil {
		t.Fatal("lost finish ACK reported success")
	}
	g.expectDrop(t, "task.finish")
	w.Close()
	w = dial()
	lease, err = w.Resume(ctx, workerTestID(t), a.Attempt())
	if err != nil {
		t.Fatal(err)
	}
	if lease.Active || lease.State != "SUCCEEDED" {
		t.Fatal("completed attempt resumed as executable")
	}
	confirmed, err := w.Finish(ctx, finishID, finish)
	if err != nil {
		t.Fatal(err)
	}
	if !confirmed.Deduplicated || confirmed.Receipt.Disposition != "APPLIED" || confirmed.Receipt.TaskState != TaskSucceeded {
		t.Fatal("missing durable finish receipt")
	}
	stale, err := w.Start(ctx, startID, a.Attempt())
	if err != nil {
		t.Fatal(err)
	}
	if !stale.Deduplicated {
		t.Fatal("missing old start receipt")
	}
	if _, active := stale.Lease.ExecutionDeadline(); active {
		t.Fatal("old start ACK resurrected execution")
	}
	conflict := finish
	conflict.Result = json.RawMessage(`{"normalized":"conflicting evidence"}`)
	evidence, err := w.Finish(ctx, workerTestID(t), conflict)
	if err != nil {
		t.Fatal(err)
	}
	if evidence.Receipt.Disposition != "EVIDENCE_ONLY" {
		t.Fatal("conflicting result replaced authority")
	}
	current, err := api.GetTask(ctx, a.Task.ID)
	if err != nil {
		t.Fatal(err)
	}
	if current.State != TaskSucceeded || current.AttemptCount != 1 || string(current.Result) != string(result) || handlerCalls != 1 {
		t.Fatal("result, precision or execution count changed")
	}
	drain, err := w.Capacity(ctx, workerTestID(t), WorkerCapacity{Sequence: 1, Capacity: 0, Drain: true})
	if err != nil || !drain.Draining {
		t.Fatalf("drain: %v", err)
	}
	w.Close()
	w = dial()
	if !w.Welcome().Draining {
		t.Fatal("drain did not persist across reconnect")
	}
	_, err = w.Capacity(ctx, workerTestID(t), WorkerCapacity{Sequence: 1, Capacity: 1})
	var control *WorkerError
	if !errors.As(err, &control) || control.Code != "STATE_CONFLICT" {
		t.Fatalf("clearing persistent drain was not rejected: %v", err)
	}
	drain, err = w.Capacity(ctx, workerTestID(t), WorkerCapacity{Sequence: 1, Capacity: 0, Drain: true})
	if err != nil || !drain.Draining {
		t.Fatalf("draining heartbeat failed: %v", err)
	}
}

func TestIntegrationWorkerConnectionExpiredLease(t *testing.T) {
	api, queue, _ := integrationWorkerClient(t)
	ctx, cancel := context.WithTimeout(context.Background(), 12*time.Second)
	defer cancel()
	r, err := NewWorkerRegistration(queue, 1, []WorkerCapability{{Handler: "sdk.expiry", HandlerVersion: "v1", PayloadVersion: "v1"}})
	if err != nil {
		t.Fatal(err)
	}
	_, err = api.SubmitTask(ctx, "sdk-expiry-"+workerTestID(t), SubmitTaskRequest{QueueID: queue, Handler: "sdk.expiry", HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{}`), TaskTimeoutMS: 30000, Retry: RetryPolicy{MaxAttempts: 1, AttemptTimeoutMS: 1100}})
	if err != nil {
		t.Fatal(err)
	}
	w, err := api.DialWorker(ctx, r)
	if err != nil {
		t.Fatal(err)
	}
	defer w.Close()
	if _, err := w.Capacity(ctx, workerTestID(t), WorkerCapacity{Sequence: 1, Capacity: 1}); err != nil {
		t.Fatal(err)
	}
	a, err := w.NextAssignment(ctx)
	if err != nil {
		t.Fatal(err)
	}
	startID := workerTestID(t)
	out, err := w.Start(ctx, startID, a.Attempt())
	if err != nil {
		t.Fatal(err)
	}
	deadline, active := out.Lease.ExecutionDeadline()
	if !active {
		t.Fatal("short valid lease was rejected")
	}
	// 用单调期限驱动协作取消；额外等待安全余量确保服务端授权也已过期。
	work, stop := context.WithDeadline(ctx, deadline)
	defer stop()
	<-work.Done()
	if !errors.Is(work.Err(), context.DeadlineExceeded) {
		t.Fatal(work.Err())
	}
	select {
	case <-ctx.Done():
		t.Fatal(ctx.Err())
	case <-time.After(350 * time.Millisecond):
	}
	stale, err := w.Start(ctx, startID, a.Attempt())
	if err != nil {
		t.Fatal(err)
	}
	if !stale.Deduplicated {
		t.Fatal("missing historical receipt")
	}
	if _, active := stale.Lease.ExecutionDeadline(); active {
		t.Fatal("expired historical ACK authorized work")
	}
	late, err := w.Finish(ctx, workerTestID(t), WorkerFinish{Ref: a.Attempt(), Outcome: WorkerUnknown, Result: json.RawMessage(`{"reason":"handler_cancelled"}`), ErrorCode: "LEASE_EXPIRED"})
	if err != nil {
		t.Fatal(err)
	}
	if late.Receipt.Disposition != "EVIDENCE_ONLY" {
		t.Fatal("late report changed authoritative result")
	}
	for {
		current, err := api.GetTask(ctx, a.Task.ID)
		if err != nil {
			t.Fatal(err)
		}
		if current.State == TaskNeedsReview {
			if current.AttemptCount != 1 || len(current.Result) != 0 {
				t.Fatal("expired execution was replayed or evidence promoted")
			}
			break
		}
		select {
		case <-ctx.Done():
			t.Fatal("server did not recover expired execution")
		case <-time.After(100 * time.Millisecond):
		}
	}
}