package authapi

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

	"github.com/itworx/pulse/internal/auth"
)

const testClientID = "pulse-test-client"

type recordingSessions struct {
	mu         sync.Mutex
	principals []auth.Principal
	failure    error
}

func (sessions *recordingSessions) Issue(response http.ResponseWriter, principal auth.Principal, _ time.Time) error {
	sessions.mu.Lock()
	defer sessions.mu.Unlock()
	if sessions.failure != nil {
		return sessions.failure
	}
	sessions.principals = append(sessions.principals, principal)
	http.SetCookie(response, &http.Cookie{Name: "pulse_session", Value: "issued", Path: "/", HttpOnly: true})
	return nil
}

func (sessions *recordingSessions) issued() []auth.Principal {
	sessions.mu.Lock()
	defer sessions.mu.Unlock()
	return append([]auth.Principal(nil), sessions.principals...)
}

type harness struct {
	handler  *Handler
	idp      *fakeIdP
	sessions *recordingSessions
	clock    func() time.Time
	offset   *time.Duration
}

func newHarness(t *testing.T, mutate func(*Options)) *harness {
	t.Helper()
	idp := newFakeIdP(t, testClientID)
	sessions := &recordingSessions{}
	offset := time.Duration(0)
	options := Options{
		OIDC: auth.OIDCConfig{
			Issuer:       idp.server.URL,
			ClientID:     testClientID,
			ClientSecret: "test-secret",
			RedirectURL:  "https://pulse.example/auth/callback",
		},
		RoleMapping: map[string]auth.Role{"pulse-viewer": auth.RoleViewer, "pulse-operator": auth.RoleOperator, "pulse-admin": auth.RoleAdministrator},
		Sessions:    sessions,
		Now:         func() time.Time { return time.Now().UTC().Add(offset) },
	}
	if mutate != nil {
		mutate(&options)
	}
	handler, err := New(options)
	if err != nil {
		t.Fatalf("New: %v", err)
	}
	return &harness{handler: handler, idp: idp, sessions: sessions, clock: options.Now, offset: &offset}
}

// begin runs GET /auth/login and returns the flow cookie value and the query the
// browser would have sent to the IdP.
func (h *harness) begin(t *testing.T, target string) (string, url.Values, *httptest.ResponseRecorder) {
	t.Helper()
	request := httptest.NewRequest(http.MethodGet, target, nil)
	response := httptest.NewRecorder()
	h.handler.LoginHandler().ServeHTTP(response, request)
	if response.Code != http.StatusFound {
		t.Fatalf("login status = %d, want %d", response.Code, http.StatusFound)
	}
	authorizationURL, err := url.Parse(response.Header().Get("Location"))
	if err != nil {
		t.Fatalf("parse authorization URL: %v", err)
	}
	query := authorizationURL.Query()
	h.idp.configure(func(idp *fakeIdP) {
		idp.nonce = query.Get("nonce")
		idp.expectedChallenge = query.Get("code_challenge")
	})
	return flowCookie(t, response), query, response
}

// complete runs GET /auth/callback with the supplied cookie and query.
func (h *harness) complete(t *testing.T, cookie string, query url.Values) *httptest.ResponseRecorder {
	t.Helper()
	request := httptest.NewRequest(http.MethodGet, "/auth/callback?"+query.Encode(), nil)
	if cookie != "" {
		request.AddCookie(&http.Cookie{Name: flowCookieName, Value: cookie})
	}
	response := httptest.NewRecorder()
	h.handler.CallbackHandler().ServeHTTP(response, request)
	return response
}

func flowCookie(t *testing.T, response *httptest.ResponseRecorder) string {
	t.Helper()
	for _, cookie := range response.Result().Cookies() {
		if cookie.Name == flowCookieName {
			return cookie.Value
		}
	}
	t.Fatal("flow cookie was not set")
	return ""
}

func cookieByName(response *httptest.ResponseRecorder, name string) *http.Cookie {
	for _, cookie := range response.Result().Cookies() {
		if cookie.Name == name {
			return cookie
		}
	}
	return nil
}

