package audit

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

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

// fakeCaptureStore is an in-memory RequestCaptureStore that records every insert.
// Goroutine-safe — the middleware dispatches InsertRequestCapture from a
// post-handler goroutine.
type fakeCaptureStore struct {
	mu   sync.Mutex
	rows []RequestCaptureRow
	err  error // if non-nil, InsertRequestCapture returns this
}

func (s *fakeCaptureStore) InsertRequestCapture(ctx context.Context, row RequestCaptureRow) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.err != nil {
		return s.err
	}
	s.rows = append(s.rows, row)
	return nil
}

func (s *fakeCaptureStore) snapshot() []RequestCaptureRow {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]RequestCaptureRow, len(s.rows))
	copy(out, s.rows)
	return out
}

// newTestRouter builds a gin router with the capture middleware wired in
// the same order the production main.go uses (after AuditEntrypoint so ctx
// has the entrypoint). Returns the router + underlying fake store.
func newTestRouter(t *testing.T, cfg RequestCaptureConfig, handler gin.HandlerFunc) *gin.Engine {
	t.Helper()
	gin.SetMode(gin.TestMode)
	r := gin.New()
	r.Use(func(c *gin.Context) {
		c.Set(ctxRequestIDKey, "req-fixed-id")
		c.Next()
	})
	r.Use(AuditEntrypointMiddleware())
	r.Use(RequestCaptureMiddleware(cfg))
	r.POST("/api/test", handler)
	return r
}

// awaitRow waits briefly for the async capture goroutine to complete.
// Returns the row if one arrives, or fails the test after 1s.
func awaitRow(t *testing.T, store *fakeCaptureStore) RequestCaptureRow {
	t.Helper()
	deadline := time.Now().Add(1 * time.Second)
	for time.Now().Before(deadline) {
		rows := store.snapshot()
		if len(rows) == 1 {
			return rows[0]
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatalf("no capture row after 1s; got %d rows", len(store.snapshot()))
	return RequestCaptureRow{}
}

// awaitNoRow waits a grace period and asserts no row was written.
func awaitNoRow(t *testing.T, store *fakeCaptureStore) {
	t.Helper()
	time.Sleep(100 * time.Millisecond)
	if rows := store.snapshot(); len(rows) != 0 {
		t.Fatalf("expected no capture rows, got %d", len(rows))
	}
}

func enabled() bool  { return true }
func disabled() bool { return false }

// --- Test 1: capture disabled → no DB write ---

func TestRequestCaptureMiddleware_Disabled_NoWrite(t *testing.T) {
	store := &fakeCaptureStore{}
	r := newTestRouter(t, RequestCaptureConfig{
		Store:     store,
		EnabledFn: disabled,
	}, func(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"ok": true}) })

	req := httptest.NewRequest(http.MethodPost, "/api/test",
		strings.NewReader(`{"password":"hunter2"}`))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status = %d, want 200", w.Code)
	}
	awaitNoRow(t, store)
}

// --- Test 2: capture enabled → JSON body scrubbed ---

