package webhooks

import (
	"context"
	"database/sql"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/google/uuid"

	"github.com/nextlevelbuilder/goclaw/internal/crypto"
	"github.com/nextlevelbuilder/goclaw/internal/providers"
	"github.com/nextlevelbuilder/goclaw/internal/security"
	"github.com/nextlevelbuilder/goclaw/internal/store"
)

// ---- stub implementations ----

// stubCallStore is an in-memory WebhookCallStore for unit tests.
// It records the last UpdateStatusCAS call for assertion.
type stubCallStore struct {
	calls       map[uuid.UUID]*store.WebhookCallData
	lastUpdate  map[string]any // last updates map passed to UpdateStatusCAS
	claimErr    error          // if non-nil, returned by ClaimNext
	reclaimN    int64          // count returned by ReclaimStale
	casLeaseErr error          // if non-nil, returned by UpdateStatusCAS
	hbCount     int32          // số lần Heartbeat được gọi (atomic)
	hbErr       error          // nếu non-nil, Heartbeat trả về lỗi này
}

func newStubCallStore(initial *store.WebhookCallData) *stubCallStore {
	s := &stubCallStore{
		calls:      make(map[uuid.UUID]*store.WebhookCallData),
		lastUpdate: nil,
	}
	if initial != nil {
		s.calls[initial.ID] = initial
	}
	return s
}

func (s *stubCallStore) Create(_ context.Context, call *store.WebhookCallData) error {
	s.calls[call.ID] = call
	return nil
}
func (s *stubCallStore) GetByID(_ context.Context, id uuid.UUID) (*store.WebhookCallData, error) {
	if c, ok := s.calls[id]; ok {
		return c, nil
	}
	return nil, sql.ErrNoRows
}
func (s *stubCallStore) GetByIdempotency(_ context.Context, _ uuid.UUID, _ string) (*store.WebhookCallData, error) {
	return nil, sql.ErrNoRows
}
func (s *stubCallStore) UpdateStatus(_ context.Context, id uuid.UUID, updates map[string]any) error {
	s.lastUpdate = updates
	if c, ok := s.calls[id]; ok {
		if st, ok := updates["status"].(string); ok {
			c.Status = st
		}
		if att, ok := updates["attempts"].(int); ok {
			c.Attempts = att
		}
	}
	return nil
}

// UpdateStatusCAS implements the K5 CAS guard. In tests it behaves like UpdateStatus
// unless casLeaseErr is set.
func (s *stubCallStore) UpdateStatusCAS(_ context.Context, id uuid.UUID, _ string, updates map[string]any) error {
	if s.casLeaseErr != nil {
		return s.casLeaseErr
	}
	s.lastUpdate = updates
	if c, ok := s.calls[id]; ok {
		if st, ok := updates["status"].(string); ok {
			c.Status = st
		}
		if att, ok := updates["attempts"].(int); ok {
			c.Attempts = att
		}
	}
	return nil
}

func (s *stubCallStore) ClaimNext(_ context.Context, _ uuid.UUID, _ time.Time) (*store.WebhookCallData, error) {
	if s.claimErr != nil {
		return nil, s.claimErr
	}
	return nil, sql.ErrNoRows
}
func (s *stubCallStore) List(_ context.Context, _ store.WebhookCallListFilter) ([]store.WebhookCallData, error) {
	return nil, nil
}
func (s *stubCallStore) Count(_ context.Context, _ store.WebhookCallListFilter) (int, error) {
	return 0, nil
}
func (s *stubCallStore) DeleteOlderThan(_ context.Context, _ uuid.UUID, _ time.Time) (int64, error) {
	return 0, nil
}
func (s *stubCallStore) ReclaimStale(_ context.Context, _ time.Time) (int64, error) {
	return s.reclaimN, nil
}
func (s *stubCallStore) Heartbeat(_ context.Context, _ uuid.UUID, _ string, _ time.Time) error {
	atomic.AddInt32(&s.hbCount, 1)
	return s.hbErr
}

// stubWebhookStore returns a fixed webhook on GetByID.
type stubWebhookStore struct {
	wh *store.WebhookData
}

