package auth

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"sync"
	"time"
)

// Sentinel errors for key validation and lookup. Match with errors.Is.
var (
	ErrInvalidKey  = errors.New("invalid API key")
	ErrKeyInactive = errors.New("API key is inactive")
	ErrKeyExpired  = errors.New("API key has expired")
	ErrKeyNotFound = errors.New("API key not found")
)

// APIKeyStore manages API keys with rate limiting.
// Keys are stored indexed by SHA-256 digest; the raw key is never retained
// after creation.
type APIKeyStore struct {
	mu         sync.RWMutex
	keys       map[string]*APIKey // keyed by KeyHash
	rateLimits map[string]*RateLimiter
}

// APIKey represents an API key with metadata.
// Key holds the raw secret only on the value returned from key creation;
// stored and persisted records carry only KeyHash and KeyPreview.
type APIKey struct {
	Key        string     `json:"key,omitempty"`
	KeyHash    string     `json:"-"`
	KeyPreview string     `json:"key_preview,omitempty"` // e.g. "nk_…ab12"
	Name       string     `json:"name"`
	ClientID   string     `json:"client_id"`
	Roles      []string   `json:"roles"`
	CreatedAt  time.Time  `json:"created_at"`
	LastUsedAt *time.Time `json:"last_used_at,omitempty"`
	ExpiresAt  *time.Time `json:"expires_at,omitempty"`
	IsActive   bool       `json:"is_active"`
	RateLimit  int        `json:"rate_limit"` // requests per minute, 0 = unlimited
}

// RateLimiter tracks rate limiting for a key
type RateLimiter struct {
	maxRequests int
	window      time.Duration
	resetTime   time.Time
	count       int
	mu          sync.Mutex
}

// AuthContext holds auth information attached to request context
type AuthContext struct {
	APIKey   *APIKey
	ClientID string
	Roles    []string
}

// HashKey returns the hex-encoded SHA-256 digest of a raw API key.
// All storage and lookups are keyed by this digest so a leaked store or
// database dump does not expose usable credentials.
func HashKey(rawKey string) string {
	sum := sha256.Sum256([]byte(rawKey))
	return hex.EncodeToString(sum[:])
}

// keyPreview returns a non-sensitive display form of a raw key ("nk_…ab12").
func keyPreview(rawKey string) string {
	if len(rawKey) < 4 {
		return ""
	}
	return "nk_…" + rawKey[len(rawKey)-4:]
}

// generateAPIKey creates a new APIKey with a random secret without touching
// any store. The returned value carries the raw secret in Key; callers decide
// where (and whether) to register it.
func generateAPIKey(clientID string, roles []string, rateLimit int, expiresIn *time.Duration) (*APIKey, error) {
	keyBytes := make([]byte, 32)
	if _, err := rand.Read(keyBytes); err != nil {
		return nil, fmt.Errorf("failed to generate key: %w", err)
	}
	key := "nk_" + hex.EncodeToString(keyBytes)

	now := time.Now().UTC()
	apiKey := &APIKey{
		Key:        key,
		KeyHash:    HashKey(key),
		KeyPreview: keyPreview(key),
		ClientID:   clientID,
		Roles:      roles,
		CreatedAt:  now,
		IsActive:   true,
		RateLimit:  rateLimit,
		Name:       fmt.Sprintf("%s-%d", clientID, now.Unix()),
	}

	if expiresIn != nil {
		expiresAt := now.Add(*expiresIn)
		apiKey.ExpiresAt = &expiresAt
	}

	return apiKey, nil
}

// NewAPIKeyStore creates a new API key store
func NewAPIKeyStore() *APIKeyStore {
	return &APIKeyStore{
		keys:       make(map[string]*APIKey),
		rateLimits: make(map[string]*RateLimiter),
	}
}

// CreateKey generates a new API key and registers it in the store.
// The returned APIKey carries the raw secret; the stored copy does not.
func (s *APIKeyStore) CreateKey(clientID string, roles []string, rateLimit int, expiresIn *time.Duration) (*APIKey, error) {
	apiKey, err := generateAPIKey(clientID, roles, rateLimit, expiresIn)
	if err != nil {
		return nil, err
	}
	s.RegisterKey(apiKey)
	return apiKey, nil
}

