package oidc

import (
	"crypto/rand"
	"encoding/base64"
	"sync"
	"time"
)

// PKCESession stores the in-flight authorization state server-side.
type PKCESession struct {
	Code                string
	ClientID            string
	RedirectURI         string
	PKCEChallenge       string // S256 challenge
	PKCEChallengeMethod string // always "S256"
	State               string
	Nonce               string
	Username            string // set after auth
	Scopes              []string
	ExpiresAt           time.Time
	// MFAVerified records whether this authorization actually required and
	// passed MFA validation (vs. MFA not being required for this user at
	// all). Feeds the assurance claim's level in token.go
	// (KEY-WP-0005-T01): aal2 when true, aal1 when false. Set once, at
	// completeAuthorization -- never re-derived from stale enrollment state
	// at token-exchange time.
	MFAVerified bool
}

// SessionStore is an in-memory PKCE session store.
type SessionStore struct {
	mu       sync.Mutex
	sessions map[string]*PKCESession // keyed by code
}

// NewSessionStore returns an initialised, empty SessionStore.
func NewSessionStore() *SessionStore {
	return &SessionStore{
		sessions: make(map[string]*PKCESession),
	}
}

// Create stores the session and returns the generated authorization code.
func (s *SessionStore) Create(sess *PKCESession) string {
	code := generateCode()
	sess.Code = code

	s.mu.Lock()
	s.sessions[code] = sess
	s.mu.Unlock()

	return code
}

// Get retrieves a session by code. Returns false if not found or expired.
func (s *SessionStore) Get(code string) (*PKCESession, bool) {
	s.mu.Lock()
	sess, ok := s.sessions[code]
	s.mu.Unlock()

	if !ok {
		return nil, false
	}
	if time.Now().After(sess.ExpiresAt) {
		s.Delete(code)
		return nil, false
	}
	return sess, true
}

// Consume retrieves a session by code and removes it in the same critical
// section, so exactly one caller can ever observe a given code. Returns false
// if the code is not present or has expired. The token endpoint must use this
// rather than Get/Delete: signing happens between those two calls, which is
// long enough for two concurrent exchanges to both observe the same session
// (KEY-WP-0016-T01). A consumed code is gone even if the exchange then fails,
// which is the intended single-use semantics -- a failed attempt must not leave
// a replayable code.
func (s *SessionStore) Consume(code string) (*PKCESession, bool) {
	s.mu.Lock()
	sess, ok := s.sessions[code]
	if ok {
		delete(s.sessions, code)
	}
	s.mu.Unlock()

	if !ok || time.Now().After(sess.ExpiresAt) {
		return nil, false
	}
	return sess, true
}

// Delete removes a session by code. No-op if the code is not present.
func (s *SessionStore) Delete(code string) {
	s.mu.Lock()
	delete(s.sessions, code)
	s.mu.Unlock()
}

// generateCode returns a cryptographically random, URL-safe string suitable
// for use as an authorization code.
func generateCode() string {
	b := make([]byte, 32)
	if _, err := rand.Read(b); err != nil {
		panic("oidc: failed to generate random code: " + err.Error())
	}
	return base64.RawURLEncoding.EncodeToString(b)
}