func (s *stubWebhookStore) Create(_ context.Context, _ *store.WebhookData) error { return nil }
func (s *stubWebhookStore) GetByID(_ context.Context, _ uuid.UUID) (*store.WebhookData, error) {
	if s.wh == nil {
		return nil, sql.ErrNoRows
	}
	return s.wh, nil
}
func (s *stubWebhookStore) GetByHash(_ context.Context, _ string) (*store.WebhookData, error) {
	return nil, sql.ErrNoRows
}
func (s *stubWebhookStore) List(_ context.Context, _ store.WebhookListFilter) ([]store.WebhookData, error) {
	return nil, nil
}
func (s *stubWebhookStore) Count(_ context.Context, _ store.WebhookListFilter) (int, error) {
	return 0, nil
}
func (s *stubWebhookStore) Update(_ context.Context, _ uuid.UUID, _ map[string]any) error { return nil }
func (s *stubWebhookStore) RotateSecret(_ context.Context, _ uuid.UUID, _, _, _ string) error {
	return nil
}
func (s *stubWebhookStore) Revoke(_ context.Context, _ uuid.UUID) error        { return nil }
func (s *stubWebhookStore) TouchLastUsed(_ context.Context, _ uuid.UUID) error { return nil }
func (s *stubWebhookStore) GetByHashUnscoped(_ context.Context, _ string) (*store.WebhookData, error) {
	return nil, sql.ErrNoRows
}
func (s *stubWebhookStore) GetByIDUnscoped(_ context.Context, id uuid.UUID) (*store.WebhookData, error) {
	if s.wh != nil && s.wh.ID == id {
		return s.wh, nil
	}
	return nil, sql.ErrNoRows
}

// ---- helpers ----

// testEncKey is a 32-byte hex key used in tests for AES-256-GCM.
const testEncKey = "0102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f20"

// newTestCall creates a minimal async webhook_calls row for testing.
func newTestCall(callbackURL string, agentID *uuid.UUID) *store.WebhookCallData {
	now := time.Now()
	deliveryID := uuid.New()
	call := &store.WebhookCallData{
		ID:         uuid.New(),
		TenantID:   uuid.New(),
		WebhookID:  uuid.New(),
		AgentID:    agentID,
		DeliveryID: deliveryID,
		Mode:       "async",
		Status:     "running", // simulating ClaimNext already set it
		Attempts:   0,
		CreatedAt:  now,
		StartedAt:  &now,
	}
	cbURL := callbackURL
	call.CallbackURL = &cbURL

	// Encode minimal request payload.
	payload := asyncPayload{
		Input:       json.RawMessage(`"hello"`),
		CallbackURL: callbackURL,
	}
	b, _ := json.Marshal(payload)
	call.RequestPayload = b
	return call
}

func TestNewCallbackUsageMapsCacheTokens(t *testing.T) {
	if got := newCallbackUsage(nil); got != nil {
		t.Fatalf("nil usage must map to nil, got %+v", got)
	}
	u := &providers.Usage{
		PromptTokens:                      10,
		CompletionTokens:                  5,
		TotalTokens:                       15,
		CacheReadTokens:                   8,
		CacheCreationTokens:               2,
		PromptTokensIncludeCachedSegments: true,
	}
	got := newCallbackUsage(u)
	if got == nil {
		t.Fatal("expected non-nil callbackUsage")
	}
	if got.PromptTokens != 10 || got.CompletionTokens != 5 || got.TotalTokens != 15 {
		t.Errorf("base tokens wrong: %+v", got)
	}
	if got.CacheReadTokens != 8 || got.CacheCreationTokens != 2 {
		t.Errorf("cache tokens not mapped: read=%d create=%d", got.CacheReadTokens, got.CacheCreationTokens)
	}
	if !got.PromptTokensIncludeCachedSegments {
		t.Error("include-cached flag not mapped")
	}
}

func TestDecodeAsyncPayload_UnwrapsAuditEnvelope(t *testing.T) {
	meta := asyncPayload{
		Input:       json.RawMessage(`"hello"`),
		CallbackURL: "https://example.com/callback",
	}
	metaBytes, err := json.Marshal(meta)
	if err != nil {
		t.Fatalf("marshal meta: %v", err)
	}
	envelope, err := json.Marshal(map[string]any{
		"body_hash": "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
		"meta":      json.RawMessage(metaBytes),
	})
	if err != nil {
		t.Fatalf("marshal envelope: %v", err)
	}

	got, err := decodeAsyncPayload(envelope)
	if err != nil {
		t.Fatalf("decodeAsyncPayload: %v", err)
	}
	if string(got.Input) != `"hello"` {
		t.Fatalf("input = %s, want %s", got.Input, `"hello"`)
	}
	if got.CallbackURL != meta.CallbackURL {
		t.Fatalf("callback_url = %q, want %q", got.CallbackURL, meta.CallbackURL)
	}
}

