package token

import (
	"encoding/json"
	"testing"
	"time"

	"github.com/tinode/chat/server/auth"
	"github.com/tinode/chat/server/store"
	"github.com/tinode/chat/server/store/types"
)

type testSecurityStore struct {
	version int
	err     error
}

func (s testSecurityStore) GetAuthVersion(types.TenantID, types.Uid) (int, error) {
	return s.version, s.err
}

func (s testSecurityStore) BumpAuthVersion(types.TenantID, types.Uid, time.Time) (int, error) {
	return s.version + 1, s.err
}

func (s testSecurityStore) GetRealName(types.TenantID, types.Uid) (*types.RealNameVerification, error) {
	return nil, types.ErrNotFound
}

func (s testSecurityStore) CreateRealNameSubmission(types.TenantID, *types.RealNameVerification) error {
	return s.err
}

func (s testSecurityStore) CreateActionToken(types.TenantID, *types.ActionToken) error {
	return s.err
}

func (s testSecurityStore) ConsumeActionToken(types.TenantID, types.Uid, string, string, string, time.Time) (*types.ActionToken, error) {
	return nil, s.err
}

func (s testSecurityStore) GetIndependentPassword(types.TenantID, types.Uid, string) (*types.IndependentPassword, error) {
	return nil, types.ErrNotFound
}

func (s testSecurityStore) UpsertIndependentPassword(types.TenantID, *types.IndependentPassword) error {
	return s.err
}

func (s testSecurityStore) UpdateIndependentPasswordAttempts(types.TenantID, types.Uid, string, int, *time.Time, time.Time) error {
	return s.err
}

func (s testSecurityStore) GetYouthMode(types.TenantID, types.Uid) (*types.YouthModeSettings, error) {
	return nil, types.ErrNotFound
}

func (s testSecurityStore) SetYouthMode(types.TenantID, types.Uid, bool, int, time.Time) (*types.YouthModeSettings, error) {
	return nil, s.err
}

func (s testSecurityStore) CreateSystemNotification(types.TenantID, *types.SystemNotification) error {
	return s.err
}

func (s testSecurityStore) ListSystemNotifications(types.TenantID, types.Uid, string, int) ([]types.SystemNotification, string, error) {
	return nil, "", s.err
}

func (s testSecurityStore) MarkSystemNotificationRead(types.TenantID, types.Uid, string, time.Time) error {
	return s.err
}

func newTestAuthenticator(t *testing.T) *authenticator {
	t.Helper()
	conf, err := json.Marshal(map[string]any{
		"key":        []byte("0123456789abcdef0123456789abcdef"),
		"serial_num": 3,
		"expire_in":  3600,
	})
	if err != nil {
		t.Fatal(err)
	}
	ta := &authenticator{}
	if err = ta.Init(conf, "token-test"); err != nil {
		t.Fatal(err)
	}
	return ta
}

func TestTokenV3TenantRoundTrip(t *testing.T) {
	oldSecurity := store.Security
	store.Security = testSecurityStore{version: 11}
	defer func() { store.Security = oldSecurity }()

	ta := newTestAuthenticator(t)
	rec := &auth.Rec{
		TenantID:    7,
		Uid:         42,
		AuthLevel:   auth.LevelAuth,
		Lifetime:    auth.Duration(time.Hour),
		Features:    auth.FeatureValidated,
		AuthVersion: 11,
	}
	token, _, err := ta.GenSecret(rec)
	if err != nil {
		t.Fatal(err)
	}
	if got, want := len(token), 63; got != want {
		t.Fatalf("token length: got %d, want %d", got, want)
	}

	got, challenge, err := ta.Authenticate(auth.AuthContext{TenantID: 7}, token)
	if err != nil {
		t.Fatal(err)
	}
	if challenge != nil {
		t.Fatal("unexpected challenge")
	}
	if got.TenantID != rec.TenantID || got.Uid != rec.Uid || got.AuthLevel != rec.AuthLevel ||
		got.Features != rec.Features || got.AuthVersion != rec.AuthVersion {
		t.Fatalf("round trip mismatch: got %+v, want %+v", got, rec)
	}

	// HTTP authentication may recover the tenant from the signed token.
	got, _, err = ta.Authenticate(auth.AuthContext{}, token)
	if err != nil || got.TenantID != rec.TenantID {
		t.Fatalf("tenant recovery failed: rec=%+v err=%v", got, err)
	}
}

