package wristband

import (
	"context"
	"crypto"
	"fmt"
	"net/http"
	"slices"
	"sync"
	"testing"
	"time"
)

// memoryOIDCBackend is the in-memory OIDCBackend/OIDCTx for wristband's
// own tests. One mutex guards the whole transaction closure; it never rolls
// back, which the engine tolerates because it only returns an error from a
// transaction for server failures.
type memoryOIDCBackend struct {
	mu         sync.Mutex
	clients    []*OIDCClientRecord
	requests   []*AuthRequestRecord
	codes      []*OIDCCodeRecord
	tokens     []*OIDCAccessTokenRecord
	refresh    []*OIDCRefreshTokenRecord
	consents   []*ConsentRecord
	subjects   map[string]string // "<user>|<sector>" -> sub
	identities map[uint]map[string]any
	nextID     uint
}

func newMemoryOIDCBackend() *memoryOIDCBackend {
	return &memoryOIDCBackend{subjects: map[string]string{}, identities: map[uint]map[string]any{}}
}

func (b *memoryOIDCBackend) WithinOIDCTx(ctx context.Context, fn func(OIDCTx) error) error {
	b.mu.Lock()
	defer b.mu.Unlock()
	return fn(&memoryOIDCTx{b: b})
}

func (b *memoryOIDCBackend) id() uint {
	b.nextID++
	return b.nextID
}

// addClient registers an OIDC client whose secret is secret.
func (b *memoryOIDCBackend) addClient(c OIDCClientRecord, secret string) *OIDCClientRecord {
	b.mu.Lock()
	defer b.mu.Unlock()
	if secret != "" {
		h := sha256Hex(secret)
		c.SecretHash = &h
	}
	c.ID = b.id()
	cp := c
	b.clients = append(b.clients, &cp)
	return &cp
}

func (b *memoryOIDCBackend) setIdentity(userID uint, claims map[string]any) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.identities[userID] = claims
}

func (b *memoryOIDCBackend) subject(userID uint, sector string) string {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.subjects[fmt.Sprintf("%d|%s", userID, sector)]
}

func (b *memoryOIDCBackend) accessTokens() []OIDCAccessTokenRecord {
	b.mu.Lock()
	defer b.mu.Unlock()
	out := make([]OIDCAccessTokenRecord, len(b.tokens))
	for i, t := range b.tokens {
		out[i] = *t
	}
	return out
}

type memoryOIDCTx struct{ b *memoryOIDCBackend }

var _ OIDCTx = (*memoryOIDCTx)(nil)

func (t *memoryOIDCTx) OIDCClientByClientID(ctx context.Context, clientID string) (*OIDCClientRecord, error) {
	for _, c := range t.b.clients {
		if c.ClientID == clientID {
			cp := *c
			return &cp, nil
		}
	}
	return nil, nil
}

func (t *memoryOIDCTx) CreateAuthRequest(ctx context.Context, rec *AuthRequestRecord) error {
	rec.ID = t.b.id()
	cp := *rec
	t.b.requests = append(t.b.requests, &cp)
	return nil
}

func (t *memoryOIDCTx) findRequest(requestID string) *AuthRequestRecord {
	for _, r := range t.b.requests {
		if r.RequestID == requestID {
			return r
		}
	}
	return nil
}

func (t *memoryOIDCTx) AuthRequestByRequestID(ctx context.Context, requestID string) (*AuthRequestRecord, error) {
	if r := t.findRequest(requestID); r != nil {
		cp := *r
		return &cp, nil
	}
	return nil, nil
}

func (t *memoryOIDCTx) AuthRequestByRequestIDForUpdate(ctx context.Context, requestID string) (*AuthRequestRecord, error) {
	return t.AuthRequestByRequestID(ctx, requestID)
}

func (t *memoryOIDCTx) ConsumeAuthRequest(ctx context.Context, id uint, now time.Time) (bool, error) {
	for _, r := range t.b.requests {
		if r.ID == id {
			if r.ConsumedAt != nil || !r.ExpiresAt.After(now) {
				return false, nil
			}
			n := now
			r.ConsumedAt = &n
			return true, nil
		}
	}
	return false, nil
}

func (t *memoryOIDCTx) SetConsentTokenHash(ctx context.Context, id uint, hash *string) error {
	for _, r := range t.b.requests {
		if r.ID == id {
			if hash == nil {
				r.ConsentTokenHash = nil
			} else {
				h := *hash
				r.ConsentTokenHash = &h
			}
		}
	}
	return nil
}

func (t *memoryOIDCTx) CreateCode(ctx context.Context, rec *OIDCCodeRecord) error {
	rec.ID = t.b.id()
	cp := *rec
	t.b.codes = append(t.b.codes, &cp)
	return nil
}

