package server

import (
	"context"
	"encoding/base32"
	"encoding/json"
	"os"
	"strings"
	"sync"
	"testing"
	"time"

	"vdi-cloud/platform/internal/store"
)

func TestTOTPRFC6238VectorsAndReplay(t *testing.T) {
	secret := base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString([]byte("12345678901234567890"))
	for _, tc := range []struct {
		unix int64
		want string
	}{{59, "94287082"}, {1111111109, "07081804"}, {1111111111, "14050471"}, {1234567890, "89005924"}, {2000000000, "69279037"}, {20000000000, "65353130"}} {
		got, err := totpCode(secret, tc.unix/30, 8)
		if err != nil || got != tc.want {
			t.Fatal(tc.unix, got, err)
		}
	}
	now := time.Unix(1234567890, 0)
	for delta := int64(-2); delta <= 2; delta++ {
		code, _ := totpCode(secret, now.Unix()/30+delta, 6)
		counter, ok := matchTOTPCounter(secret, code, now, -1)
		if ok != (delta >= -1 && delta <= 1) {
			t.Fatal("wrong time window", delta)
		}
		if ok {
			if _, replay := matchTOTPCounter(secret, code, now, counter); replay {
				t.Fatal("replay accepted")
			}
		}
	}
}

func setupTestMFA(t *testing.T, f *fixture, token, password string) (string, []string) {
	t.Helper()
	code, out := f.request("POST", "/auth/mfa/setup", token, map[string]any{"password": password}, "")
	expect(t, 200, code, out)
	secret := out["secret"].(string)
	if !strings.HasPrefix(out["provisioningUri"].(string), "otpauth://totp/") {
		t.Fatal("invalid provisioning URI")
	}
	value, _ := totpCode(secret, time.Now().Unix()/30, 6)
	code, out = f.request("POST", "/auth/mfa/confirm", token, map[string]any{"code": value}, "")
	expect(t, 200, code, out)
	items := out["recoveryCodes"].([]any)
	if len(items) != 8 {
		t.Fatal("missing recovery codes")
	}
	codes := make([]string, len(items))
	for i, v := range items {
		codes[i] = v.(string)
	}
	return secret, codes
}

func TestMFALifecycleLoginRecoveryRotationAndDisable(t *testing.T) {
	f := newFixtureWorker(t, false)
	owner := f.login("admin")
	id := createTestAccount(t, f, owner, "mfa-user", "viewer")
	user := accountLogin(t, f, "mfa-user", testAccountPassword)
	old := accountLogin(t, f, "mfa-user", testAccountPassword)
	secret, codes := setupTestMFA(t, f, user, testAccountPassword)
	code, out := f.request("GET", "/auth/me", old, nil, "")
	expect(t, 401, code, out)
	code, out = f.request("GET", "/auth/mfa", user, nil, "")
	expect(t, 200, code, out)
	if out["enabled"] != true || out["recoveryRemaining"] != float64(8) {
		t.Fatal(out)
	}
	encoded, _ := json.Marshal(out)
	if strings.Contains(string(encoded), secret) {
		t.Fatal("status leaks secret")
	}
	var record mfaRecord
	if err := f.db.Get(&record, "SELECT * FROM admin_mfa WHERE admin_id=$1", id); err != nil {
		t.Fatal(err)
	}
	if strings.Contains(record.Secret, secret) || record.Pending != "" {
		t.Fatal("secret not encrypted or pending not cleared")
	}
	if _, err := f.server.cryptMFA("wrong-admin", record.Secret, false); err == nil {
		t.Fatal("secret not bound to account")
	}
	code, out = f.request("POST", "/auth/login", "", map[string]any{"username": "mfa-user", "password": testAccountPassword}, "")
	expect(t, 401, code, out)
	if out["code"] != "mfa_required" || out["accessToken"] != nil {
		t.Fatal(out)
	}
	// Failed password must not reveal the MFA requirement or consume a recovery code.
	code, out = f.request("POST", "/auth/login", "", map[string]any{"username": "mfa-user", "password": "incorrect", "code": codes[0]}, "")
	expect(t, 401, code, out)
	if out["code"] != "invalid_credentials" {
		t.Fatal(out)
	}
	code, out = f.request("POST", "/auth/login", "", map[string]any{"username": "mfa-user", "password": testAccountPassword, "code": strings.ToLower(codes[0])}, "")
	expect(t, 200, code, out)
	newToken := out["accessToken"].(string)
	code, out = f.request("POST", "/auth/login", "", map[string]any{"username": "mfa-user", "password": testAccountPassword, "code": codes[0]}, "")
	expect(t, 401, code, out)
	// Rotate consumes one code, invalidates all old recovery codes and other sessions.
	code, out = f.request("POST", "/auth/mfa/recovery", newToken, map[string]any{"password": testAccountPassword, "code": codes[1]}, "")
	expect(t, 200, code, out)
	newCodes := out["recoveryCodes"].([]any)
	code, out = f.request("GET", "/auth/me", user, nil, "")
	expect(t, 401, code, out)
	code, out = f.request("DELETE", "/auth/mfa", newToken, map[string]any{"password": testAccountPassword, "code": codes[2]}, "")
	expect(t, 400, code, out)
	code, out = f.request("DELETE", "/auth/mfa", newToken, map[string]any{"password": testAccountPassword, "code": newCodes[0]}, "")
	expect(t, 204, code, out)
	code, out = f.request("GET", "/auth/mfa", newToken, nil, "")
	expect(t, 200, code, out)
	if out["enabled"] != false || out["recoveryRemaining"] != float64(0) {
		t.Fatal(out)
	}
	_ = accountLogin(t, f, "mfa-user", testAccountPassword)
	var actions []string
	f.db.Select(&actions, "SELECT action FROM audit_logs WHERE actor='mfa-user' AND action IN ('enable_mfa','disable_mfa','rotate_mfa_recovery') ORDER BY created_at,id")
	if len(actions) != 3 {
		t.Fatal("MFA mutations not audited", actions)
	}
}

