package middleware

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

	"github.com/gin-contrib/sessions"
	"github.com/gin-contrib/sessions/cookie"
	"github.com/gin-gonic/gin"
)

func newCSRFTestRouter() *gin.Engine {
	gin.SetMode(gin.TestMode)
	r := gin.New()
	store := cookie.NewStore([]byte("test-secret"))
	r.Use(sessions.Sessions("test_session", store))
	r.Use(CSRFProtect())

	r.GET("/form", func(c *gin.Context) {
		c.String(http.StatusOK, "TOKEN="+c.GetString(CSRFContextKey))
	})
	r.HEAD("/form", func(c *gin.Context) {
		c.String(http.StatusOK, "TOKEN="+c.GetString(CSRFContextKey))
	})
	r.POST("/action", func(c *gin.Context) {
		c.String(http.StatusOK, "ok")
	})
	return r
}

// tokenFromForm performs GET /form with the given session cookie and returns
// the issued CSRF token plus the (possibly new) session cookie.
func tokenFromForm(t *testing.T, r *gin.Engine, sessionCookie string) (token, cookie string) {
	t.Helper()
	req := httptest.NewRequest(http.MethodGet, "/form", nil)
	if sessionCookie != "" {
		req.Header.Set("Cookie", sessionCookie)
	}
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("GET /form: status = %d, want 200", w.Code)
	}
	body := w.Body.String()
	const prefix = "TOKEN="
	if !strings.HasPrefix(body, prefix) {
		t.Fatalf("GET /form: unexpected body %q", body)
	}
	token = strings.TrimPrefix(body, prefix)
	cookie = w.Header().Get("Set-Cookie")
	return token, cookie
}

func postAction(r *gin.Engine, sessionCookie, token string, useHeader bool) *httptest.ResponseRecorder {
	form := url.Values{}
	if !useHeader {
		form.Set(CSRFFieldName, token)
	}
	req := httptest.NewRequest(http.MethodPost, "/action", strings.NewReader(form.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	if sessionCookie != "" {
		req.Header.Set("Cookie", sessionCookie)
	}
	if useHeader {
		req.Header.Set(CSRFHeaderName, token)
	}
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}

func TestCSRFTokenIssuedOnGET(t *testing.T) {
	r := newCSRFTestRouter()
	token, cookie := tokenFromForm(t, r, "")
	if token == "" {
		t.Fatal("expected a token to be issued on GET")
	}
	if !strings.Contains(cookie, "test_session=") {
		t.Fatalf("expected session cookie to be set, got %q", cookie)
	}

	// A second GET with the same session must return the same token.
	token2, _ := tokenFromForm(t, r, cookie)
	if token2 != token {
		t.Fatalf("token changed between requests: %q vs %q", token, token2)
	}
}

func TestCSRFPostRejectedWithoutToken(t *testing.T) {
	r := newCSRFTestRouter()
	_, cookie := tokenFromForm(t, r, "")

	w := postAction(r, cookie, "", false)
	if w.Code != http.StatusForbidden {
		t.Fatalf("POST without token: status = %d, want 403", w.Code)
	}
}

func TestCSRFPostRejectedWithWrongToken(t *testing.T) {
	r := newCSRFTestRouter()
	_, cookie := tokenFromForm(t, r, "")

	w := postAction(r, cookie, "bogus-token", false)
	if w.Code != http.StatusForbidden {
		t.Fatalf("POST with wrong token: status = %d, want 403", w.Code)
	}
}

func TestCSRFPostRejectedWithoutSession(t *testing.T) {
	r := newCSRFTestRouter()
	// No prior GET: no session, no token issued.
	w := postAction(r, "", "some-token", false)
	if w.Code != http.StatusForbidden {
		t.Fatalf("POST without session: status = %d, want 403", w.Code)
	}
}

func TestCSRFPostAcceptedWithFormField(t *testing.T) {
	r := newCSRFTestRouter()
	token, cookie := tokenFromForm(t, r, "")

	w := postAction(r, cookie, token, false)
	if w.Code != http.StatusOK {
		t.Fatalf("POST with valid token: status = %d, want 200 (body: %s)", w.Code, w.Body.String())
	}
}

func TestCSRFPostAcceptedWithHeader(t *testing.T) {
	r := newCSRFTestRouter()
	token, cookie := tokenFromForm(t, r, "")

	w := postAction(r, cookie, token, true)
	if w.Code != http.StatusOK {
		t.Fatalf("POST with token in header: status = %d, want 200 (body: %s)", w.Code, w.Body.String())
	}
}

func TestCSRFSafeMethodsPassWithoutToken(t *testing.T) {
	r := newCSRFTestRouter()
	// GET and HEAD are registered routes; OPTIONS is not (gin does not
	// auto-register it), so it falls to noRoute - but in all cases the CSRF
	// middleware itself must not reject with 403.
	for _, method := range []string{http.MethodGet, http.MethodHead, http.MethodOptions} {
		req := httptest.NewRequest(method, "/form", nil)
		w := httptest.NewRecorder()
		r.ServeHTTP(w, req)
		if w.Code == http.StatusForbidden {
			t.Fatalf("%s /form: status = 403, CSRF middleware must not reject safe methods", method)
		}
	}
}