func errorReason(t *testing.T, response *httptest.ResponseRecorder) string {
	t.Helper()
	if response.Code != http.StatusFound {
		t.Fatalf("status = %d, want %d", response.Code, http.StatusFound)
	}
	location, err := url.Parse(response.Header().Get("Location"))
	if err != nil {
		t.Fatalf("parse location: %v", err)
	}
	if location.Path != defaultErrorPath {
		t.Fatalf("location path = %q, want %q", location.Path, defaultErrorPath)
	}
	return location.Query().Get("reason")
}

func TestLoginStartsBoundedServerSideFlow(t *testing.T) {
	harness := newHarness(t, func(options *Options) { options.Secure = true })
	cookie, query, response := harness.begin(t, "/auth/login")

	for _, key := range []string{"state", "nonce", "code_challenge", "client_id", "redirect_uri"} {
		if query.Get(key) == "" {
			t.Fatalf("authorization URL missing %s", key)
		}
	}
	if query.Get("code_challenge_method") != "S256" {
		t.Fatalf("code_challenge_method = %q", query.Get("code_challenge_method"))
	}
	if strings.Contains(response.Header().Get("Location"), cookie) {
		t.Fatal("flow identifier leaked into the authorization URL")
	}
	if cookieByName(response, "pulse_session") != nil {
		t.Fatal("login issued a session")
	}
	flowCookie := cookieByName(response, flowCookieName)
	if !flowCookie.HttpOnly || !flowCookie.Secure || flowCookie.SameSite != http.SameSiteLaxMode {
		t.Fatalf("flow cookie attributes are unsafe: %#v", flowCookie)
	}
	if flowCookie.MaxAge <= 0 || flowCookie.MaxAge > int(defaultFlowTTL.Seconds()) {
		t.Fatalf("flow cookie MaxAge = %d", flowCookie.MaxAge)
	}
	if harness.handler.flows.size() != 1 {
		t.Fatalf("pending flows = %d, want 1", harness.handler.flows.size())
	}
	pending, ok := harness.handler.flows.take(cookie, harness.clock())
	if !ok || pending.authorization.State != query.Get("state") || pending.authorization.Nonce != query.Get("nonce") || pending.authorization.CodeVerifier == "" {
		t.Fatal("flow state was not stored server-side")
	}
}

func TestCallbackHappyPathIssuesSession(t *testing.T) {
	harness := newHarness(t, nil)
	cookie, query, _ := harness.begin(t, "/auth/login?redirect=%2Fincidents")
	response := harness.complete(t, cookie, url.Values{"state": {query.Get("state")}, "code": {"authorization-code"}})

	if response.Code != http.StatusFound || response.Header().Get("Location") != "/incidents" {
		t.Fatalf("status = %d, location = %q", response.Code, response.Header().Get("Location"))
	}
	issued := harness.sessions.issued()
	if len(issued) != 1 || issued[0].Subject != "user-1" || issued[0].Role != auth.RoleOperator {
		t.Fatalf("issued sessions = %#v", issued)
	}
	if verifier := harness.idp.codeVerifier(); verifier == "" {
		t.Fatal("PKCE verifier was not sent to the token endpoint")
	}
	cleared := cookieByName(response, flowCookieName)
	if cleared == nil || cleared.MaxAge >= 0 || cleared.Value != "" {
		t.Fatalf("flow cookie was not cleared: %#v", cleared)
	}
	if harness.handler.flows.size() != 0 {
		t.Fatal("flow state survived the callback")
	}
}