// newTestWebhook creates a webhook with an encrypted raw secret.
// Returns the webhook and the raw secret bytes for signature verification.
// encKey is the AES-256-GCM key (same as testEncKey).
func newTestWebhook(id uuid.UUID, encKey string) (*store.WebhookData, []byte) {
	rawSecret := make([]byte, 32)
	for i := range rawSecret {
		rawSecret[i] = byte(i)
	}
	enc, err := crypto.Encrypt(string(rawSecret), encKey)
	if err != nil {
		panic("newTestWebhook: encrypt failed: " + err.Error())
	}
	return &store.WebhookData{
		ID:              id,
		EncryptedSecret: enc,
	}, rawSecret
}

// newTestWorker builds a worker wired with stub stores (no agent router needed for
// tests that don't invoke agent).
func newTestWorker(calls *stubCallStore, webhooks *stubWebhookStore) *WebhookWorker {
	return &WebhookWorker{
		calls:    calls,
		webhooks: webhooks,
		router:   nil, // nil OK when Response is pre-populated
		limiter:  NewCallbackLimiter(4),
		cfg:      WorkerConfig{WorkerConcurrency: 1, PerTenantConcurrency: 4},
		encKey:   testEncKey,
	}
}

// ---- tests ----

// TestHMACHeaderPresent verifies X-Webhook-Signature and X-Webhook-Delivery-Id
// are present and correctly signed on the outbound POST.
func TestHMACHeaderPresent(t *testing.T) {
	security.SetAllowLoopbackForTest(true)
	defer security.SetAllowLoopbackForTest(false)

	var gotSig, gotDelivery string
	var gotBody []byte

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotSig = r.Header.Get("X-Webhook-Signature")
		gotDelivery = r.Header.Get("X-Webhook-Delivery-Id")
		gotBody, _ = io.ReadAll(r.Body)
		w.WriteHeader(http.StatusOK)
	}))
	defer srv.Close()

	agentID := uuid.New()
	call := newTestCall(srv.URL, &agentID)
	// Pre-populate response so agent invocation is skipped.
	prevResp, _ := json.Marshal(callbackPayload{Output: "test output"})
	call.Response = prevResp

	wh, rawSecret := newTestWebhook(call.WebhookID, testEncKey)
	callStore := newStubCallStore(call)
	whStore := &stubWebhookStore{wh: wh}

	w := newTestWorker(callStore, whStore)
	w.execute(context.Background(), call, call.TenantID, "test-lease")

	if gotSig == "" {
		t.Fatal("X-Webhook-Signature header missing")
	}
	if !startsWith(gotSig, "t=") {
		t.Errorf("unexpected signature format: %q", gotSig)
	}
	if gotDelivery != call.DeliveryID.String() {
		t.Errorf("delivery_id: got %q want %q", gotDelivery, call.DeliveryID.String())
	}

	// Verify signature is valid using Sign() with the raw secret.
	var ts int64
	for _, part := range splitComma(gotSig) {
		if len(part) > 2 && part[:2] == "t=" {
			ts = parseInt64(part[2:])
		}
	}
	if ts == 0 {
		t.Fatal("could not parse t= from signature header")
	}
	expected := Sign(rawSecret, ts, gotBody)
	if gotSig != expected {
		t.Errorf("HMAC mismatch\ngot:  %s\nwant: %s", gotSig, expected)
	}
}

