package bouncer

import (
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

// cookieSource is a CookieConfigSource over a flat map of dotted keys.
type cookieSource struct {
	values map[string]any
	env    string
}

func (s cookieSource) Lookup(path string) (any, bool) {
	v, ok := s.values[path]
	return v, ok
}

func (s cookieSource) String(path string) string {
	if v, ok := s.values[path].(string); ok {
		return v
	}
	return ""
}

func (s cookieSource) Environment() string { return s.env }

func TestLoadCookieConfigDefaults(t *testing.T) {
	cfg, err := LoadCookieConfig(nil)
	if err != nil {
		t.Fatal(err)
	}
	if cfg.Name != DefaultCookieName || cfg.Path != "/" || cfg.Domain != "" || cfg.SameSite != http.SameSiteLaxMode || cfg.Secure {
		t.Fatalf("nil source = %+v", cfg)
	}

	// An http or unset app.url infers Secure false outside production only;
	// production refuses it (TestLoadCookieConfigErrors).
	for _, tc := range []struct {
		url    string
		env    string
		secure bool
	}{
		{"https://app.example.com", "production", true},
		{"HTTPS://app.example.com", "production", true},
		{"http://localhost:8000", "development", false},
		{"", "development", false},
	} {
		cfg, err := LoadCookieConfig(cookieSource{values: map[string]any{"app.url": tc.url}, env: tc.env})
		if err != nil {
			t.Fatalf("app.url %q: %v", tc.url, err)
		}
		if cfg.Name != "token" || cfg.Path != "/" || cfg.Domain != "" || cfg.SameSite != http.SameSiteLaxMode {
			t.Fatalf("app.url %q: defaults = %+v", tc.url, cfg)
		}
		if cfg.Secure != tc.secure {
			t.Fatalf("app.url %q: Secure = %t, want %t", tc.url, cfg.Secure, tc.secure)
		}
	}

	cfg, err = LoadCookieConfig(cookieSource{values: map[string]any{
		"auth.cookie.name":      "sid",
		"auth.cookie.path":      "/api",
		"auth.cookie.domain":    "example.com",
		"auth.cookie.same_site": "Strict",
		"auth.cookie.secure":    "true",
	}, env: "production"})
	if err != nil {
		t.Fatal(err)
	}
	if cfg.Name != "sid" || cfg.Path != "/api" || cfg.Domain != "example.com" || cfg.SameSite != http.SameSiteStrictMode || !cfg.Secure {
		t.Fatalf("explicit = %+v", cfg)
	}

	cfg, err = LoadCookieConfig(cookieSource{values: map[string]any{
		"auth.cookie.same_site": "none",
		"auth.cookie.secure":    true,
	}})
	if err != nil || cfg.SameSite != http.SameSiteNoneMode || !cfg.Secure {
		t.Fatalf("none with secure = %+v, %v", cfg, err)
	}

	// secure false outside production is allowed.
	cfg, err = LoadCookieConfig(cookieSource{values: map[string]any{
		"app.url":            "https://app.example.com",
		"auth.cookie.secure": false,
	}, env: "development"})
	if err != nil || cfg.Secure {
		t.Fatalf("secure false in development = %+v, %v", cfg, err)
	}
}

func TestLoadCookieConfigErrors(t *testing.T) {
	for _, tc := range []struct {
		name   string
		values map[string]any
		env    string
		key    string
	}{
		{"invalid domain", map[string]any{"auth.cookie.domain": "exa mple.com"}, "", "auth.cookie.domain"},
		{"same_site none without secure", map[string]any{"auth.cookie.same_site": "none", "auth.cookie.secure": false}, "development", "auth.cookie.same_site"},
		{"same_site none on http app", map[string]any{"auth.cookie.same_site": "none", "app.url": "http://localhost"}, "", "auth.cookie.same_site"},
		{"secure false in production", map[string]any{"auth.cookie.secure": false}, "production", "auth.cookie.secure"},
		{"secure string false in production", map[string]any{"auth.cookie.secure": "false"}, "Production", "auth.cookie.secure"},
		{"secure inferred from http app.url in production", map[string]any{"app.url": "http://api.example.com"}, "production", "app.url"},
		{"secure inferred from unset app.url in production", map[string]any{}, "production", "auth.cookie.secure resolves to false"},
		{"secure not a bool", map[string]any{"auth.cookie.secure": "maybe"}, "", "auth.cookie.secure"},
		{"secure wrong type", map[string]any{"auth.cookie.secure": 1}, "", "auth.cookie.secure"},
		{"unknown same_site", map[string]any{"auth.cookie.same_site": "relaxed"}, "", "auth.cookie.same_site"},
		{"path without leading slash", map[string]any{"auth.cookie.path": "api"}, "", "auth.cookie.path"},
		{"invalid name", map[string]any{"auth.cookie.name": "to ken"}, "", "auth.cookie.name"},
		{"empty name", map[string]any{"auth.cookie.name": " "}, "", "auth.cookie.name"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			_, err := LoadCookieConfig(cookieSource{values: tc.values, env: tc.env})
			if err == nil {
				t.Fatal("want an error")
			}
			if !strings.Contains(err.Error(), tc.key) {
				t.Fatalf("error %q does not name %s", err, tc.key)
			}
		})
	}
}

