package server_test

import (
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base32"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"html"
	"net/http"
	"net/http/httptest"
	"net/url"
	"path/filepath"
	"regexp"
	"strings"
	"testing"
	"time"

	"misaki/db"
	"misaki/server"
)

func code(secret string, offset int64) string {
	key, err := base32.StdEncoding.WithPadding(base32.NoPadding).DecodeString(strings.ReplaceAll(secret, " ", ""))
	if err != nil {
		panic(err)
	}
	var msg [8]byte
	binary.BigEndian.PutUint64(msg[:], uint64(time.Now().Unix()/30+offset))
	m := hmac.New(sha1.New, key)
	m.Write(msg[:])
	sum := m.Sum(nil)
	off := sum[len(sum)-1] & 0xf
	return fmt.Sprintf("%06d", (binary.BigEndian.Uint32(sum[off:off+4])&0x7fffffff)%1000000)
}

func cookieNamed(rr *httptest.ResponseRecorder, name string) *http.Cookie {
	for _, c := range rr.Result().Cookies() {
		if c.Name == name && c.Value != "" {
			return c
		}
	}
	return nil
}

func newTOTPServer(t *testing.T, smtp server.SMTPConfig) (http.Handler, *db.DB) {
	t.Helper()
	d := newTestDB(t)
	srv := server.New(d, server.ServiceURLs{}, nil, smtp, server.Config{
		TOTPKeyFile: filepath.Join(t.TempDir(), "totp.key"),
		PublicURL:   "https://misaki.test",
	})
	return srv, d
}

var (
	secretRe   = regexp.MustCompile(`<code[^>]*>([A-Z2-7 ]+)</code>`)
	blobRe     = regexp.MustCompile(`name="blob" value="([^"]+)"`)
	recoveryRe = regexp.MustCompile(`[a-z2-7]{4}-[a-z2-7]{4}-[a-z2-7]{4}-[a-z2-7]{4}`)
)

// enable2FA turns on TOTP from the dashboard and returns the secret and recovery codes.
func enable2FA(t *testing.T, srv http.Handler, c *http.Cookie) (string, []string) {
	t.Helper()
	rr := doPost(srv, "/dashboard/2fa/setup", nil, c)
	assertStatus(t, rr, http.StatusOK)
	body := rr.Body.String()
	sm, bm := secretRe.FindStringSubmatch(body), blobRe.FindStringSubmatch(body)
	if sm == nil || bm == nil {
		t.Fatalf("setup page missing secret/blob:\n%s", body)
	}
	if !strings.Contains(body, "otpauth://totp/misaki:") {
		t.Errorf("setup page has no otpauth link")
	}
	if !strings.Contains(body, `<svg class="qr"`) {
		t.Errorf("setup page has no QR code")
	}
	secret, blob := sm[1], html.UnescapeString(bm[1])

	rr = doPost(srv, "/dashboard/2fa/enable", url.Values{"blob": {blob}, "code": {"000000"}}, c)
	if !strings.Contains(rr.Body.String(), "wrong code") {
		t.Errorf("wrong setup code accepted")
	}
	rr = doPost(srv, "/dashboard/2fa/enable", url.Values{"blob": {blob}, "code": {code(secret, 0)}}, c)
	assertStatus(t, rr, http.StatusOK)
	codes := recoveryRe.FindAllString(rr.Body.String(), -1)
	if len(codes) != 10 {
		t.Fatalf("want 10 recovery codes, got %d", len(codes))
	}
	return secret, codes
}

func login2FA(t *testing.T, srv http.Handler, user, pass, otp string) *httptest.ResponseRecorder {
	t.Helper()
	rr := doPost(srv, "/login", url.Values{"username": {user}, "password": {pass}}, nil)
	assertRedirectTo(t, rr, "/login/code")
	if cookieNamed(rr, "session") != nil {
		t.Fatalf("session issued before the code")
	}
	p := cookieNamed(rr, "login_code")
	if p == nil {
		t.Fatalf("no pending cookie")
	}
	return doPost(srv, "/login/code", url.Values{"code": {otp}}, p)
}