// TestExecuteReDeliveryCarriesCallsBreakdown verifies that when a call already has
// a stored callbackPayload (call.Response set, simulating a re-delivery/retry), the
// re-sent HTTP callback body carries forward the per-call breakdown (`calls`) and
// `total_cost_usd` fields from the previously-stored payload.
func TestExecuteReDeliveryCarriesCallsBreakdown(t *testing.T) {
	security.SetAllowLoopbackForTest(true)
	defer security.SetAllowLoopbackForTest(false)

	var gotBody []byte

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotBody, _ = io.ReadAll(r.Body)
		w.WriteHeader(http.StatusOK)
	}))
	defer srv.Close()

	agentID := uuid.New()
	call := newTestCall(srv.URL, &agentID)
	// Pre-populate response so agent invocation is skipped and the re-delivery
	// path (execute's "else if len(call.Response) > 0" branch) is exercised.
	prevResp, _ := json.Marshal(callbackPayload{
		Output: "prior output",
		Status: "done",
		Calls: []providers.CallUsage{
			{
				Type:     "tool_call",
				Name:     "read_image",
				Provider: "9router",
				Model:    "cx/gpt-5.5",
				Usage:    providers.Usage{PromptTokens: 100, TotalTokens: 100},
				CostUSD:  0.02,
			},
		},
		TotalCostUSD: 0.02,
	})
	call.Response = prevResp

	wh, _ := newTestWebhook(call.WebhookID, testEncKey)
	callStore := newStubCallStore(call)
	whStore := &stubWebhookStore{wh: wh}

	w := newTestWorker(callStore, whStore)
	w.execute(context.Background(), call, call.TenantID, "test-lease")

	if len(gotBody) == 0 {
		t.Fatal("no callback body captured — request never reached server")
	}

	var payload callbackPayload
	if err := json.Unmarshal(gotBody, &payload); err != nil {
		t.Fatalf("unmarshal captured body: %v\nbody: %s", err, gotBody)
	}

	if len(payload.Calls) != 1 {
		t.Fatalf("payload.Calls length: got %d, want 1 (body: %s)", len(payload.Calls), gotBody)
	}
	if payload.Calls[0].Name != "read_image" {
		t.Errorf("payload.Calls[0].Name: got %q, want %q", payload.Calls[0].Name, "read_image")
	}
	if payload.TotalCostUSD < 0.0199 || payload.TotalCostUSD > 0.0201 {
		t.Errorf("payload.TotalCostUSD: got %f, want ~0.02", payload.TotalCostUSD)
	}

	// Belt-and-suspenders: also confirm the raw JSON carries the expected keys.
	bodyStr := string(gotBody)
	for _, want := range []string{`"calls"`, `"read_image"`, `"total_cost_usd"`} {
		if !strings.Contains(bodyStr, want) {
			t.Errorf("callback body missing %s: %s", want, bodyStr)
		}
	}
}

// TestDeliveryIDStableAcrossRetries verifies same delivery_id sent on attempt 1 and 3.
func TestDeliveryIDStableAcrossRetries(t *testing.T) {
	security.SetAllowLoopbackForTest(true)
	defer security.SetAllowLoopbackForTest(false)

	var deliveries []string
	var attempt int32

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		deliveries = append(deliveries, r.Header.Get("X-Webhook-Delivery-Id"))
		n := atomic.AddInt32(&attempt, 1)
		if n < 3 {
			w.WriteHeader(http.StatusInternalServerError)
		} else {
			w.WriteHeader(http.StatusOK)
		}
	}))
	defer srv.Close()

	agentID := uuid.New()
	call := newTestCall(srv.URL, &agentID)
	prevResp, _ := json.Marshal(callbackPayload{Output: "output"})
	call.Response = prevResp

	wh, _ := newTestWebhook(call.WebhookID, testEncKey)
	callStore := newStubCallStore(call)
	whStore := &stubWebhookStore{wh: wh}
	w := newTestWorker(callStore, whStore)

	// Simulate 3 execute calls (retries) — each must send same delivery_id.
	deliveryID := call.DeliveryID
	for range 3 {
		w.execute(context.Background(), call, call.TenantID, "test-lease")
	}

	if len(deliveries) != 3 {
		t.Fatalf("expected 3 delivery attempts, got %d", len(deliveries))
	}
	for i, d := range deliveries {
		if d != deliveryID.String() {
			t.Errorf("attempt %d: delivery_id %q != %q", i+1, d, deliveryID.String())
		}
	}
}