func TestTokenCookieAttributes(t *testing.T) {
	cfg := CookieConfig{Name: "sid", Domain: "example.com", Path: "/api", Secure: true, SameSite: http.SameSiteLaxMode}
	c := cfg.TokenCookie("tok", time.Hour)
	if c.Name != "sid" || c.Value != "tok" || c.Domain != "example.com" || c.Path != "/api" {
		t.Fatalf("cookie = %+v", c)
	}
	if !c.HttpOnly || !c.Secure || c.SameSite != http.SameSiteLaxMode {
		t.Fatalf("flags = %+v", c)
	}
	if c.Valid() != nil {
		t.Fatalf("cookie invalid: %v", c.Valid())
	}
	header := c.String()
	for _, part := range []string{"sid=tok", "Path=/api", "Domain=example.com", "Max-Age=3600", "HttpOnly", "Secure", "SameSite=Lax"} {
		if !strings.Contains(header, part) {
			t.Fatalf("Set-Cookie %q lacks %q", header, part)
		}
	}

	// A zero config writes the default cookie on "/", host-only, with
	// SameSite=Lax and Secure.
	z := CookieConfig{}.TokenCookie("v", time.Minute)
	if z.Name != DefaultCookieName || z.Path != "/" || z.Domain != "" || !z.HttpOnly {
		t.Fatalf("zero config cookie = %+v", z)
	}
	if z.SameSite != http.SameSiteLaxMode || !z.Secure {
		t.Fatalf("zero config cookie SameSite %v Secure %t, want Lax and Secure", z.SameSite, z.Secure)
	}
	// A partial config gets SameSite=Lax and keeps its Secure value.
	p := CookieConfig{Name: "sid"}.TokenCookie("v", time.Minute)
	if p.SameSite != http.SameSiteLaxMode || p.Secure {
		t.Fatalf("partial config cookie SameSite %v Secure %t, want Lax and not Secure", p.SameSite, p.Secure)
	}
	if strings.Contains(z.String(), "Domain=") {
		t.Fatalf("host-only cookie carries a domain: %s", z.String())
	}
}

func TestTokenCookieMaxAgeIsRefreshWindow(t *testing.T) {
	const (
		accessTTL  = 60 * time.Minute
		refreshTTL = 20160 * time.Minute
	)
	c := CookieConfig{}.TokenCookie("tok", refreshTTL)
	if c.MaxAge != 1209600 {
		t.Fatalf("MaxAge = %d, want 1209600 (refresh_ttl*60)", c.MaxAge)
	}
	if c.MaxAge == int(accessTTL/time.Second) {
		t.Fatal("MaxAge equals the access ttl")
	}
	if d := time.Until(c.Expires); d < refreshTTL-time.Minute || d > refreshTTL+time.Minute {
		t.Fatalf("Expires in %s, want about %s", d, refreshTTL)
	}
}

func TestForgetTokenCookie(t *testing.T) {
	cfg := CookieConfig{Name: "sid", Domain: "example.com", Path: "/api", Secure: true, SameSite: http.SameSiteNoneMode}
	set := cfg.TokenCookie("tok", time.Hour)
	f := cfg.ForgetTokenCookie()
	if f.Name != set.Name || f.Domain != set.Domain || f.Path != set.Path || f.Secure != set.Secure || f.SameSite != set.SameSite || f.HttpOnly != set.HttpOnly {
		t.Fatalf("forget %+v does not match set %+v", f, set)
	}
	if f.Value != "" || f.MaxAge != -1 || !f.Expires.Before(time.Now()) {
		t.Fatalf("forget cookie = %+v", f)
	}
	if !strings.Contains(f.String(), "Max-Age=0") {
		t.Fatalf("forget header %q does not expire the cookie", f.String())
	}
}

// recordWriter captures what a cookie handler passes to the app's writer.
type recordWriter struct {
	calls int
	res   CookieResult
}