func TestRequestCaptureMiddleware_Enabled_BodyScrubbed(t *testing.T) {
	store := &fakeCaptureStore{}
	r := newTestRouter(t, RequestCaptureConfig{
		Store:     store,
		EnabledFn: enabled,
	}, func(c *gin.Context) {
		// Ensure handler can still read body (not consumed by middleware).
		raw, _ := io.ReadAll(c.Request.Body)
		if !strings.Contains(string(raw), "hunter2") {
			t.Errorf("handler body lost: %q", raw)
		}
		c.JSON(http.StatusOK, gin.H{"ok": true})
	})

	body := `{"password":"hunter2","name":"alice","nested":{"accessToken":"abc"}}`
	req := httptest.NewRequest(http.MethodPost, "/api/test", strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	row := awaitRow(t, store)
	if strings.Contains(row.Body, "hunter2") {
		t.Errorf("body still contains password: %q", row.Body)
	}
	if strings.Contains(row.Body, "abc") {
		t.Errorf("body still contains nested accessToken: %q", row.Body)
	}
	if !strings.Contains(row.Body, PIIScrubSentinel) {
		t.Errorf("body missing sentinel: %q", row.Body)
	}
	if !strings.Contains(row.Body, "alice") {
		t.Errorf("non-PII field dropped: %q", row.Body)
	}
	if row.BodyTruncated {
		t.Errorf("BodyTruncated should be false for small payload")
	}
	if row.BodySize != len(body) {
		t.Errorf("BodySize = %d, want %d", row.BodySize, len(body))
	}
	if row.Method != http.MethodPost {
		t.Errorf("Method = %q, want POST", row.Method)
	}
	if row.Path != "/api/test" {
		t.Errorf("Path = %q", row.Path)
	}
	if row.RequestID != "req-fixed-id" {
		t.Errorf("RequestID = %q, want req-fixed-id (from ctx)", row.RequestID)
	}
	if row.Entrypoint != string(EntrypointRestAPI) {
		t.Errorf("Entrypoint = %q, want REST_API", row.Entrypoint)
	}
}

// --- Test 3: headers scrubbed ---

func TestRequestCaptureMiddleware_Enabled_HeadersScrubbed(t *testing.T) {
	store := &fakeCaptureStore{}
	r := newTestRouter(t, RequestCaptureConfig{
		Store:     store,
		EnabledFn: enabled,
	}, func(c *gin.Context) { c.Status(http.StatusOK) })

	req := httptest.NewRequest(http.MethodPost, "/api/test", strings.NewReader(`{}`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer supersecret")
	req.Header.Set("Cookie", "sid=abc")
	req.Header.Set("X-Custom-Benign", "keep-me")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	row := awaitRow(t, store)
	if strings.Contains(row.Headers, "supersecret") {
		t.Errorf("headers still contain token: %q", row.Headers)
	}
	if strings.Contains(row.Headers, "sid=abc") {
		t.Errorf("headers still contain cookie: %q", row.Headers)
	}
	if !strings.Contains(row.Headers, PIIScrubSentinel) {
		t.Errorf("headers missing sentinel: %q", row.Headers)
	}
	if !strings.Contains(row.Headers, "keep-me") {
		t.Errorf("benign header dropped: %q", row.Headers)
	}

	// Verify it's valid JSON so downstream export can parse it.
	var out map[string][]string
	if err := json.Unmarshal([]byte(row.Headers), &out); err != nil {
		t.Fatalf("headers not JSON: %v", err)
	}
}

// --- Test 4: response capture default-off → ResponseBody NULL ---

func TestRequestCaptureMiddleware_ResponseCapture_DefaultOff(t *testing.T) {
	store := &fakeCaptureStore{}
	r := newTestRouter(t, RequestCaptureConfig{
		Store:     store,
		EnabledFn: enabled,
		// ResponseCaptureFn omitted → production default (env kill-switch OFF).
	}, func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{"secret": "dont-leak"})
	})

	req := httptest.NewRequest(http.MethodPost, "/api/test", strings.NewReader(`{}`))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	row := awaitRow(t, store)
	if row.ResponseBody.Valid {
		t.Errorf("ResponseBody should be NULL by default, got %q", row.ResponseBody.String)
	}
	if row.ResponseStatus != http.StatusOK {
		t.Errorf("ResponseStatus = %d, want 200", row.ResponseStatus)
	}
}

// --- Test 5: response capture opted-in → populated + scrubbed ---

func TestRequestCaptureMiddleware_ResponseCapture_OptIn(t *testing.T) {
	store := &fakeCaptureStore{}
	r := newTestRouter(t, RequestCaptureConfig{
		Store:             store,
		EnabledFn:         enabled,
		ResponseCaptureFn: func(path string) bool { return path == "/api/test" },
	}, func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{"token": "will-be-scrubbed", "ok": true})
	})

	req := httptest.NewRequest(http.MethodPost, "/api/test", strings.NewReader(`{}`))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	row := awaitRow(t, store)
	if !row.ResponseBody.Valid {
		t.Fatalf("ResponseBody not captured despite opt-in")
	}
	if strings.Contains(row.ResponseBody.String, "will-be-scrubbed") {
		t.Errorf("response body not scrubbed: %q", row.ResponseBody.String)
	}
	if !strings.Contains(row.ResponseBody.String, PIIScrubSentinel) {
		t.Errorf("response body missing sentinel: %q", row.ResponseBody.String)
	}
}

// --- Test 6: body exceeds cap → truncated + BodyTruncated=true ---

func TestRequestCaptureMiddleware_BodyCap_Truncation(t *testing.T) {
	store := &fakeCaptureStore{}
	cap := 64
	r := newTestRouter(t, RequestCaptureConfig{
		Store:        store,
		EnabledFn:    enabled,
		BodyCapBytes: cap,
	}, func(c *gin.Context) { c.Status(http.StatusOK) })

	// Non-JSON so no scrub attempted — easier to assert on raw bytes.
	payload := bytes.Repeat([]byte("A"), cap*3)
	req := httptest.NewRequest(http.MethodPost, "/api/test", bytes.NewReader(payload))
	req.Header.Set("Content-Type", "text/plain")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	row := awaitRow(t, store)
	if !row.BodyTruncated {
		t.Errorf("BodyTruncated should be true, got false")
	}
	if len(row.Body) != cap {
		t.Errorf("Body length = %d, want cap %d", len(row.Body), cap)
	}
	if row.BodySize <= cap {
		t.Errorf("BodySize = %d, want > cap %d", row.BodySize, cap)
	}
}

// --- Test 7: DB write failure → request still succeeds ---