// TestAttemptsIncrementPostSend verifies attempts is NOT set during ClaimNext
// but IS incremented after send completes.
func TestAttemptsIncrementPostSend(t *testing.T) {
	security.SetAllowLoopbackForTest(true)
	defer security.SetAllowLoopbackForTest(false)

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	defer srv.Close()

	agentID := uuid.New()
	call := newTestCall(srv.URL, &agentID)
	call.Attempts = 0 // as set by ClaimNext — NOT incremented
	prevResp, _ := json.Marshal(callbackPayload{Output: "output"})
	call.Response = prevResp

	wh, _ := newTestWebhook(call.WebhookID, testEncKey)
	callStore := newStubCallStore(call)
	whStore := &stubWebhookStore{wh: wh}
	w := newTestWorker(callStore, whStore)

	w.execute(context.Background(), call, call.TenantID, "test-lease")

	// UpdateStatusCAS should have been called with attempts=1.
	if callStore.lastUpdate == nil {
		t.Fatal("UpdateStatusCAS never called")
	}
	gotAttempts, _ := callStore.lastUpdate["attempts"].(int)
	if gotAttempts != 1 {
		t.Errorf("attempts after send: got %d, want 1", gotAttempts)
	}
	gotStatus, _ := callStore.lastUpdate["status"].(string)
	if gotStatus != "done" {
		t.Errorf("status after 200: got %q, want done", gotStatus)
	}
}

// TestSSRFBlockedCallback verifies a private-IP callback_url leads to status=failed.
func TestSSRFBlockedCallback(t *testing.T) {
	// Do NOT enable loopback bypass — private IPs must be blocked.
	agentID := uuid.New()
	call := newTestCall("http://192.168.1.1/callback", &agentID)
	prevResp, _ := json.Marshal(callbackPayload{Output: "output"})
	call.Response = prevResp

	wh, _ := newTestWebhook(call.WebhookID, testEncKey)
	callStore := newStubCallStore(call)
	whStore := &stubWebhookStore{wh: wh}
	w := newTestWorker(callStore, whStore)

	w.execute(context.Background(), call, call.TenantID, "test-lease")

	if callStore.lastUpdate == nil {
		t.Fatal("UpdateStatusCAS never called for SSRF-blocked URL")
	}
	gotStatus, _ := callStore.lastUpdate["status"].(string)
	if gotStatus != "failed" {
		t.Errorf("SSRF-blocked URL: status=%q, want failed", gotStatus)
	}
}

// TestBackoffSchedule verifies the delay table values and jitter bounds.
func TestBackoffSchedule(t *testing.T) {
	cases := []struct {
		attempt int
		minDur  time.Duration
		maxDur  time.Duration
	}{
		{0, 27 * time.Second, 33 * time.Second},    // 30s ±10%
		{1, 108 * time.Second, 132 * time.Second},  // 2m ±10%
		{2, 9 * time.Minute, 11 * time.Minute},     // 10m ±10%
		{3, 54 * time.Minute, 66 * time.Minute},    // 1h ±10%
		{4, 324 * time.Minute, 396 * time.Minute},  // 6h ±10%
		{99, 324 * time.Minute, 396 * time.Minute}, // capped at 6h
	}
	for _, tc := range cases {
		for range 50 { // sample many times to cover jitter
			d := DelayFor(tc.attempt)
			if d < tc.minDur || d > tc.maxDur {
				t.Errorf("DelayFor(%d)=%v, want [%v, %v]", tc.attempt, d, tc.minDur, tc.maxDur)
				break
			}
		}
	}
}

// TestRetryAfterHonored verifies 429 Retry-After header is respected.
func TestRetryAfterHonored(t *testing.T) {
	security.SetAllowLoopbackForTest(true)
	defer security.SetAllowLoopbackForTest(false)

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Retry-After", "60")
		w.WriteHeader(http.StatusTooManyRequests)
	}))
	defer srv.Close()

	agentID := uuid.New()
	call := newTestCall(srv.URL, &agentID)
	prevResp, _ := json.Marshal(callbackPayload{Output: "output"})
	call.Response = prevResp

	wh, _ := newTestWebhook(call.WebhookID, testEncKey)
	callStore := newStubCallStore(call)
	whStore := &stubWebhookStore{wh: wh}
	w := newTestWorker(callStore, whStore)

	before := time.Now()
	w.execute(context.Background(), call, call.TenantID, "test-lease")

	if callStore.lastUpdate == nil {
		t.Fatal("UpdateStatusCAS never called")
	}
	gotStatus, _ := callStore.lastUpdate["status"].(string)
	if gotStatus != "queued" {
		t.Errorf("429: status=%q, want queued", gotStatus)
	}
	nextAt, _ := callStore.lastUpdate["next_attempt_at"].(time.Time)
	delay := nextAt.Sub(before)
	// Should be ≈60s (± a few ms for test execution).
	if delay < 55*time.Second || delay > 70*time.Second {
		t.Errorf("Retry-After=60 → delay=%v, want ~60s", delay)
	}
}