func TestCallbackWorksWithRealSessionManager(t *testing.T) {
	manager := auth.NewSessionManager("pulse_session", time.Hour, false)
	harness := newHarness(t, func(options *Options) { options.Sessions = manager })
	cookie, query, _ := harness.begin(t, "/auth/login")
	response := harness.complete(t, cookie, url.Values{"state": {query.Get("state")}, "code": {"authorization-code"}})

	sessionCookie := cookieByName(response, "pulse_session")
	if sessionCookie == nil {
		t.Fatal("session cookie was not set")
	}
	next := httptest.NewRequest(http.MethodGet, "/api/v1/system/status", nil)
	next.AddCookie(sessionCookie)
	principal, ok := manager.Principal(next, time.Now().UTC())
	if !ok || principal.Subject != "user-1" || principal.Role != auth.RoleOperator {
		t.Fatalf("principal = %#v, ok = %v", principal, ok)
	}
}

func TestCallbackFailurePathsIssueNoSession(t *testing.T) {
	for _, test := range []struct {
		name string
		// arrange starts a flow and returns the callback cookie and query.
		arrange func(t *testing.T, h *harness) (string, url.Values)
		reason  string
		// pendingFlows is what may still be stored afterwards: an untouched flow
		// stays pending until it expires, an identified one is always destroyed.
		pendingFlows int
	}{
		{
			name: "missing flow cookie",
			arrange: func(t *testing.T, h *harness) (string, url.Values) {
				_, query, _ := h.begin(t, "/auth/login")
				return "", url.Values{"state": {query.Get("state")}, "code": {"authorization-code"}}
			},
			reason:       reasonInvalidRequest,
			pendingFlows: 1,
		},
		{
			name: "unknown flow identifier",
			arrange: func(t *testing.T, h *harness) (string, url.Values) {
				_, query, _ := h.begin(t, "/auth/login")
				return "not-a-known-flow", url.Values{"state": {query.Get("state")}, "code": {"authorization-code"}}
			},
			reason:       reasonExpired,
			pendingFlows: 1,
		},
		{
			name: "replayed flow identifier",
			arrange: func(t *testing.T, h *harness) (string, url.Values) {
				cookie, query, _ := h.begin(t, "/auth/login")
				values := url.Values{"state": {query.Get("state")}, "code": {"authorization-code"}}
				if first := h.complete(t, cookie, values); first.Header().Get("Location") != "/" {
					t.Fatalf("first callback did not succeed: %q", first.Header().Get("Location"))
				}
				h.sessions.mu.Lock()
				h.sessions.principals = nil
				h.sessions.mu.Unlock()
				return cookie, values
			},
			reason: reasonExpired,
		},
		{
			name: "expired flow",
			arrange: func(t *testing.T, h *harness) (string, url.Values) {
				cookie, query, _ := h.begin(t, "/auth/login")
				*h.offset = defaultFlowTTL + time.Minute
				return cookie, url.Values{"state": {query.Get("state")}, "code": {"authorization-code"}}
			},
			reason: reasonExpired,
		},
		{
			name: "state mismatch",
			arrange: func(t *testing.T, h *harness) (string, url.Values) {
				cookie, _, _ := h.begin(t, "/auth/login")
				return cookie, url.Values{"state": {"forged-state"}, "code": {"authorization-code"}}
			},
			reason: reasonInvalidRequest,
		},
		{
			name: "missing authorization code",
			arrange: func(t *testing.T, h *harness) (string, url.Values) {
				cookie, query, _ := h.begin(t, "/auth/login")
				return cookie, url.Values{"state": {query.Get("state")}}
			},
			reason: reasonInvalidRequest,
		},
		{
			name: "nonce mismatch",
			arrange: func(t *testing.T, h *harness) (string, url.Values) {
				cookie, query, _ := h.begin(t, "/auth/login")
				h.idp.configure(func(idp *fakeIdP) { idp.nonce = "replayed-nonce" })
				return cookie, url.Values{"state": {query.Get("state")}, "code": {"authorization-code"}}
			},
			reason: reasonInvalidRequest,
		},
		{
			name: "provider access denied",
			arrange: func(t *testing.T, h *harness) (string, url.Values) {
				cookie, query, _ := h.begin(t, "/auth/login")
				return cookie, url.Values{"state": {query.Get("state")}, "error": {"access_denied"}, "error_description": {"<script>alert(1)</script> denied by policy"}}
			},
			reason: reasonDenied,
		},
		{
			name: "provider error response",
			arrange: func(t *testing.T, h *harness) (string, url.Values) {
				cookie, query, _ := h.begin(t, "/auth/login")
				return cookie, url.Values{"state": {query.Get("state")}, "error": {"server_error"}}
			},
			reason: reasonProviderUnavailable,
		},
		{
			name: "token exchange failure",
			arrange: func(t *testing.T, h *harness) (string, url.Values) {
				cookie, query, _ := h.begin(t, "/auth/login")
				h.idp.configure(func(idp *fakeIdP) { idp.tokenFails = true })
				return cookie, url.Values{"state": {query.Get("state")}, "code": {"authorization-code"}}
			},
			reason: reasonProviderUnavailable,
		},
		{
			name: "id token missing",
			arrange: func(t *testing.T, h *harness) (string, url.Values) {
				cookie, query, _ := h.begin(t, "/auth/login")
				h.idp.configure(func(idp *fakeIdP) { idp.omitIDToken = true })
				return cookie, url.Values{"state": {query.Get("state")}, "code": {"authorization-code"}}
			},
			reason: reasonProviderUnavailable,
		},
		{
			name: "id token from wrong issuer",
			arrange: func(t *testing.T, h *harness) (string, url.Values) {
				cookie, query, _ := h.begin(t, "/auth/login")
				h.idp.configure(func(idp *fakeIdP) { idp.issuerOverride = "https://attacker.example" })
				return cookie, url.Values{"state": {query.Get("state")}, "code": {"authorization-code"}}
			},
			reason: reasonInvalidRequest,
		},
		{
			name: "no mapped role",
			arrange: func(t *testing.T, h *harness) (string, url.Values) {
				cookie, query, _ := h.begin(t, "/auth/login")
				h.idp.configure(func(idp *fakeIdP) { idp.groups = []string{"some-other-group"} })
				return cookie, url.Values{"state": {query.Get("state")}, "code": {"authorization-code"}}
			},
			reason: reasonNotAuthorized,
		},
		{
			name: "session issue failure",
			arrange: func(t *testing.T, h *harness) (string, url.Values) {
				cookie, query, _ := h.begin(t, "/auth/login")
				h.sessions.mu.Lock()
				h.sessions.failure = errors.New("session store unavailable")
				h.sessions.mu.Unlock()
				return cookie, url.Values{"state": {query.Get("state")}, "code": {"authorization-code"}}
			},
			reason: reasonUnavailable,
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			harness := newHarness(t, nil)
			cookie, query := test.arrange(t, harness)
			response := harness.complete(t, cookie, query)

			if reason := errorReason(t, response); reason != test.reason {
				t.Fatalf("reason = %q, want %q", reason, test.reason)
			}
			if issued := harness.sessions.issued(); len(issued) != 0 {
				t.Fatalf("a session was issued on a failure path: %#v", issued)
			}
			if cookieByName(response, "pulse_session") != nil {
				t.Fatal("a session cookie was set on a failure path")
			}
			if harness.handler.flows.size() != test.pendingFlows {
				t.Fatalf("pending flows = %d, want %d", harness.handler.flows.size(), test.pendingFlows)
			}
			body := response.Body.String()
			for _, forbidden := range []string{"<script>", "denied by policy", "server_error", "access_denied", "authorization-code"} {
				if strings.Contains(body, forbidden) {
					t.Fatalf("response body reflected untrusted text %q: %s", forbidden, body)
				}
			}
		})
	}
}