func TestTOTPLogin(t *testing.T) {
	srv, d := newTOTPServer(t, server.SMTPConfig{})
	makeApprovedUser(t, d, "alice", "password123")
	c := doLogin(t, srv, "alice", "password123")
	other := doLogin(t, srv, "alice", "password123")
	secret, codes := enable2FA(t, srv, c)

	assertStatus(t, doGet(srv, "/dashboard", c), http.StatusOK)
	assertRedirectTo(t, doGet(srv, "/dashboard", other), "/login")

	assertRedirectTo(t, login2FA(t, srv, "alice", "password123", "123456"), "/login/code?err=")

	// the setup step was used, so the next step's code is the only fresh TOTP
	rr := login2FA(t, srv, "alice", "password123", code(secret, 1))
	assertRedirectTo(t, rr, "/dashboard")
	sc := cookieNamed(rr, "session")
	if sc == nil {
		t.Fatal("no session after code")
	}
	if s, err := d.GetSession(sc.Value); err != nil || s.AMR != "pwd otp" {
		t.Errorf("amr: %+v %v", s, err)
	}

	for _, replay := range []string{code(secret, 1), code(secret, 0)} {
		assertRedirectTo(t, login2FA(t, srv, "alice", "password123", replay), "/login/code?err=")
	}

	assertRedirectTo(t, login2FA(t, srv, "alice", "password123", strings.ToUpper(codes[0])), "/dashboard")
	assertRedirectTo(t, login2FA(t, srv, "alice", "password123", codes[0]), "/login")
}

func TestTOTPPendingCookieIsNotASession(t *testing.T) {
	srv, d := newTOTPServer(t, server.SMTPConfig{})
	makeApprovedUser(t, d, "bob", "password123")
	enable2FA(t, srv, doLogin(t, srv, "bob", "password123"))
	rr := doPost(srv, "/login", url.Values{"username": {"bob"}, "password": {"password123"}}, nil)
	p := cookieNamed(rr, "login_code")
	p.Name = "session"
	assertRedirectTo(t, doGet(srv, "/dashboard", p), "/login")
	assertRedirectTo(t, doGet(srv, "/login/code", nil), "/login")
}

func TestTOTPCodeLimit(t *testing.T) {
	srv, d := newTOTPServer(t, server.SMTPConfig{})
	makeApprovedUser(t, d, "carol", "password123")
	secret, _ := enable2FA(t, srv, doLogin(t, srv, "carol", "password123"))
	for range 5 {
		login2FA(t, srv, "carol", "password123", "999999")
	}
	rr := login2FA(t, srv, "carol", "password123", code(secret, 1))
	assertRedirectTo(t, rr, "/login?err=")
	if cookieNamed(rr, "session") != nil {
		t.Error("correct code accepted while locked")
	}
}

func TestTOTPDisable(t *testing.T) {
	srv, d := newTOTPServer(t, server.SMTPConfig{})
	id := makeApprovedUser(t, d, "dave", "password123")
	c := doLogin(t, srv, "dave", "password123")
	_, codes := enable2FA(t, srv, c)
	assertRedirectTo(t, doPost(srv, "/dashboard/2fa/disable", url.Values{"code": {"nope"}}, c), "/dashboard?err=")
	assertRedirectTo(t, doPost(srv, "/dashboard/2fa/disable", url.Values{"code": {codes[3]}}, c), "/dashboard?msg=")
	if sealed, _, _ := d.GetTOTP(id); sealed != "" {
		t.Error("secret still stored")
	}
	if n, _ := d.CountRecoveryCodes(id); n != 0 {
		t.Errorf("%d recovery codes left", n)
	}
	doLogin(t, srv, "dave", "password123")
}

func TestTOTPSecretSealed(t *testing.T) {
	srv, d := newTOTPServer(t, server.SMTPConfig{})
	id := makeApprovedUser(t, d, "erin", "password123")
	secret, _ := enable2FA(t, srv, doLogin(t, srv, "erin", "password123"))
	sealed, _, _ := d.GetTOTP(id)
	if sealed == "" || strings.Contains(sealed, strings.ReplaceAll(secret, " ", "")) {
		t.Errorf("secret not sealed: %q", sealed)
	}
}