// TestFourXxPermanentFailed verifies non-429 4xx leads to status=failed (no retry).
func TestFourXxPermanentFailed(t *testing.T) {
	security.SetAllowLoopbackForTest(true)
	defer security.SetAllowLoopbackForTest(false)

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusUnauthorized)
	}))
	defer srv.Close()

	agentID := uuid.New()
	call := newTestCall(srv.URL, &agentID)
	prevResp, _ := json.Marshal(callbackPayload{Output: "output"})
	call.Response = prevResp

	wh, _ := newTestWebhook(call.WebhookID, testEncKey)
	callStore := newStubCallStore(call)
	whStore := &stubWebhookStore{wh: wh}
	w := newTestWorker(callStore, whStore)

	w.execute(context.Background(), call, call.TenantID, "test-lease")

	gotStatus, _ := callStore.lastUpdate["status"].(string)
	if gotStatus != "failed" {
		t.Errorf("401: status=%q, want failed", gotStatus)
	}
}

// TestFiveConsecutive5xxLeadsToDead verifies MaxAttempts=5 consecutive 5xx → dead.
func TestFiveConsecutive5xxLeadsToDead(t *testing.T) {
	security.SetAllowLoopbackForTest(true)
	defer security.SetAllowLoopbackForTest(false)

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
	}))
	defer srv.Close()

	agentID := uuid.New()
	call := newTestCall(srv.URL, &agentID)
	prevResp, _ := json.Marshal(callbackPayload{Output: "output"})
	call.Response = prevResp

	wh, _ := newTestWebhook(call.WebhookID, testEncKey)
	callStore := newStubCallStore(call)
	whStore := &stubWebhookStore{wh: wh}
	w := newTestWorker(callStore, whStore)

	// Simulate MaxAttempts - 1 prior failures (call.Attempts tracks pre-send count).
	call.Attempts = MaxAttempts - 1

	w.execute(context.Background(), call, call.TenantID, "test-lease")

	gotStatus, _ := callStore.lastUpdate["status"].(string)
	if gotStatus != "dead" {
		t.Errorf("5th 500: status=%q, want dead", gotStatus)
	}
	gotAttempts, _ := callStore.lastUpdate["attempts"].(int)
	if gotAttempts != MaxAttempts {
		t.Errorf("5th 500: attempts=%d, want %d", gotAttempts, MaxAttempts)
	}
}

// TestPanicInExecuteRecovered verifies a panic inside execute is recovered and the
// row is retried (not left in running state).
func TestPanicInExecuteRecovered(t *testing.T) {
	agentID := uuid.New()
	call := newTestCall("http://should-not-reach", &agentID)
	// Pre-populate response so agent step is skipped; no callback_url after SSRF check.
	call.Response = []byte(`{"output":"test"}`)

	// Webhook with empty encrypted_secret causes "no HMAC" path — but callback_url is
	// 192.168.1.1 which is blocked by SSRF, so status=failed is set before HMAC step.
	// Use a private-IP URL to hit the SSRF-blocked path deterministically.
	cbURL := "http://192.168.1.1/callback"
	call.CallbackURL = &cbURL

	wh := &store.WebhookData{ID: call.WebhookID}
	callStore := newStubCallStore(call)
	whStore := &stubWebhookStore{wh: wh}
	w := newTestWorker(callStore, whStore)

	// Should not panic; recover() catches it and calls updateRetry.
	defer func() {
		if r := recover(); r != nil {
			t.Fatalf("panic escaped execute: %v", r)
		}
	}()

	w.execute(context.Background(), call, call.TenantID, "test-lease")

	// Row should be in failed state (SSRF blocked).
	if callStore.lastUpdate == nil {
		t.Fatal("UpdateStatusCAS never called after SSRF-blocked URL")
	}
	gotStatus, _ := callStore.lastUpdate["status"].(string)
	if gotStatus != "failed" && gotStatus != "queued" {
		t.Errorf("SSRF-blocked: status=%q, want failed or queued", gotStatus)
	}
}