func TestAuditFailureBlocksSession(t *testing.T) {
	harness := newHarness(t, func(options *Options) {
		options.Audit = func(context.Context, string, string) error { return errors.New("audit unavailable") }
	})
	cookie, query, _ := harness.begin(t, "/auth/login")
	response := harness.complete(t, cookie, url.Values{"state": {query.Get("state")}, "code": {"authorization-code"}})
	if reason := errorReason(t, response); reason != reasonUnavailable {
		t.Fatalf("reason = %q, want %q", reason, reasonUnavailable)
	}
	if issued := harness.sessions.issued(); len(issued) != 0 {
		t.Fatalf("session issued despite audit failure: %#v", issued)
	}
}

func TestAuditRecordsSuccessfulLogin(t *testing.T) {
	var actor, result string
	harness := newHarness(t, func(options *Options) {
		options.Audit = func(_ context.Context, recordedActor, recordedResult string) error {
			actor, result = recordedActor, recordedResult
			return nil
		}
	})
	cookie, query, _ := harness.begin(t, "/auth/login")
	harness.complete(t, cookie, url.Values{"state": {query.Get("state")}, "code": {"authorization-code"}})
	if actor != "user-1" || result != "success" {
		t.Fatalf("audit actor = %q, result = %q", actor, result)
	}
}