func TestPasswordChange(t *testing.T) {
	srv, d := newTOTPServer(t, server.SMTPConfig{})
	makeApprovedUser(t, d, "fay", "password123")
	c := doLogin(t, srv, "fay", "password123")
	other := doLogin(t, srv, "fay", "password123")
	form := url.Values{"current": {"wrong"}, "password": {"newpass123"}, "password_confirm": {"newpass123"}}
	assertRedirectTo(t, doPost(srv, "/dashboard/password", form, c), "/dashboard?err=")
	form.Set("current", "password123")
	assertRedirectTo(t, doPost(srv, "/dashboard/password", form, c), "/dashboard?msg=")
	assertStatus(t, doGet(srv, "/dashboard", c), http.StatusOK)
	assertRedirectTo(t, doGet(srv, "/dashboard", other), "/login")
	doLogin(t, srv, "fay", "newpass123")
}

func TestResetToken(t *testing.T) {
	srv, d := newTOTPServer(t, server.SMTPConfig{})
	id := makeApprovedUser(t, d, "gus", "password123")
	old := doLogin(t, srv, "gus", "password123")

	dead, _ := d.CreateResetToken(id, -time.Minute)
	if !strings.Contains(doGet(srv, "/reset?t="+dead, nil).Body.String(), "expired") {
		t.Error("expired token accepted")
	}

	tok, _ := d.CreateResetToken(id, 30*time.Minute)
	rr := doGet(srv, "/reset?t="+tok, nil)
	if !strings.Contains(rr.Body.String(), `action="/reset/confirm"`) {
		t.Fatal("no reset form for a valid token")
	}
	if rr.Header().Get("Referrer-Policy") != "no-referrer" || rr.Header().Get("Cache-Control") != "no-store" {
		t.Error("reset page leaks via referrer or cache")
	}

	form := url.Values{"t": {tok}, "password": {"brandnew1"}, "password_confirm": {"different"}}
	assertRedirectTo(t, doPost(srv, "/reset/confirm", form, nil), "/reset?t=")
	form.Set("password_confirm", "brandnew1")
	assertRedirectTo(t, doPost(srv, "/reset/confirm", form, nil), "/login?msg=")

	assertRedirectTo(t, doGet(srv, "/dashboard", old), "/login")
	doLogin(t, srv, "gus", "brandnew1")
	if rr := doPost(srv, "/reset/confirm", form, nil); rr.Header().Get("Location") != "/reset?t=expired" {
		t.Errorf("token reused: %s", rr.Header().Get("Location"))
	}
}

func TestResetStillNeeds2FA(t *testing.T) {
	srv, d := newTOTPServer(t, server.SMTPConfig{})
	id := makeApprovedUser(t, d, "hana", "password123")
	secret, _ := enable2FA(t, srv, doLogin(t, srv, "hana", "password123"))
	tok, _ := d.CreateResetToken(id, 30*time.Minute)
	form := url.Values{"t": {tok}, "password": {"brandnew1"}, "password_confirm": {"brandnew1"}}
	assertRedirectTo(t, doPost(srv, "/reset/confirm", form, nil), "/reset?t=")
	if _, err := d.GetResetToken(tok); err != nil {
		t.Fatal("failed code burnt the token")
	}
	form.Set("code", code(secret, 1))
	assertRedirectTo(t, doPost(srv, "/reset/confirm", form, nil), "/login?msg=")
	login2FA(t, srv, "hana", "brandnew1", "000000")
}

func TestResetNoEnumeration(t *testing.T) {
	srv, d := newTOTPServer(t, server.SMTPConfig{Host: "127.0.0.1:1", From: "misaki@misaki.test"})
	id := makeApprovedUser(t, d, "ivy", "password123")
	d.SetUserEmail(id, "ivy@misaki.test")
	makeApprovedUser(t, d, "jun", "password123")
	var bodies []string
	for _, u := range []string{"ivy", "jun", "nobody"} {
		rr := doPost(srv, "/reset", url.Values{"username": {u}}, nil)
		assertStatus(t, rr, http.StatusOK)
		bodies = append(bodies, rr.Body.String())
	}
	if bodies[0] != bodies[1] || bodies[1] != bodies[2] {
		t.Error("reset response differs by account")
	}
	time.Sleep(100 * time.Millisecond) // let the background mail attempts finish
}

