package middleware

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
)

func TestWriteAdmissionConfigFromEnv(t *testing.T) {
	env := map[string]string{
		writeAdmissionMaxInFlightEnv:  "25",
		writeAdmissionWaitMillisEnv:   "150",
		writeAdmissionRetrySecondsEnv: "9",
	}

	cfg := WriteAdmissionConfigFromEnv(func(key string) string { return env[key] })

	if cfg.MaxInFlight != 25 {
		t.Fatalf("MaxInFlight = %d, want 25", cfg.MaxInFlight)
	}
	if cfg.WaitTimeout != 150*time.Millisecond {
		t.Fatalf("WaitTimeout = %s, want 150ms", cfg.WaitTimeout)
	}
	if cfg.RetryAfterSeconds != 9 {
		t.Fatalf("RetryAfterSeconds = %d, want 9", cfg.RetryAfterSeconds)
	}
}

func TestWriteAdmissionConfigFallsBackForInvalidValues(t *testing.T) {
	env := map[string]string{
		writeAdmissionMaxInFlightEnv:  "not-a-number",
		writeAdmissionWaitMillisEnv:   "-1",
		writeAdmissionRetrySecondsEnv: "0",
	}

	cfg := WriteAdmissionConfigFromEnv(func(key string) string { return env[key] })

	if cfg.MaxInFlight != 0 {
		t.Fatalf("MaxInFlight = %d, want disabled fallback 0", cfg.MaxInFlight)
	}
	if cfg.WaitTimeout != 0 {
		t.Fatalf("WaitTimeout = %s, want 0", cfg.WaitTimeout)
	}
	if cfg.RetryAfterSeconds != defaultWriteAdmissionRetrySeconds {
		t.Fatalf("RetryAfterSeconds = %d, want default %d", cfg.RetryAfterSeconds, defaultWriteAdmissionRetrySeconds)
	}
}

func TestWriteAdmissionDisabledReturnsNil(t *testing.T) {
	if got := NewWriteAdmission(WriteAdmissionConfig{}); got != nil {
		t.Fatalf("NewWriteAdmission disabled = %#v, want nil", got)
	}
}

func TestWriteAdmissionSkipsReadRequests(t *testing.T) {
	gin.SetMode(gin.TestMode)
	router := gin.New()
	router.Use(NewWriteAdmission(WriteAdmissionConfig{MaxInFlight: 1}))
	router.GET("/objects", func(c *gin.Context) { c.Status(http.StatusNoContent) })

	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/objects", nil))

	if rec.Code != http.StatusNoContent {
		t.Fatalf("GET status = %d body=%s, want %d", rec.Code, rec.Body.String(), http.StatusNoContent)
	}
}

func TestWriteAdmissionSkipsPostQueryRequests(t *testing.T) {
	gin.SetMode(gin.TestMode)
	router := gin.New()
	router.Use(NewWriteAdmission(WriteAdmissionConfig{MaxInFlight: 1}))

	started := make(chan struct{})
	release := make(chan struct{})
	router.POST("/services/data/:version/sobjects/:sobject", func(c *gin.Context) {
		close(started)
		<-release
		c.Status(http.StatusCreated)
	})
	router.POST("/services/data/:version/query", func(c *gin.Context) {
		c.Status(http.StatusNoContent)
	})

	firstDone := make(chan int, 1)
	go func() {
		rec := httptest.NewRecorder()
		router.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/services/data/v63.0/sobjects/Account", nil))
		firstDone <- rec.Code
	}()

	select {
	case <-started:
	case <-time.After(2 * time.Second):
		t.Fatal("first request did not acquire the write admission slot")
	}

	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/services/data/v63.0/query", nil))
	if rec.Code != http.StatusNoContent {
		t.Fatalf("POST query status = %d body=%s, want %d", rec.Code, rec.Body.String(), http.StatusNoContent)
	}

	close(release)
	select {
	case code := <-firstDone:
		if code != http.StatusCreated {
			t.Fatalf("first POST status = %d, want %d", code, http.StatusCreated)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("first request did not finish after release")
	}
}

func TestWriteAdmissionRejectsWhenFull(t *testing.T) {
	gin.SetMode(gin.TestMode)
	router := gin.New()
	router.Use(NewWriteAdmission(WriteAdmissionConfig{
		MaxInFlight:       1,
		RetryAfterSeconds: 7,
	}))

	started := make(chan struct{})
	release := make(chan struct{})
	router.POST("/objects", func(c *gin.Context) {
		close(started)
		<-release
		c.Status(http.StatusNoContent)
	})

	firstDone := make(chan int, 1)
	go func() {
		rec := httptest.NewRecorder()
		router.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/objects", nil))
		firstDone <- rec.Code
	}()

	select {
	case <-started:
	case <-time.After(2 * time.Second):
		t.Fatal("first request did not acquire the write admission slot")
	}

	rec := httptest.NewRecorder()
	router.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/objects", nil))

	if rec.Code != http.StatusServiceUnavailable {
		t.Fatalf("second POST status = %d body=%s, want %d", rec.Code, rec.Body.String(), http.StatusServiceUnavailable)
	}
	if got := rec.Header().Get("Retry-After"); got != "7" {
		t.Fatalf("Retry-After = %q, want 7", got)
	}
	if body := rec.Body.String(); !strings.Contains(body, "write_admission_full") {
		t.Fatalf("body = %s, want write_admission_full", body)
	}

	close(release)
	select {
	case code := <-firstDone:
		if code != http.StatusNoContent {
			t.Fatalf("first POST status = %d, want %d", code, http.StatusNoContent)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("first request did not finish after release")
	}
}