func TestPostLoginRedirectRejectsUnsafeTargets(t *testing.T) {
	for _, test := range []struct {
		name     string
		redirect string
		want     string
	}{
		{"in-app path", "/incidents", "/incidents"},
		{"in-app path with query and fragment", "/dashboards/1?tab=live#top", "/dashboards/1?tab=live#top"},
		{"absent", "", "/"},
		{"protocol relative", "//evil.example/phish", "/"},
		{"triple slash", "///evil.example", "/"},
		{"absolute http", "http://evil.example", "/"},
		{"absolute https", "https://evil.example/x", "/"},
		{"scheme relative backslash", "/\\evil.example", "/"},
		{"backslashes", "\\\\evil.example", "/"},
		{"relative path", "incidents", "/"},
		{"javascript scheme", "javascript:alert(1)", "/"},
		{"data scheme", "data:text/html,<script>alert(1)</script>", "/"},
		{"newline injection", "/incidents\r\nSet-Cookie: x=1", "/"},
		{"userinfo authority", "//user:pass@evil.example/", "/"},
		{"overlong", "/" + strings.Repeat("a", maxRedirectLength), "/"},
	} {
		t.Run(test.name, func(t *testing.T) {
			harness := newHarness(t, nil)
			cookie, query, _ := harness.begin(t, "/auth/login?redirect="+url.QueryEscape(test.redirect))
			response := harness.complete(t, cookie, url.Values{"state": {query.Get("state")}, "code": {"authorization-code"}})
			if response.Code != http.StatusFound {
				t.Fatalf("status = %d", response.Code)
			}
			if location := response.Header().Get("Location"); location != test.want {
				t.Fatalf("location = %q, want %q", location, test.want)
			}
		})
	}
}

func TestConcurrentFlowCreationIsSafeAndBounded(t *testing.T) {
	const workers = 64
	harness := newHarness(t, func(options *Options) { options.MaxFlows = 16 })
	var wait sync.WaitGroup
	cookies := make([]string, workers)
	for index := range workers {
		wait.Add(1)
		go func() {
			defer wait.Done()
			request := httptest.NewRequest(http.MethodGet, "/auth/login", nil)
			response := httptest.NewRecorder()
			harness.handler.LoginHandler().ServeHTTP(response, request)
			for _, cookie := range response.Result().Cookies() {
				if cookie.Name == flowCookieName {
					cookies[index] = cookie.Value
				}
			}
		}()
	}
	wait.Wait()

	unique := make(map[string]struct{}, workers)
	for _, cookie := range cookies {
		if cookie == "" {
			t.Fatal("a concurrent login produced no flow cookie")
		}
		unique[cookie] = struct{}{}
	}
	if len(unique) != workers {
		t.Fatalf("unique flow identifiers = %d, want %d", len(unique), workers)
	}
	if size := harness.handler.flows.size(); size > 16 {
		t.Fatalf("pending flows = %d, want at most 16", size)
	}
}