func TestOIDCAMR(t *testing.T) {
	srv, d := newOIDCServer(t)
	c := grantedUser(t, srv, d, "kai")
	amr := func(c *http.Cookie) any {
		rr := exchange(srv, url.Values{"code": {getCode(t, srv, c, nil)}}, true)
		var tr tokenResp
		json.NewDecoder(rr.Body).Decode(&tr)
		return verifyIDToken(t, srv, tr.IDToken)["amr"]
	}
	if got := fmt.Sprint(amr(c)); got != "[pwd]" {
		t.Errorf("password login amr: %s", got)
	}
	secret, _ := enable2FA(t, srv, c)
	rr := doPost(srv, "/login", url.Values{"username": {"kai"}, "password": {"hunter2hunter2"}}, nil)
	rr = doPost(srv, "/login/code", url.Values{"code": {code(secret, 1)}}, cookieNamed(rr, "login_code"))
	if got := fmt.Sprint(amr(cookieNamed(rr, "session"))); got != "[pwd otp]" {
		t.Errorf("2FA login amr: %s", got)
	}
}

func TestRename(t *testing.T) {
	srv, d := newTestServer(t)
	uid := makeApprovedUser(t, d, "alice", "password1")
	makeApprovedUser(t, d, "bob", "password1")
	c := doLogin(t, srv, "alice", "password1")
	rename := func(name, pass string) string {
		rr := doPost(srv, "/dashboard/username", url.Values{"username": {name}, "current": {pass}}, c)
		return rr.Header().Get("Location")
	}

	for name, want := range map[string]string{
		"carol": "password+is+wrong",
		"BOB":   "taken",
		"a b":   "may+only",
		"admin": "not+available",
	} {
		pass := "password1"
		if name == "carol" {
			pass = "nope"
		}
		if loc := rename(name, pass); !strings.Contains(loc, want) {
			t.Errorf("rename %q: got %s, want %s", name, loc, want)
		}
	}

	if loc := rename("carol", "password1"); !strings.Contains(loc, "msg=") {
		t.Fatalf("rename failed: %s", loc)
	}
	if u, err := d.GetUserByID(uid); err != nil || u.Username != "carol" {
		t.Fatalf("username not changed: %v %v", u, err)
	}
	assertStatus(t, doGet(srv, "/dashboard", c), http.StatusOK)
	if taken, _ := d.NameTaken("ALICE", 0); !taken {
		t.Error("old name not reserved")
	}
	if taken, _ := d.NameTaken("alice", uid); taken {
		t.Error("old name should stay free for its owner")
	}
	if loc := rename("alice", "password1"); !strings.Contains(loc, "again+in+30+days") {
		t.Errorf("cooldown not enforced: %s", loc)
	}
	rr := doPost(srv, "/apply", url.Values{"username": {"Alice"}, "password": {"password1"}, "password_confirm": {"password1"},
		"email": {"x@example.com"}, "bio": {"hi"}, "agree": {"on"}}, nil)
	if !strings.Contains(rr.Header().Get("Location"), "taken") {
		t.Errorf("apply could claim a retired name: %s", rr.Header().Get("Location"))
	}
}

func TestRenameNeeds2FA(t *testing.T) {
	srv, d := newTOTPServer(t, server.SMTPConfig{})
	uid := makeApprovedUser(t, d, "alice", "password1")
	c := doLogin(t, srv, "alice", "password1")
	secret, _ := enable2FA(t, srv, c)
	form := url.Values{"username": {"carol"}, "current": {"password1"}}
	if loc := doPost(srv, "/dashboard/username", form, c).Header().Get("Location"); !strings.Contains(loc, "wrong+code") {
		t.Fatalf("rename without code: %s", loc)
	}
	form.Set("code", code(secret, 1))
	if loc := doPost(srv, "/dashboard/username", form, c).Header().Get("Location"); !strings.Contains(loc, "msg=") {
		t.Fatalf("rename with code: %s", loc)
	}
	if u, _ := d.GetUserByID(uid); u.Username != "carol" {
		t.Errorf("not renamed: %s", u.Username)
	}
}