package middleware

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

	"git.yuntuo.hk/k2safe/go-omni/httpx"
)

func invoke(h httpx.Handler, method, body, scope, key string) *httptest.ResponseRecorder {
	r := httptest.NewRequest(method, "/records?version=1", strings.NewReader(body))
	r.Header.Set("Content-Type", "application/json")
	if key != "" {
		r.Header.Set("Idempotency-Key", key)
	}
	w := httptest.NewRecorder()
	c := httpx.NewContext(r, w, nil)
	c.Set("verifiedIdentity", scope)
	if err := h(c); err != nil {
		c.WriteError(err)
	}
	return w
}

func identity(c *httpx.Context) (string, error) {
	v, _ := c.Get("verifiedIdentity")
	s, _ := v.(string)
	return s, nil
}

func TestRateLimitRefillCapacityAndIdleExpiry(t *testing.T) {
	l, err := NewRateLimiter(RateLimitConfig{Rate: 2, Burst: 2, MaxKeys: 1, IdleTTL: time.Second, Key: func(c *httpx.Context) string { s, _ := identity(c); return s }})
	if err != nil {
		t.Fatal(err)
	}
	now := time.Unix(100, 0)
	l.now = func() time.Time { return now }
	h := l.Middleware()(func(c *httpx.Context) error { return c.NoContent() })
	for _, status := range []int{204, 204, 429} {
		w := invoke(h, "GET", "", "a", "")
		if w.Code != status {
			t.Fatal(w.Code, w.Body.String())
		}
		if status == 429 && w.Header().Get("Retry-After") != "1" {
			t.Fatal(w.Header())
		}
	}
	if w := invoke(h, "GET", "", "b", ""); w.Code != 503 {
		t.Fatal(w.Code)
	}
	now = now.Add(500 * time.Millisecond)
	if w := invoke(h, "GET", "", "a", ""); w.Code != 204 {
		t.Fatal(w.Code)
	}
	now = now.Add(time.Second)
	if w := invoke(h, "GET", "", "b", ""); w.Code != 204 {
		t.Fatal(w.Code)
	}
	if _, err := NewRateLimiter(RateLimitConfig{Rate: 1, Burst: 2, IdleTTL: time.Second}); err == nil {
		t.Fatal("early bucket expiry accepted")
	}
}

func TestRateLimitConcurrentBurst(t *testing.T) {
	l, _ := NewRateLimiter(RateLimitConfig{Rate: 1, Burst: 7})
	l.now = func() time.Time { return time.Unix(100, 0) }
	var allowed atomic.Int32
	var wg sync.WaitGroup
	for n := 0; n < 100; n++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			delay, full := l.allow("same")
			if delay == 0 && !full {
				allowed.Add(1)
			}
		}()
	}
	wg.Wait()
	if allowed.Load() != 7 {
		t.Fatal(allowed.Load())
	}
}

func TestTimeoutCooperativeAndCommittedResponse(t *testing.T) {
	for _, committed := range []bool{false, true} {
		h := Timeout(5 * time.Millisecond)(func(c *httpx.Context) error {
			if committed {
				_ = c.Accepted("accepted")
			}
			<-c.Context().Done()
			return c.Context().Err()
		})
		w := invoke(h, "GET", "", "", "")
		want := 504
		if committed {
			want = 202
		}
		if w.Code != want {
			t.Fatal(w.Code, w.Body.String())
		}
	}
}

func TestIdempotencyReplayScopeConflictAndExpiry(t *testing.T) {
	i, err := NewIdempotency(IdempotencyConfig{Scope: identity, TTL: time.Minute})
	if err != nil {
		t.Fatal(err)
	}
	now := time.Unix(100, 0)
	i.now = func() time.Time { return now }
	calls := 0
	h := i.Middleware()(func(c *httpx.Context) error { calls++; return c.Created(map[string]int{"calls": calls}) })
	first := invoke(h, "POST", `{"name":"one"}`, "alice", "key1")
	replay := invoke(h, "POST", `{"name":"one"}`, "alice", "key1")
	if first.Code != 201 || replay.Code != 201 || first.Body.String() != replay.Body.String() || replay.Header().Get("Idempotency-Replayed") != "true" || calls != 1 {
		t.Fatal(first, replay, calls)
	}
	if w := invoke(h, "POST", `{"name":"two"}`, "alice", "key1"); w.Code != 409 {
		t.Fatal(w.Code)
	}
	if w := invoke(h, "POST", `{"name":"one"}`, "bob", "key1"); w.Code != 201 || calls != 2 {
		t.Fatal(w.Code, calls)
	}
	if w := invoke(h, "POST", `{}`, "", "key1"); w.Code != 401 {
		t.Fatal(w.Code)
	}
	if w := invoke(h, "POST", `{}`, "alice", ""); w.Code != 400 {
		t.Fatal(w.Code)
	}
	now = now.Add(time.Minute)
	if w := invoke(h, "POST", `{"name":"one"}`, "alice", "key1"); w.Code != 201 || calls != 3 {
		t.Fatal(w.Code, calls)
	}
}