func TestRequestCaptureMiddleware_DBFailure_RequestSucceeds(t *testing.T) {
	store := &fakeCaptureStore{err: errAlwaysFail}
	r := newTestRouter(t, RequestCaptureConfig{
		Store:     store,
		EnabledFn: enabled,
	}, func(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"ok": true}) })

	req := httptest.NewRequest(http.MethodPost, "/api/test", strings.NewReader(`{}`))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("request should succeed despite DB failure; got %d", w.Code)
	}
	// Give the goroutine time to fail silently.
	time.Sleep(50 * time.Millisecond)
	// No rows were successfully stored (store.err short-circuits the append).
	if rows := store.snapshot(); len(rows) != 0 {
		t.Errorf("expected 0 rows on DB failure, got %d", len(rows))
	}
}

// errAlwaysFail is a named error so test output is legible.
type sentinelErr string

func (e sentinelErr) Error() string { return string(e) }

var errAlwaysFail sentinelErr = "simulated DB outage"

// --- Test 8: form-urlencoded body → scrubbed via ScrubFormBody (HIGH-2) ---

func TestRequestCaptureMiddleware_FormBody_Scrubbed(t *testing.T) {
	store := &fakeCaptureStore{}
	r := newTestRouter(t, RequestCaptureConfig{
		Store:     store,
		EnabledFn: enabled,
	}, func(c *gin.Context) { c.Status(http.StatusOK) })

	body := "password=hunter2&name=alice"
	req := httptest.NewRequest(http.MethodPost, "/api/test", strings.NewReader(body))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	row := awaitRow(t, store)
	if strings.Contains(row.Body, "hunter2") {
		t.Errorf("form body still contains password: %q", row.Body)
	}
	if !strings.Contains(row.Body, "name=alice") {
		t.Errorf("non-PII field dropped: %q", row.Body)
	}
	if !strings.Contains(row.Body, "SCRUBBED") {
		t.Errorf("form body missing sentinel: %q", row.Body)
	}
}

// --- Test 8b: XML / text/plain still pass-through (Phase 2B TODO) ---

func TestRequestCaptureMiddleware_XMLBody_PassThrough(t *testing.T) {
	store := &fakeCaptureStore{}
	r := newTestRouter(t, RequestCaptureConfig{
		Store:     store,
		EnabledFn: enabled,
	}, func(c *gin.Context) { c.Status(http.StatusOK) })

	body := "<root><password>hunter2</password></root>"
	req := httptest.NewRequest(http.MethodPost, "/api/test", strings.NewReader(body))
	req.Header.Set("Content-Type", "application/xml")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	row := awaitRow(t, store)
	if row.Body != body {
		t.Errorf("XML body altered: got %q want %q", row.Body, body)
	}
}

// --- Test 8c: query string scrubbed (HIGH-1) ---

func TestRequestCaptureMiddleware_QueryString_Scrubbed(t *testing.T) {
	store := &fakeCaptureStore{}
	r := newTestRouter(t, RequestCaptureConfig{
		Store:     store,
		EnabledFn: enabled,
	}, func(c *gin.Context) { c.Status(http.StatusOK) })

	req := httptest.NewRequest(http.MethodPost,
		"/api/test?token=abc123&name=alice&password=hunter2",
		strings.NewReader(`{}`))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	row := awaitRow(t, store)
	if strings.Contains(row.Query, "abc123") {
		t.Errorf("query still contains token: %q", row.Query)
	}
	if strings.Contains(row.Query, "hunter2") {
		t.Errorf("query still contains password: %q", row.Query)
	}
	if !strings.Contains(row.Query, "name=alice") {
		t.Errorf("benign query param dropped: %q", row.Query)
	}
	if !strings.Contains(row.Query, "SCRUBBED") {
		t.Errorf("query missing sentinel: %q", row.Query)
	}
}

// --- Test 8d: STRICT mode drops non-scrubbable content type ---

func TestRequestCaptureMiddleware_Strict_RejectsUnknownCT(t *testing.T) {
	store := &fakeCaptureStore{}
	strict := func() bool { return true }
	r := newTestRouter(t, RequestCaptureConfig{
		Store:     store,
		EnabledFn: enabled,
		StrictFn:  strict,
	}, func(c *gin.Context) { c.Status(http.StatusOK) })

	// XML body + strict mode → no row written, but handler still runs.
	body := "<root><password>hunter2</password></root>"
	req := httptest.NewRequest(http.MethodPost, "/api/test", strings.NewReader(body))
	req.Header.Set("Content-Type", "application/xml")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("handler should run; got %d", w.Code)
	}
	awaitNoRow(t, store)
}

// --- Test 8e: STRICT mode still allows scrubbable CTs ---

