package middleware

import "net/http"

// SecurityHeadersConfig configures the SecurityHeaders middleware. All headers
// are opt-in via cfg so they can be tuned (or disabled) without code changes;
// the defaults are safe for an SPA + WebSocket app. HSTS is gated on TLS — it is
// only emitted when TLSEnabled is true, since advertising HSTS over plain HTTP
// either does nothing or (once cached) breaks local http access.
type SecurityHeadersConfig struct {
	Enabled    bool
	TLSEnabled bool
	// HSTSMaxAgeSeconds is the Strict-Transport-Security max-age. 0 → default
	// (1 year) when TLSEnabled.
	HSTSMaxAgeSeconds int
	// FrameOptions sets X-Frame-Options; empty → "DENY".
	FrameOptions string
	// ReferrerPolicy sets Referrer-Policy; empty → "no-referrer".
	ReferrerPolicy string
	// ContentSecurityPolicy sets CSP; empty → header omitted (CSP can break the
	// SPA + WS handshake if mis-set, so it is off by default and opt-in).
	ContentSecurityPolicy string
}

// SecurityHeaders returns middleware that stamps standard security response
// headers. When cfg.Enabled is false it is a transparent pass-through, so it can
// be wired unconditionally and toggled by config.
func SecurityHeaders(cfg SecurityHeadersConfig) func(http.Handler) http.Handler {
	frame := cfg.FrameOptions
	if frame == "" {
		frame = "DENY"
	}
	referrer := cfg.ReferrerPolicy
	if referrer == "" {
		referrer = "no-referrer"
	}
	hstsAge := cfg.HSTSMaxAgeSeconds
	if hstsAge <= 0 {
		hstsAge = 31536000 // 1 year
	}
	hstsValue := "max-age=" + itoa(hstsAge) + "; includeSubDomains"

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if !cfg.Enabled {
				next.ServeHTTP(w, r)
				return
			}
			h := w.Header()
			h.Set("X-Content-Type-Options", "nosniff")
			h.Set("X-Frame-Options", frame)
			h.Set("Referrer-Policy", referrer)
			if cfg.ContentSecurityPolicy != "" {
				h.Set("Content-Security-Policy", cfg.ContentSecurityPolicy)
			}
			if cfg.TLSEnabled {
				h.Set("Strict-Transport-Security", hstsValue)
			}
			next.ServeHTTP(w, r)
		})
	}
}

// itoa is a tiny non-allocating-ish int->string for positive values, avoiding a
// strconv import for a single use.
func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	var buf [20]byte
	i := len(buf)
	for n > 0 {
		i--
		buf[i] = byte('0' + n%10)
		n /= 10
	}
	return string(buf[i:])
}