func (rw *recordWriter) write(w http.ResponseWriter, _ *http.Request, res CookieResult) {
	rw.calls++
	rw.res = res
	w.WriteHeader(res.Status)
	_, _ = w.Write([]byte(res.Token))
}

func refreshRequest(name, value string) *http.Request {
	r := httptest.NewRequest(http.MethodPost, "https://api.example.com/auth/refresh", nil)
	if name != "" {
		r.AddCookie(&http.Cookie{Name: name, Value: value})
	}
	return r
}

func TestCookieRefreshHandlerRotatesAndBlacklists(t *testing.T) {
	const (
		issuer     = "https://api.example.com/auth/login"
		refreshTTL = 20160 * time.Minute
	)
	old, oldJTI, err := Mint(secret, "42", issuer, time.Hour)
	if err != nil {
		t.Fatal(err)
	}
	bl := NewMemoryBlacklist()
	cfg := CookieConfig{Name: "token", Domain: "example.com", Path: "/", Secure: true, SameSite: http.SameSiteLaxMode}
	rw := &recordWriter{}
	h, err := CookieRefreshHandler(CookieRefreshOptions{
		Cookie:     cfg,
		RefreshTTL: refreshTTL,
		Refresh: func(_ context.Context, token string) (string, error) {
			return Refresh(secret, token, refreshTTL, bl, 0, "https://api.example.com/auth/refresh")
		},
		Write: rw.write,
	})
	if err != nil {
		t.Fatal(err)
	}
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, refreshRequest("token", old))

	if rec.Code != http.StatusOK || rw.calls != 1 || rw.res.Status != http.StatusOK || rw.res.Err != nil {
		t.Fatalf("status %d, writer %d calls, res %+v", rec.Code, rw.calls, rw.res)
	}
	if rw.res.Token == "" || rw.res.Token == old || rec.Body.String() != rw.res.Token {
		t.Fatalf("writer token %q", rw.res.Token)
	}
	cookies := rec.Result().Cookies()
	if len(cookies) != 1 {
		t.Fatalf("Set-Cookie count = %d", len(cookies))
	}
	c := cookies[0]
	if c.Name != "token" || c.Value != rw.res.Token || c.Domain != "example.com" || !c.HttpOnly || !c.Secure || c.MaxAge != 1209600 {
		t.Fatalf("rotated cookie = %+v", c)
	}
	claims := decodeClaims(t, c.Value, secret)
	if claims["jti"] == oldJTI || claims["sub"] != "42" {
		t.Fatalf("new claims = %#v", claims)
	}
	blocked, err := bl.IsBlacklisted(t.Context(), oldJTI)
	if err != nil || !blocked {
		t.Fatalf("old jti blacklisted=%t err=%v", blocked, err)
	}

	// The old token no longer refreshes.
	rw2 := &recordWriter{}
	h2, _ := CookieRefreshHandler(CookieRefreshOptions{
		Cookie: cfg, RefreshTTL: refreshTTL,
		Refresh: func(_ context.Context, token string) (string, error) {
			return Refresh(secret, token, refreshTTL, bl, 0, "https://api.example.com/auth/refresh")
		},
		Write: rw2.write,
	})
	rec2 := httptest.NewRecorder()
	h2.ServeHTTP(rec2, refreshRequest("token", old))
	if rec2.Code != http.StatusUnauthorized || len(rec2.Result().Cookies()) != 0 {
		t.Fatalf("replayed refresh: %d, cookies %v", rec2.Code, rec2.Result().Cookies())
	}
}

func TestCookieRefreshHandlerMissingCookie(t *testing.T) {
	for _, tc := range []struct{ name, cookie, value string }{
		{"no cookie", "", ""},
		{"other cookie", "session", "abc"},
		{"empty cookie", "token", ""},
		{"blank cookie", "token", "   "},
	} {
		t.Run(tc.name, func(t *testing.T) {
			called := false
			rw := &recordWriter{}
			h, err := CookieRefreshHandler(CookieRefreshOptions{
				RefreshTTL: time.Hour,
				Refresh: func(context.Context, string) (string, error) {
					called = true
					return "x", nil
				},
				Write: rw.write,
			})
			if err != nil {
				t.Fatal(err)
			}
			rec := httptest.NewRecorder()
			h.ServeHTTP(rec, refreshRequest(tc.cookie, tc.value))
			if called {
				t.Fatal("Refresh called without a token cookie")
			}
			if rec.Code != http.StatusUnauthorized || !errors.Is(rw.res.Err, ErrNoTokenCookie) {
				t.Fatalf("status %d err %v", rec.Code, rw.res.Err)
			}
			if rec.Header().Get("Set-Cookie") != "" {
				t.Fatalf("Set-Cookie sent: %q", rec.Header().Get("Set-Cookie"))
			}
		})
	}
}

