package auth

import (
	"crypto/rand"
	"fmt"
	"sync"
	"time"
)

// sessionEntry 存储一次 Phase 1 协商的会话密钥。
type sessionEntry struct {
	key    []byte // AES-256 会话密钥
	expiry time.Time
}

// SessionManager 管理内存中的会话密钥（Nonce → AES Key）。
// 每个 Phase 1 请求生成一个 session，Phase 2 消费后删除。
type SessionManager struct {
	mu    sync.Mutex
	store map[string]*sessionEntry
	ttl   time.Duration
}

// NewSessionManager 创建会话密钥管理器。
// ttl: 密钥过期时间（如 30 秒）
func NewSessionManager(ttl time.Duration) *SessionManager {
	sm := &SessionManager{
		store: make(map[string]*sessionEntry),
		ttl:   ttl,
	}
	// 后台协程定期清理过期密钥
	go sm.cleanupLoop()
	return sm
}

// GenerateKey 生成 AES-256 会话密钥，用指定的 keyID 存储（通常用 nonce）。
func (sm *SessionManager) GenerateKey(keyID string) ([]byte, error) {
	key := make([]byte, 32)
	if _, err := rand.Read(key); err != nil {
		return nil, fmt.Errorf("generate session key: %w", err)
	}

	sm.mu.Lock()
	sm.store[keyID] = &sessionEntry{
		key:    key,
		expiry: time.Now().Add(sm.ttl),
	}
	sm.mu.Unlock()

	return key, nil
}

// GetKey 获取并删除会话密钥（一次性使用）。
// 返回 nil 表示 keyID 不存在或已过期。
func (sm *SessionManager) GetKey(keyID string) []byte {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	entry, ok := sm.store[keyID]
	if !ok {
		return nil
	}

	delete(sm.store, keyID)

	if time.Now().After(entry.expiry) {
		return nil
	}

	return entry.key
}

// cleanupLoop 定期清理过期密钥。
func (sm *SessionManager) cleanupLoop() {
	ticker := time.NewTicker(sm.ttl)
	defer ticker.Stop()

	for range ticker.C {
		sm.mu.Lock()
		now := time.Now()
		for id, entry := range sm.store {
			if now.After(entry.expiry) {
				delete(sm.store, id)
			}
		}
		sm.mu.Unlock()
	}
}

// NonceCache 用于防重放攻击的 Nonce 缓存。
type NonceCache struct {
	mu    sync.Mutex
	store map[string]time.Time
	ttl   time.Duration
}

// NewNonceCache 创建 Nonce 缓存。
// ttl: Nonce 有效时间（如 60 秒）
func NewNonceCache(ttl time.Duration) *NonceCache {
	nc := &NonceCache{
		store: make(map[string]time.Time),
		ttl:   ttl,
	}
	go nc.cleanupLoop()
	return nc
}

// Check 检查并记录 Nonce。返回 true 表示 Nonce 有效（未使用过）。
// 返回 false 表示 Nonce 已存在（重放攻击）。
func (nc *NonceCache) Check(nonce string) bool {
	nc.mu.Lock()
	defer nc.mu.Unlock()

	if _, exists := nc.store[nonce]; exists {
		return false
	}

	nc.store[nonce] = time.Now().Add(nc.ttl)
	return true
}

func (nc *NonceCache) cleanupLoop() {
	ticker := time.NewTicker(nc.ttl)
	defer ticker.Stop()

	for range ticker.C {
		nc.mu.Lock()
		now := time.Now()
		for n, expiry := range nc.store {
			if now.After(expiry) {
				delete(nc.store, n)
			}
		}
		nc.mu.Unlock()
	}
}