// TestSlotDrainFixed verifies K4: the semaphore slot is released after every
// goroutine dispatch, including successful ones. With concurrency=1 and a
// non-blocking pollOneTenant mock, a second poll must be able to acquire the slot.
func TestSlotDrainFixed(t *testing.T) {
	// This is a unit-level slot test — we invoke pollOneTenant indirectly
	// by checking that slotCh has room after the goroutine runs.
	slotCh := make(chan struct{}, 1)

	// Simulate acquiring the slot.
	slotCh <- struct{}{}
	slotRelease := func() { <-slotCh }

	// Simulate a goroutine that runs and calls slotRelease.
	done := make(chan struct{})
	go func() {
		slotRelease()
		close(done)
	}()

	<-done

	// After the goroutine exits the slot should be free.
	select {
	case slotCh <- struct{}{}:
		// Success — slot was properly released (K4 fix works).
		<-slotCh
	default:
		t.Error("K4: slot not released after goroutine exit — worker would wedge")
	}
}

// TestLeaseExpiredIgnored verifies K5: when UpdateStatusCAS returns ErrLeaseExpired,
// the worker logs a warning and does not return an error to the caller.
func TestLeaseExpiredIgnored(t *testing.T) {
	security.SetAllowLoopbackForTest(true)
	defer security.SetAllowLoopbackForTest(false)

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	defer srv.Close()

	agentID := uuid.New()
	call := newTestCall(srv.URL, &agentID)
	prevResp, _ := json.Marshal(callbackPayload{Output: "output"})
	call.Response = prevResp

	wh, _ := newTestWebhook(call.WebhookID, testEncKey)
	callStore := newStubCallStore(call)
	callStore.casLeaseErr = store.ErrLeaseExpired // simulate stale lease
	whStore := &stubWebhookStore{wh: wh}
	w := newTestWorker(callStore, whStore)

	// Should not panic or error — lease expiry is a normal concurrent race condition.
	defer func() {
		if r := recover(); r != nil {
			t.Fatalf("K5: panic on ErrLeaseExpired: %v", r)
		}
	}()

	w.execute(context.Background(), call, call.TenantID, "stale-lease")
	// No assertions on lastUpdate — the CAS was rejected so lastUpdate stays nil.
}

// TestCallbackLimiterNonBlocking verifies TryAcquire returns false when at capacity.
func TestCallbackLimiterNonBlocking(t *testing.T) {
	limiter := NewCallbackLimiter(2)
	defer limiter.Stop()

	tid := "tenant-abc"

	// Acquire all slots.
	if !limiter.TryAcquire(tid) {
		t.Fatal("first TryAcquire should succeed")
	}
	if !limiter.TryAcquire(tid) {
		t.Fatal("second TryAcquire should succeed")
	}

	// Third should fail (cap=2).
	if limiter.TryAcquire(tid) {
		t.Error("third TryAcquire should return false when at capacity")
	}

	// Release one and retry.
	limiter.Release(tid)
	if !limiter.TryAcquire(tid) {
		t.Error("TryAcquire should succeed after Release")
	}
}

// TestStaleReclaimThreshold verifies that ReclaimStale is called with correct threshold.
func TestStaleReclaimThreshold(t *testing.T) {
	callStore := newStubCallStore(nil)
	callStore.reclaimN = 3
	w := &WebhookWorker{
		calls:   callStore,
		limiter: NewCallbackLimiter(4),
		cfg:     WorkerConfig{WorkerConcurrency: 1},
	}

	before := time.Now()
	w.reclaimStale(context.Background())
	after := time.Now()

	// The reclaim should complete without error (stub returns reclaimN=3).
	// We can't directly assert the threshold without more instrumentation, but we
	// verify the call completes and we haven't crashed.
	_ = before
	_ = after
	// The stub doesn't record the threshold, so just validate the method runs.
}