func TestIdempotencyActiveClaimNeverExpires(t *testing.T) {
	i, _ := NewIdempotency(IdempotencyConfig{Scope: identity, TTL: time.Nanosecond, MaxInFlight: 2})
	started, release := make(chan struct{}), make(chan struct{})
	var calls atomic.Int32
	h := i.Middleware()(func(c *httpx.Context) error { calls.Add(1); close(started); <-release; return c.Created("saved") })
	done := make(chan *httptest.ResponseRecorder, 1)
	go func() { done <- invoke(h, "POST", "{}", "alice", "key") }()
	<-started
	for n := 0; n < 30; n++ {
		if w := invoke(h, "POST", "{}", "alice", "key"); w.Code != 409 {
			t.Errorf("duplicate status %d", w.Code)
		}
	}
	close(release)
	if w := <-done; w.Code != 201 || calls.Load() != 1 {
		t.Fatal(w.Code, calls.Load())
	}
}

func TestIdempotencyBoundsAndAmbiguousOutcomes(t *testing.T) {
	for _, mode := range []string{"panic", "error", "overflow", "cookie", "noStore", "memory"} {
		t.Run(mode, func(t *testing.T) {
			cfg := IdempotencyConfig{Scope: identity, MaxRequestBytes: 8, MaxResponseBytes: 128, MaxEntries: 1}
			if mode == "memory" {
				cfg.MaxBytes = 1
			}
			i, _ := NewIdempotency(cfg)
			calls := 0
			h := Recovery()(i.Middleware()(func(c *httpx.Context) error {
				calls++
				switch mode {
				case "panic":
					panic("failure")
				case "error":
					return errors.New("failed")
				case "overflow":
					return c.OK(strings.Repeat("x", 200))
				case "cookie":
					c.SetResponseHeader("Set-Cookie", "session=secret")
				case "noStore":
					c.SetResponseHeader("Cache-Control", "no-store")
				}
				return c.Created("ok")
			}))
			if w := invoke(h, "POST", "123456789", "alice", "key"); w.Code != 413 || calls != 0 {
				t.Fatal(w.Code, calls)
			}
			first := invoke(h, "POST", "{}", "alice", "key")
			if mode == "panic" || mode == "error" || mode == "overflow" {
				if first.Code != 500 {
					t.Fatal(first.Code)
				}
			} else if first.Code != 201 {
				t.Fatal(first.Code)
			}
			if w := invoke(h, "POST", "{}", "alice", "key"); w.Code != 409 || calls != 1 {
				t.Fatal(w.Code, calls)
			}
			if w := invoke(h, "POST", "{}", "alice", "newKey"); w.Code != 503 {
				t.Fatal(w.Code)
			}
		})
	}
}

func TestIdempotencyInFlightLimitAndTimeoutComposition(t *testing.T) {
	i, _ := NewIdempotency(IdempotencyConfig{Scope: identity, MaxInFlight: 1})
	started, release := make(chan struct{}), make(chan struct{})
	h := i.Middleware()(func(c *httpx.Context) error { close(started); <-release; return c.NoContent() })
	done := make(chan struct{})
	go func() { defer close(done); invoke(h, "POST", "{}", "alice", "one") }()
	<-started
	if w := invoke(h, "POST", "{}", "alice", "two"); w.Code != 503 {
		t.Error(w.Code)
	}
	close(release)
	<-done
	timed := Timeout(5 * time.Millisecond)(i.Middleware()(func(c *httpx.Context) error { <-c.Context().Done(); return context.DeadlineExceeded }))
	if w := invoke(timed, "POST", "{}", "alice", "timeout"); w.Code != 504 {
		t.Fatal(w.Code, w.Body.String())
	}
	if w := invoke(timed, "POST", "{}", "alice", "timeout"); w.Code != 409 {
		t.Fatal(w.Code)
	}
}