func TestCookieRefreshHandlerRefreshError(t *testing.T) {
	refusal := errors.New("token expired past refresh window")
	rw := &recordWriter{}
	h, err := CookieRefreshHandler(CookieRefreshOptions{
		Cookie:     CookieConfig{Name: "sid"},
		RefreshTTL: time.Hour,
		Refresh: func(_ context.Context, token string) (string, error) {
			if token != "stale" {
				t.Fatalf("Refresh got %q", token)
			}
			return "", refusal
		},
		Write: rw.write,
	})
	if err != nil {
		t.Fatal(err)
	}
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, refreshRequest("sid", "stale"))
	if rec.Code != http.StatusUnauthorized || !errors.Is(rw.res.Err, refusal) || rw.res.Token != "" {
		t.Fatalf("status %d res %+v", rec.Code, rw.res)
	}
	if rec.Header().Get("Set-Cookie") != "" {
		t.Fatalf("cookie touched on a refused refresh: %q", rec.Header().Get("Set-Cookie"))
	}
}

func TestCookieLogoutHandler(t *testing.T) {
	cfg := CookieConfig{Name: "token", Domain: "example.com", Path: "/", Secure: true, SameSite: http.SameSiteLaxMode}
	revokeErr := errors.New("blacklist down")
	for _, tc := range []struct {
		name      string
		cookie    string
		revokeErr error
		wantCall  bool
	}{
		{"revokes the cookie token", "tok-1", nil, true},
		{"revoke failure still forgets", "tok-2", revokeErr, true},
		{"no cookie still forgets", "", nil, false},
	} {
		t.Run(tc.name, func(t *testing.T) {
			var got []string
			rw := &recordWriter{}
			h, err := CookieLogoutHandler(CookieLogoutOptions{
				Cookie: cfg,
				Revoke: func(_ context.Context, token string) error {
					got = append(got, token)
					return tc.revokeErr
				},
				Write: rw.write,
			})
			if err != nil {
				t.Fatal(err)
			}
			name := ""
			if tc.cookie != "" {
				name = "token"
			}
			rec := httptest.NewRecorder()
			h.ServeHTTP(rec, refreshRequest(name, tc.cookie))
			if tc.wantCall && (len(got) != 1 || got[0] != tc.cookie) {
				t.Fatalf("Revoke calls = %v", got)
			}
			if !tc.wantCall && len(got) != 0 {
				t.Fatalf("Revoke called without a cookie: %v", got)
			}
			if rec.Code != http.StatusOK || rw.res.Status != http.StatusOK || !errors.Is(rw.res.Err, tc.revokeErr) {
				t.Fatalf("status %d res %+v", rec.Code, rw.res)
			}
			cookies := rec.Result().Cookies()
			if len(cookies) != 1 {
				t.Fatalf("Set-Cookie count = %d", len(cookies))
			}
			f := cookies[0]
			if f.Name != "token" || f.Value != "" || f.MaxAge != -1 || f.Domain != "example.com" || f.Path != "/" || !f.Secure || !f.HttpOnly {
				t.Fatalf("forget cookie = %+v", f)
			}
		})
	}
}

func TestCookieHandlerOptionsValidated(t *testing.T) {
	refresh := func(context.Context, string) (string, error) { return "", nil }
	revoke := func(context.Context, string) error { return nil }
	write := func(http.ResponseWriter, *http.Request, CookieResult) {}

	for _, tc := range []struct {
		name string
		opts CookieRefreshOptions
		want string
	}{
		{"nil refresh", CookieRefreshOptions{RefreshTTL: time.Hour, Write: write}, "Refresh"},
		{"nil write", CookieRefreshOptions{RefreshTTL: time.Hour, Refresh: refresh}, "Write"},
		{"zero ttl", CookieRefreshOptions{Refresh: refresh, Write: write}, "RefreshTTL"},
		{"negative ttl", CookieRefreshOptions{RefreshTTL: -time.Second, Refresh: refresh, Write: write}, "RefreshTTL"},
	} {
		h, err := CookieRefreshHandler(tc.opts)
		if err == nil || h != nil || !strings.Contains(err.Error(), tc.want) {
			t.Fatalf("refresh %s: handler %v err %v", tc.name, h, err)
		}
	}
	if h, err := CookieRefreshHandler(CookieRefreshOptions{RefreshTTL: time.Hour, Refresh: refresh, Write: write}); err != nil || h == nil {
		t.Fatalf("valid refresh options: %v", err)
	}

	for _, tc := range []struct {
		name string
		opts CookieLogoutOptions
		want string
	}{
		{"nil revoke", CookieLogoutOptions{Write: write}, "Revoke"},
		{"nil write", CookieLogoutOptions{Revoke: revoke}, "Write"},
	} {
		h, err := CookieLogoutHandler(tc.opts)
		if err == nil || h != nil || !strings.Contains(err.Error(), tc.want) {
			t.Fatalf("logout %s: handler %v err %v", tc.name, h, err)
		}
	}
	if h, err := CookieLogoutHandler(CookieLogoutOptions{Revoke: revoke, Write: write}); err != nil || h == nil {
		t.Fatalf("valid logout options: %v", err)
	}
}

