package auth

import (
	"context"
	"sync"
	"sync/atomic"

	"github.com/google/uuid"
)

// Cache is the in-memory write-through projection of the active rows in
// `backend.device_sessions`. Reads (Get) are RLocked; writes (Add,
// Remove, RemoveByUser) are Locked. The cache holds two maps:
//
//   - byID    maps device_session_id → Session.
//   - byUser  maps user_id → set of device_session_ids belonging to that
//     user, used to satisfy bulk revoke without scanning byID.
//
// Both maps are updated atomically inside one Lock per mutation. The
// caller is expected to commit the corresponding database write *before*
// invoking Add or Remove so that the cache stays consistent under crash:
// a Postgres commit failure leaves the cache untouched, matching the
// previous DB state.
type Cache struct {
	mu     sync.RWMutex
	byID   map[uuid.UUID]Session
	byUser map[uuid.UUID]map[uuid.UUID]struct{}
	ready  atomic.Bool
}

// NewCache constructs an empty Cache. The cache reports Ready() == false
// until Warm completes successfully.
func NewCache() *Cache {
	return &Cache{
		byID:   make(map[uuid.UUID]Session),
		byUser: make(map[uuid.UUID]map[uuid.UUID]struct{}),
	}
}

// Warm replaces the cache contents with every active session loaded from
// store. It is intended to be called exactly once at process boot before
// the HTTP listener accepts traffic; successful completion flips Ready
// to true. Subsequent calls re-warm the cache (useful in tests).
func (c *Cache) Warm(ctx context.Context, store *Store) error {
	sessions, err := store.ListActiveSessions(ctx)
	if err != nil {
		return err
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	c.byID = make(map[uuid.UUID]Session, len(sessions))
	c.byUser = make(map[uuid.UUID]map[uuid.UUID]struct{})
	for _, s := range sessions {
		c.byID[s.DeviceSessionID] = s
		set, ok := c.byUser[s.UserID]
		if !ok {
			set = make(map[uuid.UUID]struct{})
			c.byUser[s.UserID] = set
		}
		set[s.DeviceSessionID] = struct{}{}
	}
	c.ready.Store(true)
	return nil
}

// Ready reports whether Warm has completed at least once. The HTTP
// readiness probe wires through this method so `/readyz` only flips to
// 200 after the cache is hydrated.
func (c *Cache) Ready() bool {
	if c == nil {
		return false
	}
	return c.ready.Load()
}

// Size returns the number of cached active sessions. Useful in startup
// logs ("auth cache warmed: N sessions") and in tests.
func (c *Cache) Size() int {
	if c == nil {
		return 0
	}
	c.mu.RLock()
	defer c.mu.RUnlock()
	return len(c.byID)
}

// Get returns the session with deviceSessionID and a presence flag.
// Misses always return the zero Session and false; callers should not
// inspect the returned value when ok is false.
func (c *Cache) Get(deviceSessionID uuid.UUID) (Session, bool) {
	if c == nil {
		return Session{}, false
	}
	c.mu.RLock()
	defer c.mu.RUnlock()
	s, ok := c.byID[deviceSessionID]
	return s, ok
}

// Add stores s in the cache. It is safe to call on an existing entry
// — both the primary map and the user index are updated to the latest
// snapshot.
func (c *Cache) Add(s Session) {
	if c == nil {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	c.byID[s.DeviceSessionID] = s
	set, ok := c.byUser[s.UserID]
	if !ok {
		set = make(map[uuid.UUID]struct{})
		c.byUser[s.UserID] = set
	}
	set[s.DeviceSessionID] = struct{}{}
}

// Remove evicts the entry for deviceSessionID from both maps. Calling
// Remove on a missing entry is a no-op.
func (c *Cache) Remove(deviceSessionID uuid.UUID) {
	if c == nil {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	s, ok := c.byID[deviceSessionID]
	if !ok {
		return
	}
	delete(c.byID, deviceSessionID)
	if set := c.byUser[s.UserID]; set != nil {
		delete(set, deviceSessionID)
		if len(set) == 0 {
			delete(c.byUser, s.UserID)
		}
	}
}

// ListByUser returns a freshly-allocated snapshot of every cached
// session belonging to userID. The user-surface "list my sessions"
// handler consumes this. An empty slice is returned for an unknown
// userID.
func (c *Cache) ListByUser(userID uuid.UUID) []Session {
	if c == nil {
		return nil
	}
	c.mu.RLock()
	defer c.mu.RUnlock()
	set, ok := c.byUser[userID]
	if !ok {
		return nil
	}
	out := make([]Session, 0, len(set))
	for id := range set {
		if sess, ok := c.byID[id]; ok {
			out = append(out, sess)
		}
	}
	return out
}

// RemoveByUser evicts every cached entry belonging to userID and returns
// the device_session_ids it removed. The returned slice is safe for the
// caller to hold past the call — it is freshly allocated.
func (c *Cache) RemoveByUser(userID uuid.UUID) []uuid.UUID {
	if c == nil {
		return nil
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	set, ok := c.byUser[userID]
	if !ok {
		return nil
	}
	removed := make([]uuid.UUID, 0, len(set))
	for id := range set {
		removed = append(removed, id)
		delete(c.byID, id)
	}
	delete(c.byUser, userID)
	return removed
}