package middleware

import (
	"crypto/tls"
	"net/http"
	"net/http/httptest"
	"testing"

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

func newHeadersTestRouter() *gin.Engine {
	gin.SetMode(gin.TestMode)
	r := gin.New()
	r.Use(SecurityHeaders())
	r.GET("/", func(c *gin.Context) { c.String(http.StatusOK, "ok") })
	return r
}

func TestSecurityHeadersPresent(t *testing.T) {
	r := newHeadersTestRouter()

	// Plain HTTP request: hardening headers present, no HSTS.
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	for _, h := range []string{
		"Content-Security-Policy",
		"X-Content-Type-Options",
		"X-Frame-Options",
		"Referrer-Policy",
		"Permissions-Policy",
	} {
		if v := w.Header().Get(h); v == "" {
			t.Errorf("missing header %s", h)
		}
	}
	if w.Header().Get("X-Content-Type-Options") != "nosniff" {
		t.Errorf("X-Content-Type-Options = %q, want nosniff", w.Header().Get("X-Content-Type-Options"))
	}
	if w.Header().Get("X-Frame-Options") != "DENY" {
		t.Errorf("X-Frame-Options = %q, want DENY", w.Header().Get("X-Frame-Options"))
	}
	if w.Header().Get("Content-Security-Policy") == "" {
		t.Error("CSP header missing")
	}
	if w.Header().Get("Strict-Transport-Security") != "" {
		t.Errorf("HSTS must be absent over plain HTTP, got %q", w.Header().Get("Strict-Transport-Security"))
	}
}

func TestSecurityHeadersHSTSOverHTTPS(t *testing.T) {
	r := newHeadersTestRouter()

	// TLS request: HSTS present.
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.TLS = &tls.ConnectionState{}
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	if v := w.Header().Get("Strict-Transport-Security"); v != "max-age=31536000" {
		t.Errorf("HSTS over TLS = %q, want max-age=31536000", v)
	}

	// Behind a trusted proxy (X-Forwarded-Proto: https): HSTS present.
	req = httptest.NewRequest(http.MethodGet, "/", nil)
	req.Header.Set("X-Forwarded-Proto", "https")
	w = httptest.NewRecorder()
	r.ServeHTTP(w, req)
	if v := w.Header().Get("Strict-Transport-Security"); v != "max-age=31536000" {
		t.Errorf("HSTS behind proxy = %q, want max-age=31536000", v)
	}
}

func TestIsHTTPSRequest(t *testing.T) {
	// TLS request.
	if !IsHTTPSRequest(&gin.Context{Request: mustTLSRequest()}) {
		t.Error("TLS request must be HTTPS")
	}
	// Plain request.
	c := &gin.Context{}
	c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
	if IsHTTPSRequest(c) {
		t.Error("plain request must not be HTTPS")
	}
	// X-Forwarded-Proto: https.
	c.Request.Header.Set("X-Forwarded-Proto", "https")
	if !IsHTTPSRequest(c) {
		t.Error("X-Forwarded-Proto https must be treated as HTTPS")
	}
	// X-Forwarded-Proto: http must not trigger HTTPS behavior.
	c.Request.Header.Set("X-Forwarded-Proto", "http")
	if IsHTTPSRequest(c) {
		t.Error("X-Forwarded-Proto http must not be treated as HTTPS")
	}
}

func mustTLSRequest() *http.Request {
	r := httptest.NewRequest(http.MethodGet, "/", nil)
	r.TLS = &tls.ConnectionState{}
	return r
}