func TestDiscoveryFailureRedirectsSafely(t *testing.T) {
	closed := httptest.NewServer(http.NewServeMux())
	issuer := closed.URL
	closed.Close()
	handler, err := New(Options{
		OIDC:     auth.OIDCConfig{Issuer: issuer, ClientID: testClientID, RedirectURL: "https://pulse.example/auth/callback"},
		Sessions: &recordingSessions{},
	})
	if err != nil {
		t.Fatalf("New: %v", err)
	}
	response := httptest.NewRecorder()
	handler.LoginHandler().ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/auth/login", nil))
	if reason := errorReason(t, response); reason != reasonProviderUnavailable {
		t.Fatalf("reason = %q, want %q", reason, reasonProviderUnavailable)
	}
	if body := response.Body.String(); strings.Contains(body, issuer) {
		t.Fatalf("response leaked the issuer: %s", body)
	}
}

func TestNonGetMethodsAreRejected(t *testing.T) {
	harness := newHarness(t, nil)
	for _, test := range []struct {
		name    string
		handler http.Handler
		target  string
	}{
		{"login", harness.handler.LoginHandler(), "/auth/login"},
		{"callback", harness.handler.CallbackHandler(), "/auth/callback"},
	} {
		t.Run(test.name, func(t *testing.T) {
			response := httptest.NewRecorder()
			test.handler.ServeHTTP(response, httptest.NewRequest(http.MethodPost, test.target, nil))
			if response.Code != http.StatusMethodNotAllowed {
				t.Fatalf("status = %d, want %d", response.Code, http.StatusMethodNotAllowed)
			}
			if contentType := response.Header().Get("Content-Type"); contentType != "application/problem+json" {
				t.Fatalf("content type = %q", contentType)
			}
		})
	}
}

func TestNewValidatesOptions(t *testing.T) {
	valid := auth.OIDCConfig{Issuer: "https://idp.example", ClientID: "pulse", RedirectURL: "https://pulse.example/auth/callback"}
	for _, test := range []struct {
		name    string
		options Options
		wantErr bool
	}{
		{"complete", Options{OIDC: valid, Sessions: &recordingSessions{}}, false},
		{"missing issuer", Options{OIDC: auth.OIDCConfig{ClientID: "pulse", RedirectURL: valid.RedirectURL}, Sessions: &recordingSessions{}}, true},
		{"missing client id", Options{OIDC: auth.OIDCConfig{Issuer: valid.Issuer, RedirectURL: valid.RedirectURL}, Sessions: &recordingSessions{}}, true},
		{"missing redirect url", Options{OIDC: auth.OIDCConfig{Issuer: valid.Issuer, ClientID: "pulse"}, Sessions: &recordingSessions{}}, true},
		{"missing sessions", Options{OIDC: valid}, true},
	} {
		t.Run(test.name, func(t *testing.T) {
			handler, err := New(test.options)
			if (err != nil) != test.wantErr {
				t.Fatalf("err = %v, wantErr = %v", err, test.wantErr)
			}
			if err != nil {
				if handler != nil {
					t.Fatal("handler returned with an error")
				}
				return
			}
			if handler.options.ErrorPath != defaultErrorPath || handler.options.DefaultRedirect != defaultRedirect || handler.options.GroupsClaim != "groups" {
				t.Fatalf("defaults not applied: %#v", handler.options)
			}
		})
	}
}

func TestNewNormalizesUnsafePaths(t *testing.T) {
	handler, err := New(Options{
		OIDC:            auth.OIDCConfig{Issuer: "https://idp.example", ClientID: "pulse", RedirectURL: "https://pulse.example/auth/callback"},
		Sessions:        &recordingSessions{},
		DefaultRedirect: "//evil.example",
		ErrorPath:       "/login/error?reason=spoofed",
	})
	if err != nil {
		t.Fatalf("New: %v", err)
	}
	if handler.options.DefaultRedirect != defaultRedirect || handler.options.ErrorPath != defaultErrorPath {
		t.Fatalf("unsafe paths were kept: %#v", handler.options)
	}
}