func TestMFASetupExpirySessionBindingValidationAndLimits(t *testing.T) {
	f := newFixtureWorker(t, false)
	user := f.login("viewer")
	other := f.login("viewer")
	for _, path := range []string{"/auth/mfa", "/auth/mfa/setup", "/auth/mfa/confirm", "/auth/mfa/recovery"} {
		method := "POST"
		if path == "/auth/mfa" {
			method = "GET"
		}
		code, out := f.request(method, path, "", map[string]any{}, "")
		expect(t, 401, code, out)
	}
	code, out := f.request("POST", "/auth/mfa/setup", user, map[string]any{"password": "incorrect"}, "")
	expect(t, 400, code, out)
	code, out = f.request("POST", "/auth/mfa/setup", user, map[string]any{"password": "Demo@2026", "adminId": "other"}, "")
	expect(t, 400, code, out)
	code, out = f.request("POST", "/auth/mfa/setup", user, map[string]any{"password": "Demo@2026"}, "")
	expect(t, 200, code, out)
	secret := out["secret"].(string)
	value, _ := totpCode(secret, time.Now().Unix()/30, 6)
	code, out = f.request("POST", "/auth/mfa/confirm", other, map[string]any{"code": value}, "")
	expect(t, 409, code, out)
	f.db.MustExec("UPDATE admin_mfa SET pending_expires=now()-interval '1 second'")
	code, out = f.request("POST", "/auth/mfa/confirm", user, map[string]any{"code": value}, "")
	expect(t, 409, code, out)
	code, out = f.request("GET", "/auth/mfa", user, nil, "")
	expect(t, 200, code, out)
	if out["enabled"] != false {
		t.Fatal(out)
	}
	f.db.MustExec("TRUNCATE login_limits")
	for i := 0; i < 11; i++ {
		code, out = f.request("POST", "/auth/mfa/setup", user, map[string]any{"password": "incorrect"}, "")
		want := 400
		if i == 10 {
			want = 429
		}
		expect(t, want, code, out)
	}
	// MFA attempt budget does not block ordinary login for an unbound account.
	_ = f.login("viewer")
}

func TestMFAConcurrentCodeConsumptionAndTOTPLogin(t *testing.T) {
	f := newFixtureWorker(t, false)
	user := f.login("operator")
	secret, codes := setupTestMFA(t, f, user, "Demo@2026")
	var id string
	f.db.Get(&id, "SELECT id FROM administrators WHERE username='operator'")
	// Clear only the test row's counter to exercise a fresh current-time code without a 30-second sleep.
	f.db.MustExec("UPDATE admin_mfa SET last_counter=-1 WHERE admin_id=$1", id)
	value, _ := totpCode(secret, time.Now().Unix()/30, 6)
	var wg sync.WaitGroup
	results := make(chan int, 2)
	for i := 0; i < 2; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			code, _ := f.request("POST", "/auth/login", "", map[string]any{"username": "operator", "password": "Demo@2026", "code": value}, "")
			results <- code
		}()
	}
	wg.Wait()
	close(results)
	sum := 0
	for code := range results {
		sum += code
	}
	if sum != 601 {
		t.Fatal("TOTP replay/concurrency", sum)
	}
	results = make(chan int, 2)
	for i := 0; i < 2; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			code, _ := f.request("POST", "/auth/login", "", map[string]any{"username": "operator", "password": "Demo@2026", "code": codes[0]}, "")
			results <- code
		}()
	}
	wg.Wait()
	close(results)
	sum = 0
	for code := range results {
		sum += code
	}
	if sum != 601 {
		t.Fatal("recovery code replay/concurrency", sum)
	}
}

func TestMFASuperAdminRestartAndPasswordChangePreserveSecondFactor(t *testing.T) {
	f := newFixtureWorker(t, false)
	owner := f.login("admin")
	_, codes := setupTestMFA(t, f, owner, "Demo@2026")
	db, err := store.Open(context.Background(), os.Getenv("VDI_TEST_DATABASE_URL"))
	if err != nil {
		t.Fatal(err)
	}
	db.Close()
	code, out := f.request("GET", "/auth/mfa", owner, nil, "")
	expect(t, 200, code, out)
	if out["enabled"] != true || out["recoveryRemaining"] != float64(8) {
		t.Fatal("restart changed MFA", out)
	}
	code, out = f.request("POST", "/auth/password", owner, map[string]any{"currentPassword": "Demo@2026", "newPassword": "Changed-mfa-2026"}, "")
	expect(t, 204, code, out)
	code, out = f.request("POST", "/auth/login", "", map[string]any{"username": "admin", "password": "Changed-mfa-2026"}, "")
	expect(t, 401, code, out)
	if out["code"] != "mfa_required" {
		t.Fatal("password change bypassed MFA", out)
	}
	code, out = f.request("POST", "/auth/login", "", map[string]any{"username": "admin", "password": "Changed-mfa-2026", "code": codes[0]}, "")
	expect(t, 200, code, out)
}