func (t *memoryOIDCTx) ConsumeCode(ctx context.Context, codeHash string, now time.Time) (*OIDCCodeRecord, error) {
	for _, c := range t.b.codes {
		if c.CodeHash == codeHash {
			if c.ConsumedAt != nil || !c.ExpiresAt.After(now) {
				return nil, nil
			}
			n := now
			c.ConsumedAt = &n
			cp := *c
			return &cp, nil
		}
	}
	return nil, nil
}

func (t *memoryOIDCTx) CodeByHash(ctx context.Context, codeHash string) (*OIDCCodeRecord, error) {
	for _, c := range t.b.codes {
		if c.CodeHash == codeHash {
			cp := *c
			return &cp, nil
		}
	}
	return nil, nil
}

func (t *memoryOIDCTx) CreateAccessToken(ctx context.Context, rec *OIDCAccessTokenRecord) error {
	rec.ID = t.b.id()
	cp := *rec
	t.b.tokens = append(t.b.tokens, &cp)
	return nil
}

func (t *memoryOIDCTx) AccessTokenByHash(ctx context.Context, tokenHash string) (*OIDCAccessTokenRecord, error) {
	for _, tok := range t.b.tokens {
		if tok.TokenHash == tokenHash {
			cp := *tok
			return &cp, nil
		}
	}
	return nil, nil
}

func (t *memoryOIDCTx) TouchAccessToken(ctx context.Context, id uint, now time.Time) error {
	for _, tok := range t.b.tokens {
		if tok.ID == id {
			n := now
			tok.LastUsedAt = &n
		}
	}
	return nil
}

func (t *memoryOIDCTx) revokeAccessToken(id uint, now time.Time) {
	for _, tok := range t.b.tokens {
		if tok.ID == id && tok.RevokedAt == nil {
			n := now
			tok.RevokedAt = &n
		}
	}
}

func (t *memoryOIDCTx) RevokeAccessTokensForCode(ctx context.Context, codeID uint, now time.Time) error {
	for _, tok := range t.b.tokens {
		if tok.AuthCodeID != nil && *tok.AuthCodeID == codeID {
			t.revokeAccessToken(tok.ID, now)
			for _, rt := range t.b.refresh {
				if rt.AccessTokenID != nil && *rt.AccessTokenID == tok.ID {
					if err := t.RevokeRefreshLineage(ctx, rt.ID, now); err != nil {
						return err
					}
				}
			}
		}
	}
	return nil
}

func (t *memoryOIDCTx) ConsentFor(ctx context.Context, userID uint, clientID string) (*ConsentRecord, error) {
	for _, c := range t.b.consents {
		if c.UserID == userID && c.ClientID == clientID {
			cp := *c
			cp.Scopes = slices.Clone(c.Scopes)
			return &cp, nil
		}
	}
	return nil, nil
}

func (t *memoryOIDCTx) SaveConsent(ctx context.Context, userID uint, clientID string, scopes []string, now time.Time) error {
	n := now
	for _, c := range t.b.consents {
		if c.UserID == userID && c.ClientID == clientID {
			for _, sc := range scopes {
				if !slices.Contains(c.Scopes, sc) {
					c.Scopes = append(c.Scopes, sc)
				}
			}
			c.LastUsedAt = &n
			return nil
		}
	}
	t.b.consents = append(t.b.consents, &ConsentRecord{
		ID: t.b.id(), UserID: userID, ClientID: clientID,
		Scopes: slices.Clone(scopes), CreatedAt: now, LastUsedAt: &n,
	})
	return nil
}

func (t *memoryOIDCTx) TouchConsent(ctx context.Context, id uint, now time.Time) error {
	for _, c := range t.b.consents {
		if c.ID == id {
			n := now
			c.LastUsedAt = &n
		}
	}
	return nil
}

func (t *memoryOIDCTx) SubjectFor(ctx context.Context, userID uint, sector, candidate string) (string, error) {
	key := fmt.Sprintf("%d|%s", userID, sector)
	if sub, ok := t.b.subjects[key]; ok {
		return sub, nil
	}
	t.b.subjects[key] = candidate
	return candidate, nil
}