func TestRequestCaptureMiddleware_Strict_AllowsJSON(t *testing.T) {
	store := &fakeCaptureStore{}
	strict := func() bool { return true }
	r := newTestRouter(t, RequestCaptureConfig{
		Store:     store,
		EnabledFn: enabled,
		StrictFn:  strict,
	}, func(c *gin.Context) { c.Status(http.StatusOK) })

	req := httptest.NewRequest(http.MethodPost, "/api/test",
		strings.NewReader(`{"password":"x"}`))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	row := awaitRow(t, store)
	if strings.Contains(row.Body, `"x"`) {
		t.Errorf("strict mode should still scrub JSON, got: %q", row.Body)
	}
}

// --- Test 9: EnableResponseCapture registry ---

func TestEnableResponseCapture_RegistryRoundTrip(t *testing.T) {
	ResetResponseCaptureForTest()
	defer ResetResponseCaptureForTest()

	if isResponseCaptureEnabledForPath("/x") {
		t.Fatalf("expected empty registry to not match")
	}
	EnableResponseCapture("/x")
	if !isResponseCaptureEnabledForPath("/x") {
		t.Fatalf("expected /x to match after enable")
	}
	if isResponseCaptureEnabledForPath("/y") {
		t.Fatalf("/y should not match exact-path registry")
	}
}

// --- Security regression: SessionId must NEVER be the raw bearer token ---
//
// Historical bug (V2-HIGH-2): the auth middleware stashed the raw bearer
// token under gin ctx key "token"; the capture middleware read it and
// wrote it straight into `_System_RequestCapture.SessionId`. A DB dump of
// the capture table therefore yielded replayable credentials. The fix
// stores the validated JWT's session identifier (claims.RegisteredClaims.ID /
// jti) instead. This test pins the invariant at the capture layer.
func TestRequestCaptureMiddleware_SessionID_UsesJWTJtiNotRawToken(t *testing.T) {
	store := &fakeCaptureStore{}
	gin.SetMode(gin.TestMode)
	r := gin.New()
	r.Use(func(c *gin.Context) {
		c.Set(ctxRequestIDKey, "req-sess-id-check")
		// Simulate what auth.RequireAuth now does: set the JWT jti under
		// ctxSessionIDKey, and NEVER set "token" (which used to hold the
		// raw bearer).
		c.Set(ctxSessionIDKey, "jwt-jti-xyz")
		c.Set("token", "RAW-BEARER-SHOULD-NEVER-BE-PERSISTED")
		c.Next()
	})
	r.Use(AuditEntrypointMiddleware())
	r.Use(RequestCaptureMiddleware(RequestCaptureConfig{
		Store:     store,
		EnabledFn: enabled,
	}))
	r.POST("/api/test", func(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"ok": true}) })

	req := httptest.NewRequest(http.MethodPost, "/api/test", strings.NewReader(`{}`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer RAW-BEARER-SHOULD-NEVER-BE-PERSISTED")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	row := awaitRow(t, store)
	if !row.SessionID.Valid || row.SessionID.String != "jwt-jti-xyz" {
		t.Fatalf("SessionId = %+v, want valid=true string=%q", row.SessionID, "jwt-jti-xyz")
	}
	if strings.Contains(row.SessionID.String, "RAW-BEARER") {
		t.Fatalf("SessionId leaked raw bearer token: %q", row.SessionID.String)
	}
}

// --- Test: replay traffic is skipped to prevent re-capture cascade ---
//
// Regression guard (PR #1850 Pass B HIGH): when cmd/replay POSTs against the
// live backend with X-Replay-RequestId, the capture middleware previously
// stamped a fresh _System_RequestCapture row with the suffixed rid. The next
// replay invocation then re-replayed those rows (with a second-level suffix),
// causing exponential growth of the capture table on repeat runs. Replay
// traffic is self-identified by the header and does not need to be captured.
func TestRequestCaptureMiddleware_SkipsReplayTraffic(t *testing.T) {
	store := &fakeCaptureStore{}
	r := newTestRouter(t, RequestCaptureConfig{
		Store:     store,
		EnabledFn: enabled,
	}, func(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"ok": true}) })

	req := httptest.NewRequest(http.MethodPost, "/api/test",
		strings.NewReader(`{"k":"v"}`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Replay-RequestId", "rid-test-replay")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status = %d, want 200 (replay traffic must still be handled)", w.Code)
	}
	awaitNoRow(t, store)
}

// --- Test 10: truthy env parsing ---

func TestIsTruthy(t *testing.T) {
	for _, v := range []string{"on", "ON", "On", "true", "1", "yes", " on "} {
		if !isTruthy(v) {
			t.Errorf("isTruthy(%q) = false, want true", v)
		}
	}
	for _, v := range []string{"", "off", "false", "0", "no", "maybe"} {
		if isTruthy(v) {
			t.Errorf("isTruthy(%q) = true, want false", v)
		}
	}
}