package webui

import (
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	"github.com/pquerna/otp/totp"
	"mailgoserver/internal/db"
)

func TestLoginSuccessWithoutMFA(t *testing.T) {
	app := newTestApp(t)
	mux := app.Mux()

	hash, err := db.HashPassword("correct-horse-battery-staple1!")
	if err != nil {
		t.Fatal(err)
	}
	if _, err := app.DB.CreateAdminUser("alice", hash, false); err != nil {
		t.Fatal(err)
	}

	form := url.Values{"username": {"alice"}, "password": {"correct-horse-battery-staple1!"}}
	req := httptest.NewRequest(http.MethodPost, Prefix+"/login", strings.NewReader(form.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)

	if rec.Code != http.StatusFound {
		t.Fatalf("expected redirect after login, got %d: %s", rec.Code, rec.Body.String())
	}
	var sessionCookie *http.Cookie
	for _, c := range rec.Result().Cookies() {
		if c.Name == sessionCookieName {
			sessionCookie = c
		}
	}
	if sessionCookie == nil || sessionCookie.Value == "" {
		t.Fatal("expected a session cookie to be set on successful login")
	}

	// That cookie must actually grant access to a protected page.
	req2 := httptest.NewRequest(http.MethodGet, Prefix+"/", nil)
	req2.AddCookie(sessionCookie)
	rec2 := httptest.NewRecorder()
	mux.ServeHTTP(rec2, req2)
	if rec2.Code != http.StatusOK {
		t.Fatalf("expected dashboard to load with the new session, got %d", rec2.Code)
	}
}

func TestLoginWrongPasswordRejected(t *testing.T) {
	app := newTestApp(t)
	mux := app.Mux()

	hash, _ := db.HashPassword("the-real-password-123!")
	if _, err := app.DB.CreateAdminUser("bob", hash, false); err != nil {
		t.Fatal(err)
	}

	form := url.Values{"username": {"bob"}, "password": {"wrong-password"}}
	req := httptest.NewRequest(http.MethodPost, Prefix+"/login", strings.NewReader(form.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("expected the login form re-rendered with an error, got status %d", rec.Code)
	}
	for _, c := range rec.Result().Cookies() {
		if c.Name == sessionCookieName && c.Value != "" {
			t.Fatal("must not set a session cookie on failed login")
		}
	}
}

func TestLoginRequiresTOTPWhenEnabled(t *testing.T) {
	app := newTestApp(t)
	mux := app.Mux()

	hash, _ := db.HashPassword("carol-password-123!")
	userID, err := app.DB.CreateAdminUser("carol", hash, false)
	if err != nil {
		t.Fatal(err)
	}
	secret := "JBSWY3DPEHPK3PXP" // fixed test secret, valid base32
	if err := app.DB.SetAdminTOTPSecret(userID, secret, true); err != nil {
		t.Fatal(err)
	}

	form := url.Values{"username": {"carol"}, "password": {"carol-password-123!"}}
	req := httptest.NewRequest(http.MethodPost, Prefix+"/login", strings.NewReader(form.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)

	if rec.Code != http.StatusFound || rec.Header().Get("Location") != Prefix+"/login/mfa?next=" {
		t.Fatalf("expected redirect to MFA step, got %d location %q", rec.Code, rec.Header().Get("Location"))
	}
	var pendingCookie *http.Cookie
	for _, c := range rec.Result().Cookies() {
		if c.Name == pendingMFACookieName {
			pendingCookie = c
		}
	}
	if pendingCookie == nil {
		t.Fatal("expected a pending-MFA cookie after correct password with TOTP enabled")
	}

	// Wrong code must not grant a session.
	code, err := totp.GenerateCode(secret, time.Now())
	if err != nil {
		t.Fatal(err)
	}
	wrongCode := "000000"
	if wrongCode == code {
		wrongCode = "111111"
	}
	badForm := url.Values{"code": {wrongCode}}
	badReq := httptest.NewRequest(http.MethodPost, Prefix+"/login/mfa", strings.NewReader(badForm.Encode()))
	badReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	badReq.AddCookie(pendingCookie)
	badRec := httptest.NewRecorder()
	mux.ServeHTTP(badRec, badReq)
	for _, c := range badRec.Result().Cookies() {
		if c.Name == sessionCookieName && c.Value != "" {
			t.Fatal("must not grant a session for a wrong TOTP code")
		}
	}

	// Correct code must grant a fully-verified session.
	goodForm := url.Values{"code": {code}}
	goodReq := httptest.NewRequest(http.MethodPost, Prefix+"/login/mfa", strings.NewReader(goodForm.Encode()))
	goodReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	goodReq.AddCookie(pendingCookie)
	goodRec := httptest.NewRecorder()
	mux.ServeHTTP(goodRec, goodReq)
	if goodRec.Code != http.StatusFound {
		t.Fatalf("expected redirect after correct TOTP code, got %d", goodRec.Code)
	}
	var sessionCookie *http.Cookie
	for _, c := range goodRec.Result().Cookies() {
		if c.Name == sessionCookieName {
			sessionCookie = c
		}
	}
	if sessionCookie == nil || sessionCookie.Value == "" {
		t.Fatal("expected a session cookie after correct TOTP code")
	}

	dashReq := httptest.NewRequest(http.MethodGet, Prefix+"/", nil)
	dashReq.AddCookie(sessionCookie)
	dashRec := httptest.NewRecorder()
	mux.ServeHTTP(dashRec, dashReq)
	if dashRec.Code != http.StatusOK {
		t.Fatalf("expected dashboard to load after full MFA login, got %d", dashRec.Code)
	}
}
