// Filename: internal/repository/key_writer.go
package repository

import (
	"fmt"
	"gemini-balancer/internal/errors"
	"gemini-balancer/internal/models"
	"strconv"
	"time"
)

func (r *gormKeyRepository) UpdateKeyUsageTimestamp(groupID, keyID uint) {
	lruKey := fmt.Sprintf(KeyGroupLRU, groupID)
	timestamp := float64(time.Now().UnixMilli())

	members := map[string]float64{
		strconv.FormatUint(uint64(keyID), 10): timestamp,
	}

	if err := r.store.ZAdd(lruKey, members); err != nil {
		r.logger.WithError(err).Warnf("Failed to update usage timestamp for key %d in group %d", keyID, groupID)
	}
}

func (r *gormKeyRepository) SyncKeyStatusInPollingCaches(groupID, keyID uint, newStatus models.APIKeyStatus) {
	r.logger.Infof("SYNC: Directly updating polling caches for G:%d K:%d -> %s", groupID, keyID, newStatus)
	r.updatePollingCachesLogic(groupID, keyID, newStatus)
}

func (r *gormKeyRepository) HandleCacheUpdateEvent(groupID, keyID uint, newStatus models.APIKeyStatus) {
	r.logger.Infof("EVENT: Updating polling caches for G:%d K:%d -> %s from an event", groupID, keyID, newStatus)
	r.updatePollingCachesLogic(groupID, keyID, newStatus)
}

func (r *gormKeyRepository) updatePollingCachesLogic(groupID, keyID uint, newStatus models.APIKeyStatus) {
	keyIDStr := strconv.FormatUint(uint64(keyID), 10)
	sequentialKey := fmt.Sprintf(KeyGroupSequential, groupID)
	lruKey := fmt.Sprintf(KeyGroupLRU, groupID)
	mainPoolKey := fmt.Sprintf(KeyGroupRandomMain, groupID)
	cooldownPoolKey := fmt.Sprintf(KeyGroupRandomCooldown, groupID)

	_ = r.store.LRem(sequentialKey, 0, keyIDStr)
	_ = r.store.ZRem(lruKey, keyIDStr)
	_ = r.store.SRem(mainPoolKey, keyIDStr)
	_ = r.store.SRem(cooldownPoolKey, keyIDStr)

	if newStatus == models.StatusActive {
		if err := r.store.LPush(sequentialKey, keyIDStr); err != nil {
			r.logger.WithError(err).Warnf("Failed to add key %d to sequential list for group %d", keyID, groupID)
		}
		members := map[string]float64{keyIDStr: 0}
		if err := r.store.ZAdd(lruKey, members); err != nil {
			r.logger.WithError(err).Warnf("Failed to add key %d to LRU zset for group %d", keyID, groupID)
		}
		if err := r.store.SAdd(mainPoolKey, keyIDStr); err != nil {
			r.logger.WithError(err).Warnf("Failed to add key %d to random main pool for group %d", keyID, groupID)
		}
	}
}

// UpdateKeyStatusAfterRequest is the new central hub for handling feedback.
func (r *gormKeyRepository) UpdateKeyStatusAfterRequest(group *models.KeyGroup, key *models.APIKey, success bool, apiErr *errors.APIError) {
	if success {
		if group.PollingStrategy == models.StrategyWeighted {
			go r.UpdateKeyUsageTimestamp(group.ID, key.ID)
		}
		return
	}
	if apiErr == nil {
		r.logger.Warnf("Request failed for KeyID %d in GroupID %d but no specific API error was provided.", key.ID, group.ID)
		return
	}
	r.logger.Warnf("Request failed for KeyID %d in GroupID %d with error: %s. Temporarily removing from active polling caches.", key.ID, group.ID, apiErr.Message)

	// This call is correct. It uses the synchronous, direct method.
	r.SyncKeyStatusInPollingCaches(group.ID, key.ID, models.StatusCooldown)
}