func TestTokenV3RejectsWrongOrTamperedTenant(t *testing.T) {
	oldSecurity := store.Security
	store.Security = testSecurityStore{version: 1}
	defer func() { store.Security = oldSecurity }()

	ta := newTestAuthenticator(t)
	token, _, err := ta.GenSecret(&auth.Rec{
		TenantID: 7, Uid: 42, AuthLevel: auth.LevelAuth, Lifetime: auth.Duration(time.Hour), AuthVersion: 1,
	})
	if err != nil {
		t.Fatal(err)
	}

	if _, _, err = ta.Authenticate(auth.AuthContext{TenantID: 8}, token); err != types.ErrFailed {
		t.Fatalf("wrong tenant: got %v, want %v", err, types.ErrFailed)
	}

	tampered := append([]byte(nil), token...)
	tampered[1] ^= 1 // First byte of tenant_id in the little-endian payload.
	if _, _, err = ta.Authenticate(auth.AuthContext{}, tampered); err != types.ErrFailed {
		t.Fatalf("tampered tenant: got %v, want %v", err, types.ErrFailed)
	}
}

func TestTokenV3RejectsLegacyMissingIdentityAndStaleVersion(t *testing.T) {
	oldSecurity := store.Security
	store.Security = testSecurityStore{version: 2}
	defer func() { store.Security = oldSecurity }()

	ta := newTestAuthenticator(t)
	if _, _, err := ta.Authenticate(auth.AuthContext{}, make([]byte, 50)); err != types.ErrMalformed {
		t.Fatalf("legacy token: got %v, want %v", err, types.ErrMalformed)
	}
	if _, _, err := ta.GenSecret(&auth.Rec{Uid: 42, AuthLevel: auth.LevelAuth}); err != types.ErrMalformed {
		t.Fatalf("missing tenant: got %v, want %v", err, types.ErrMalformed)
	}
	if _, _, err := ta.GenSecret(&auth.Rec{TenantID: 7, AuthLevel: auth.LevelAuth}); err != types.ErrMalformed {
		t.Fatalf("missing uid: got %v, want %v", err, types.ErrMalformed)
	}

	token, _, err := ta.GenSecret(&auth.Rec{
		TenantID: 7, Uid: 42, AuthLevel: auth.LevelAuth, Lifetime: auth.Duration(time.Hour), AuthVersion: 1,
	})
	if err != nil {
		t.Fatal(err)
	}
	if _, _, err = ta.Authenticate(auth.AuthContext{TenantID: 7}, token); err != types.ErrFailed {
		t.Fatalf("stale auth version: got %v, want %v", err, types.ErrFailed)
	}
}

func TestTokenV3NoLoginDoesNotRequireAuthVersion(t *testing.T) {
	oldSecurity := store.Security
	store.Security = testSecurityStore{version: 99}
	defer func() { store.Security = oldSecurity }()

	ta := newTestAuthenticator(t)
	token, _, err := ta.GenSecret(&auth.Rec{
		TenantID: 7, Uid: 42, AuthLevel: auth.LevelAuth, Lifetime: auth.Duration(time.Hour), Features: auth.FeatureNoLogin,
	})
	if err != nil {
		t.Fatal(err)
	}
	got, _, err := ta.Authenticate(auth.AuthContext{TenantID: 7}, token)
	if err != nil {
		t.Fatal(err)
	}
	if got.AuthVersion != 0 || got.Features&auth.FeatureNoLogin == 0 {
		t.Fatalf("unexpected no-login token auth: %+v", got)
	}
}