func (t *memoryOIDCTx) IdentityClaims(ctx context.Context, userID uint, scopes []string) (map[string]any, bool, error) {
	all, ok := t.b.identities[userID]
	if !ok {
		return nil, false, nil
	}
	out := map[string]any{}
	if slices.Contains(scopes, "profile") {
		for _, k := range []string{"name", "picture", "groups", "communities"} {
			if v, ok := all[k]; ok {
				out[k] = v
			}
		}
	}
	if slices.Contains(scopes, "email") {
		for _, k := range []string{"email", "email_verified"} {
			if v, ok := all[k]; ok {
				out[k] = v
			}
		}
	}
	// Extra keys let tests prove the allowlist and protocol-claim guard.
	for k, v := range all {
		if k == "iss" || k == "aud" || k == "x_extra" {
			out[k] = v
		}
	}
	return out, true, nil
}

func (t *memoryOIDCTx) CreateRefreshToken(ctx context.Context, rec *OIDCRefreshTokenRecord) error {
	rec.ID = t.b.id()
	cp := *rec
	t.b.refresh = append(t.b.refresh, &cp)
	return nil
}

func (t *memoryOIDCTx) RefreshTokenByHashForUpdate(ctx context.Context, tokenHash string) (*OIDCRefreshTokenRecord, error) {
	for _, rt := range t.b.refresh {
		if rt.TokenHash == tokenHash {
			cp := *rt
			return &cp, nil
		}
	}
	return nil, nil
}

func (t *memoryOIDCTx) MarkRefreshRotated(ctx context.Context, id, successorID uint) error {
	for _, rt := range t.b.refresh {
		if rt.ID == id {
			s := successorID
			rt.RotatedToID = &s
		}
	}
	return nil
}

func (t *memoryOIDCTx) RevokeRefreshLineage(ctx context.Context, startID uint, now time.Time) error {
	next := &startID
	for next != nil {
		var cur *OIDCRefreshTokenRecord
		for _, rt := range t.b.refresh {
			if rt.ID == *next {
				cur = rt
			}
		}
		if cur == nil {
			return nil
		}
		if cur.RevokedAt == nil {
			n := now
			cur.RevokedAt = &n
		}
		if cur.AccessTokenID != nil {
			t.revokeAccessToken(*cur.AccessTokenID, now)
		}
		next = cur.RotatedToID
	}
	return nil
}

// memoryKeyStore is an in-memory KeyStore built from GenerateSigningKey.
type memoryKeyStore struct {
	mu    sync.Mutex
	keys  []*memoryKey
	grace time.Duration
	// alg is the algorithm of a lazily generated first key.
	alg string
}

type memoryKey struct {
	kid       string
	alg       string
	signer    crypto.Signer
	public    crypto.PublicKey
	active    bool
	retiredAt *time.Time
}

func newMemoryKeyStore() *memoryKeyStore {
	return &memoryKeyStore{grace: 24 * time.Hour, alg: AlgRS256}
}

// rotate generates a new active key for alg and retires the old one at
// retiredAt.
func (k *memoryKeyStore) rotate(t *testing.T, alg string, retiredAt time.Time) string {
	t.Helper()
	k.mu.Lock()
	defer k.mu.Unlock()
	key, err := generateMemoryKey(alg)
	if err != nil {
		t.Fatal(err)
	}
	for _, old := range k.keys {
		if old.active {
			old.active = false
			r := retiredAt
			old.retiredAt = &r
		}
	}
	k.keys = append(k.keys, key)
	return key.kid
}

func generateMemoryKey(alg string) (*memoryKey, error) {
	kid, privPEM, pubPEM, err := GenerateSigningKey(alg)
	if err != nil {
		return nil, err
	}
	signer, err := ParsePrivateKeyPEM(privPEM)
	if err != nil {
		return nil, err
	}
	pub, err := ParsePublicKeyPEM(pubPEM)
	if err != nil {
		return nil, err
	}
	return &memoryKey{kid: kid, alg: alg, signer: signer, public: pub, active: true}, nil
}

func (k *memoryKeyStore) ActiveSigningKey(ctx context.Context) (SigningKey, error) {
	k.mu.Lock()
	defer k.mu.Unlock()
	for _, key := range k.keys {
		if key.active {
			return SigningKey{Kid: key.kid, Alg: key.alg, Signer: key.signer}, nil
		}
	}
	key, err := generateMemoryKey(k.alg)
	if err != nil {
		return SigningKey{}, err
	}
	k.keys = append(k.keys, key)
	return SigningKey{Kid: key.kid, Alg: key.alg, Signer: key.signer}, nil
}

func (k *memoryKeyStore) PublishedKeys(ctx context.Context, now time.Time) ([]PublicKey, error) {
	k.mu.Lock()
	defer k.mu.Unlock()
	var out []PublicKey
	for _, key := range k.keys {
		if key.active || (key.retiredAt != nil && now.Before(key.retiredAt.Add(k.grace))) {
			out = append(out, PublicKey{Kid: key.kid, Alg: key.alg, Key: key.public})
		}
	}
	return out, nil
}