func TestLoadCookieConfigProductionInferredSecure(t *testing.T) {
	for _, tc := range []struct {
		name    string
		env     string
		values  map[string]any
		wantErr bool
		secure  bool
	}{
		{"production, http app.url", "production", map[string]any{"app.url": "http://api.example.com"}, true, false},
		{"production, unset app.url", "production", map[string]any{}, true, false},
		{"production, malformed app.url", "production", map[string]any{"app.url": "://not a url"}, true, false},
		{"production, https app.url", "production", map[string]any{"app.url": "https://api.example.com"}, false, true},
		{"production, http app.url and secure true", "production", map[string]any{"app.url": "http://api.example.com", "auth.cookie.secure": true}, false, true},
		{"development, http app.url", "development", map[string]any{"app.url": "http://localhost:8000"}, false, false},
	} {
		t.Run(tc.name, func(t *testing.T) {
			cfg, err := LoadCookieConfig(cookieSource{values: tc.values, env: tc.env})
			if tc.wantErr {
				if err == nil {
					t.Fatalf("config = %+v, want an error", cfg)
				}
				for _, key := range []string{"auth.cookie.secure", "app.url"} {
					if !strings.Contains(err.Error(), key) {
						t.Errorf("error %q does not name %s", err, key)
					}
				}
				return
			}
			if err != nil {
				t.Fatal(err)
			}
			if cfg.Secure != tc.secure {
				t.Fatalf("Secure = %t, want %t", cfg.Secure, tc.secure)
			}
		})
	}
}

func TestZeroCookieConfigDefaults(t *testing.T) {
	var zero CookieConfig
	for name, c := range map[string]*http.Cookie{
		"TokenCookie":       zero.TokenCookie("tok", time.Hour),
		"ForgetTokenCookie": zero.ForgetTokenCookie(),
	} {
		if c.Name != DefaultCookieName || c.Path != "/" || c.Domain != "" || !c.HttpOnly || c.SameSite != http.SameSiteLaxMode || !c.Secure {
			t.Errorf("%s of the zero config = %+v, want token on / with HttpOnly, SameSite=Lax and Secure", name, c)
		}
	}
	partial := CookieConfig{Name: "sid"}.TokenCookie("tok", time.Hour)
	if partial.SameSite != http.SameSiteLaxMode || partial.Secure {
		t.Errorf("{Name: sid} cookie SameSite %v Secure %t, want Lax and not Secure", partial.SameSite, partial.Secure)
	}

	write := func(w http.ResponseWriter, _ *http.Request, res CookieResult) { w.WriteHeader(res.Status) }
	refresh, err := CookieRefreshHandler(CookieRefreshOptions{
		RefreshTTL: time.Hour,
		Refresh:    func(context.Context, string) (string, error) { return "rotated", nil },
		Write:      write,
	})
	if err != nil {
		t.Fatal(err)
	}
	logout, err := CookieLogoutHandler(CookieLogoutOptions{
		Revoke: func(context.Context, string) error { return nil },
		Write:  write,
	})
	if err != nil {
		t.Fatal(err)
	}
	for name, h := range map[string]http.Handler{"refresh": refresh, "logout": logout} {
		req := httptest.NewRequest(http.MethodPost, "/auth/"+name, nil)
		req.AddCookie(&http.Cookie{Name: DefaultCookieName, Value: "old"})
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, req)
		set := rec.Header().Get("Set-Cookie")
		if set == "" {
			t.Fatalf("%s: no Set-Cookie (status %d)", name, rec.Code)
		}
		for _, want := range []string{DefaultCookieName + "=", "Path=/", "HttpOnly", "SameSite=Lax", "Secure"} {
			if !strings.Contains(set, want) {
				t.Errorf("%s: Set-Cookie %q lacks %s", name, set, want)
			}
		}
	}
}