// TestSign verifies the sign function produces the expected format.
func TestSign(t *testing.T) {
	key := make([]byte, 32)
	ts := int64(1700000000)
	body := []byte(`{"hello":"world"}`)

	sig := Sign(key, ts, body)

	if !startsWith(sig, "t=1700000000,v1=") {
		t.Errorf("unexpected sign output: %q", sig)
	}
	// v1= part should be 64 hex chars (SHA-256 = 32 bytes).
	parts := splitComma(sig)
	var v1 string
	for _, p := range parts {
		if startsWith(p, "v1=") {
			v1 = p[3:]
		}
	}
	if len(v1) != 64 {
		t.Errorf("v1 hex length: got %d, want 64", len(v1))
	}
}

// TestHeartbeatLoopRenews verifies the loop calls Heartbeat repeatedly and does NOT
// cancel the run while the lease stays valid.
func TestHeartbeatLoopRenews(t *testing.T) {
	callStore := newStubCallStore(nil)
	w := newTestWorker(callStore, &stubWebhookStore{})

	runCtx, cancelRun := context.WithCancel(context.Background())
	defer cancelRun()
	stop := make(chan struct{})

	go w.heartbeatLoop(context.Background(), uuid.New(), "lease-1", 5*time.Millisecond, stop, cancelRun)

	time.Sleep(40 * time.Millisecond)
	close(stop)
	time.Sleep(10 * time.Millisecond)

	if atomic.LoadInt32(&callStore.hbCount) < 3 {
		t.Errorf("expected ≥3 heartbeats, got %d", atomic.LoadInt32(&callStore.hbCount))
	}
	if runCtx.Err() != nil {
		t.Errorf("runCtx must NOT be cancelled while lease valid; err=%v", runCtx.Err())
	}
}

func TestBuildCallbackBreakdown(t *testing.T) {
	calls := []providers.CallUsage{
		{Type: "llm_call", Provider: "p", Model: "m", Usage: providers.Usage{PromptTokens: 10, TotalTokens: 10}, CostUSD: 0.01},
		{Type: "tool_call", Name: "read_image", Provider: "9router", Model: "cx/gpt-5.5",
			Usage: providers.Usage{PromptTokens: 100, TotalTokens: 100}, CostUSD: 0.02},
	}
	usage, cost := callbackBreakdown(calls)
	if usage == nil || usage.PromptTokens != 110 {
		t.Errorf("usage sum wrong: %+v", usage)
	}
	if cost < 0.0299 || cost > 0.0301 {
		t.Errorf("cost = %f, want ~0.03", cost)
	}
	if usage2, _ := callbackBreakdown(nil); usage2 != nil {
		t.Error("nil calls → nil usage")
	}
}

// TestHeartbeatLoopCancelsOnLeaseLost verifies that when Heartbeat returns ErrLeaseExpired,
// the loop cancels runCtx and returns.
func TestHeartbeatLoopCancelsOnLeaseLost(t *testing.T) {
	callStore := newStubCallStore(nil)
	callStore.hbErr = store.ErrLeaseExpired // mất lease ngay tick đầu
	w := newTestWorker(callStore, &stubWebhookStore{})

	runCtx, cancelRun := context.WithCancel(context.Background())
	defer cancelRun()
	stop := make(chan struct{})
	defer close(stop)

	go w.heartbeatLoop(context.Background(), uuid.New(), "lease-1", 5*time.Millisecond, stop, cancelRun)

	select {
	case <-runCtx.Done():
		// đúng — run bị hủy do mất lease
	case <-time.After(200 * time.Millisecond):
		t.Fatal("runCtx not cancelled after lease lost within 200ms")
	}
}

// ---- test helpers ----

func startsWith(s, pfx string) bool {
	return len(s) >= len(pfx) && s[:len(pfx)] == pfx
}

func splitComma(s string) []string {
	var parts []string
	for _, p := range splitBytes([]byte(s), ',') {
		parts = append(parts, string(p))
	}
	return parts
}

func splitBytes(b []byte, sep byte) [][]byte {
	var out [][]byte
	start := 0
	for i, c := range b {
		if c == sep {
			out = append(out, b[start:i])
			start = i + 1
		}
	}
	out = append(out, b[start:])
	return out
}

func parseInt64(s string) int64 {
	var n int64
	for _, c := range s {
		if c < '0' || c > '9' {
			break
		}
		n = n*10 + int64(c-'0')
	}
	return n
}