// fakeOIDCHost is a settable OIDCHost that records what the engine asked
// it to render.
type fakeOIDCHost struct {
	mu          sync.Mutex
	userID      uint
	signedIn    bool
	views       []ConsentView
	errors      []string
	endSessions int
}

func (h *fakeOIDCHost) signIn(userID uint) {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.userID, h.signedIn = userID, true
}

func (h *fakeOIDCHost) signOut() {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.signedIn = false
}

func (h *fakeOIDCHost) SignedInUser(r *http.Request) (uint, bool, error) {
	h.mu.Lock()
	defer h.mu.Unlock()
	return h.userID, h.signedIn, nil
}

func (h *fakeOIDCHost) RenderAuthorizeError(w http.ResponseWriter, r *http.Request, code string) {
	h.mu.Lock()
	h.errors = append(h.errors, code)
	h.mu.Unlock()
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.WriteHeader(http.StatusBadRequest)
	_, _ = w.Write([]byte(code))
}

func (h *fakeOIDCHost) RenderConsent(w http.ResponseWriter, r *http.Request, view ConsentView) {
	h.mu.Lock()
	h.views = append(h.views, view)
	h.mu.Unlock()
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.WriteHeader(http.StatusOK)
	_, _ = w.Write([]byte("consent"))
}

func (h *fakeOIDCHost) EndSession(w http.ResponseWriter, r *http.Request) error {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.endSessions++
	h.signedIn = false
	return nil
}

func (h *fakeOIDCHost) lastView(t *testing.T) ConsentView {
	t.Helper()
	h.mu.Lock()
	defer h.mu.Unlock()
	if len(h.views) == 0 {
		t.Fatal("no consent view was rendered")
	}
	return h.views[len(h.views)-1]
}

// oidcFixture is an OIDC server over the memory fakes plus one active,
// confidential, public-subject client.
type oidcFixture struct {
	srv          *Server
	backend      *memoryOIDCBackend
	keys         *memoryKeyStore
	host         *fakeOIDCHost
	client       *OIDCClientRecord
	clientSecret string
}

const (
	testOIDCIssuer   = "https://issuer.example.test"
	testOIDCRedirect = "https://rp.example.test/callback"
	testOIDCSecret   = "oas_cs_test-secret-0123456789abcdefghijklmnopqrstuvwx"
)

func newOIDCFixture(t *testing.T, issuer string, mutate func(*OIDCOptions)) *oidcFixture {
	t.Helper()
	oidc := DefaultOIDCOptions()
	if mutate != nil {
		mutate(&oidc)
	}
	opts := DefaultOptions()
	opts.Issuer = issuer
	srv := NewOIDCServer(opts, oidc)
	f := &oidcFixture{
		srv:          srv,
		backend:      newMemoryOIDCBackend(),
		keys:         newMemoryKeyStore(),
		host:         &fakeOIDCHost{},
		clientSecret: testOIDCSecret,
	}
	if err := srv.SetOIDCBackend(f.backend, f.keys, f.host); err != nil {
		t.Fatal(err)
	}
	f.client = f.backend.addClient(OIDCClientRecord{
		ClientID:     "oas_0123456789abcdef01234567",
		Name:         "Example RP",
		Kind:         ClientKindOIDC,
		RedirectURIs: []string{testOIDCRedirect},
		Scopes:       []string{"openid", "profile", "email"},
		IsActive:     true,
		SubjectType:  SubjectTypePublic,
	}, testOIDCSecret)
	f.backend.setIdentity(7, map[string]any{
		"name": "Ada Example", "email": "ada@example.test", "email_verified": true,
		"groups": []string{"editors"}, "communities": []any{},
		"iss": "https://evil.example.test", "x_extra": "dropped",
	})
	return f
}

// mux mounts the OIDC routes the way a host application would.
func (f *oidcFixture) mux() *http.ServeMux {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /.well-known/openid-configuration", f.srv.Discovery)
	mux.HandleFunc("GET /oauth/jwks", f.srv.JWKS)
	mux.HandleFunc("GET /oauth/authorize", f.srv.Authorize)
	mux.HandleFunc("POST /oauth/authorize", f.srv.Authorize)
	mux.HandleFunc("GET /oauth/authorize/resume/{request}", f.srv.Resume)
	mux.HandleFunc("POST /oauth/authorize/switch-account", f.srv.SwitchAccount)
	mux.HandleFunc("POST /oauth/token", f.srv.Token)
	mux.HandleFunc("GET /oauth/userinfo", f.srv.UserInfo)
	return mux
}