// RegisterKey adds a key to the store, indexed by digest. The stored copy has
// the raw secret stripped. If KeyHash is unset it is computed from Key.
func (s *APIKeyStore) RegisterKey(apiKey *APIKey) {
	stored := *apiKey
	if stored.KeyHash == "" && stored.Key != "" {
		stored.KeyHash = HashKey(stored.Key)
	}
	if stored.KeyPreview == "" && stored.Key != "" {
		stored.KeyPreview = keyPreview(stored.Key)
	}
	stored.Key = "" // never retain the raw secret

	s.mu.Lock()
	defer s.mu.Unlock()
	s.keys[stored.KeyHash] = &stored
	s.rateLimits[stored.KeyHash] = &RateLimiter{
		maxRequests: stored.RateLimit,
		window:      time.Minute,
		resetTime:   time.Now().Add(time.Minute),
		count:       0,
	}
}

// RemoveKeyByHash deletes a key and its rate limiter from the store.
func (s *APIKeyStore) RemoveKeyByHash(keyHash string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.keys, keyHash)
	delete(s.rateLimits, keyHash)
}

// ValidateKey checks if a raw API key is valid and returns the key metadata
func (s *APIKeyStore) ValidateKey(keyStr string) (*APIKey, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	key, exists := s.keys[HashKey(keyStr)]
	if !exists {
		return nil, ErrInvalidKey
	}

	if !key.IsActive {
		return nil, ErrKeyInactive
	}

	if key.ExpiresAt != nil && time.Now().After(*key.ExpiresAt) {
		return nil, ErrKeyExpired
	}

	return key, nil
}

// CheckRateLimit checks if a key has exceeded its rate limit
func (s *APIKeyStore) CheckRateLimit(keyStr string) (bool, error) {
	keyHash := HashKey(keyStr)

	// Look up key and limiter under the store lock, then release it
	// before acquiring the per-key limiter lock to avoid nested locking.
	s.mu.RLock()
	key, exists := s.keys[keyHash]
	if !exists {
		s.mu.RUnlock()
		return false, ErrInvalidKey
	}

	// Unlimited rate limit
	if key.RateLimit <= 0 {
		s.mu.RUnlock()
		return true, nil
	}

	limiter, exists := s.rateLimits[keyHash]
	if !exists {
		s.mu.RUnlock()
		return false, fmt.Errorf("rate limiter not found")
	}
	s.mu.RUnlock()

	// Now lock only the per-key rate limiter
	limiter.mu.Lock()
	defer limiter.mu.Unlock()

	now := time.Now()
	if now.After(limiter.resetTime) {
		limiter.count = 0
		limiter.resetTime = now.Add(limiter.window)
	}

	if limiter.count >= limiter.maxRequests {
		return false, nil
	}

	limiter.count++
	return true, nil
}

// UpdateLastUsed updates the last used timestamp for a key
func (s *APIKeyStore) UpdateLastUsed(keyStr string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	key, exists := s.keys[HashKey(keyStr)]
	if !exists {
		return ErrInvalidKey
	}

	now := time.Now().UTC()
	key.LastUsedAt = &now
	return nil
}

// DeactivateKey deactivates an API key by its raw value
func (s *APIKeyStore) DeactivateKey(keyStr string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	key, exists := s.keys[HashKey(keyStr)]
	if !exists {
		return ErrInvalidKey
	}

	key.IsActive = false
	return nil
}

// GetKey retrieves key metadata by raw key (for management purposes)
func (s *APIKeyStore) GetKey(keyStr string) (*APIKey, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	key, exists := s.keys[HashKey(keyStr)]
	if !exists {
		return nil, ErrKeyNotFound
	}

	return key, nil
}

// ListKeys lists all API keys for a client
func (s *APIKeyStore) ListKeys(clientID string) []*APIKey {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var keys []*APIKey
	for _, key := range s.keys {
		if key.ClientID == clientID {
			keys = append(keys, key)
		}
	}
	return keys
}

// authContextKey is an unexported type for context keys to avoid collisions.
type authContextKey struct{}

// ContextWithAuth adds auth context to a request context
func ContextWithAuth(ctx context.Context, auth *AuthContext) context.Context {
	return context.WithValue(ctx, authContextKey{}, auth)
}

// GetAuthContext retrieves auth context from a request context
func GetAuthContext(ctx context.Context) (*AuthContext, bool) {
	auth, ok := ctx.Value(authContextKey{}).(*AuthContext)
	return auth, ok
}