ez-api/internal/service/sync.go

package service

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"

	"github.com/ez-api/ez-api/internal/model"
	"github.com/redis/go-redis/v9"
	"gorm.io/gorm"
)

type SyncService struct {
	rdb *redis.Client
}

func NewSyncService(rdb *redis.Client) *SyncService {
	return &SyncService{rdb: rdb}
}

func (s *SyncService) SyncKey(key *model.Key) error {
	tokenHash := hashToken(key.KeySecret)
	redisKey := fmt.Sprintf("auth:token:%s", tokenHash)

	fields := map[string]interface{}{
		"status":  key.Status,
		"balance": key.Balance,
	}
	if key.ProviderID != nil {
		fields["provider_id"] = *key.ProviderID
	} else {
		fields["provider_id"] = 0
	}

	return s.rdb.HSet(context.Background(), redisKey, fields).Err()
}

type providerSnapshot struct {
	ID      uint   `json:"id"`
	Name    string `json:"name"`
	Type    string `json:"type"`
	BaseURL string `json:"base_url"`
	APIKey  string `json:"api_key"`
}

type keySnapshot struct {
	ID         uint    `json:"id"`
	ProviderID uint    `json:"provider_id"`
	TokenHash  string  `json:"token_hash"`
	Status     string  `json:"status"`
	Weight     int     `json:"weight"`
	Balance    float64 `json:"balance"`
}

type modelSnapshot struct {
	Name               string  `json:"name"`
	ContextWindow      int     `json:"context_window"`
	CostPerToken       float64 `json:"cost_per_token"`
	SupportsVision     bool    `json:"supports_vision"`
	SupportsFunction   bool    `json:"supports_functions"`
	SupportsToolChoice bool    `json:"supports_tool_choice"`
	SupportsFIM        bool    `json:"supports_fim"`
	MaxOutputTokens    int     `json:"max_output_tokens"`
}

// SyncAll writes full snapshots (providers/keys/models) into Redis for DP consumption.
func (s *SyncService) SyncAll(db *gorm.DB) error {
	ctx := context.Background()

	// Providers snapshot
	var providers []model.Provider
	if err := db.Find(&providers).Error; err != nil {
		return fmt.Errorf("load providers: %w", err)
	}
	providerSnap := make([]providerSnapshot, 0, len(providers))
	for _, p := range providers {
		providerSnap = append(providerSnap, providerSnapshot{
			ID:      p.ID,
			Name:    p.Name,
			Type:    p.Type,
			BaseURL: p.BaseURL,
			APIKey:  p.APIKey,
		})
	}
	if err := s.storeJSON(ctx, "config:providers", providerSnap); err != nil {
		return err
	}

	// Keys snapshot + auth hashes
	var keys []model.Key
	if err := db.Find(&keys).Error; err != nil {
		return fmt.Errorf("load keys: %w", err)
	}
	keySnap := make([]keySnapshot, 0, len(keys))
	for _, k := range keys {
		tokenHash := hashToken(k.KeySecret)
		keySnap = append(keySnap, keySnapshot{
			ID:         k.ID,
			ProviderID: firstID(k.ProviderID),
			TokenHash:  tokenHash,
			Status:     k.Status,
			Weight:     k.Weight,
			Balance:    k.Balance,
		})

		// Maintain per-token auth hash for quick checks
		fields := map[string]interface{}{
			"status":      k.Status,
			"provider_id": firstID(k.ProviderID),
			"weight":      k.Weight,
			"balance":     k.Balance,
		}
		if err := s.rdb.HSet(ctx, fmt.Sprintf("auth:token:%s", tokenHash), fields).Err(); err != nil {
			return fmt.Errorf("write auth token: %w", err)
		}
	}
	if err := s.storeJSON(ctx, "config:keys", keySnap); err != nil {
		return err
	}

	// Models snapshot
	var models []model.Model
	if err := db.Find(&models).Error; err != nil {
		return fmt.Errorf("load models: %w", err)
	}
	modelSnap := make([]modelSnapshot, 0, len(models))
	for _, m := range models {
		modelSnap = append(modelSnap, modelSnapshot{
			Name:               m.Name,
			ContextWindow:      m.ContextWindow,
			CostPerToken:       m.CostPerToken,
			SupportsVision:     m.SupportsVision,
			SupportsFunction:   m.SupportsFunctions,
			SupportsToolChoice: m.SupportsToolChoice,
			SupportsFIM:        m.SupportsFIM,
			MaxOutputTokens:    m.MaxOutputTokens,
		})
	}
	if err := s.storeJSON(ctx, "meta:models", modelSnap); err != nil {
		return err
	}

	return nil
}

func (s *SyncService) storeJSON(ctx context.Context, key string, val interface{}) error {
	payload, err := json.Marshal(val)
	if err != nil {
		return fmt.Errorf("marshal %s: %w", key, err)
	}
	if err := s.rdb.Set(ctx, key, payload, 0).Err(); err != nil {
		return fmt.Errorf("write %s: %w", key, err)
	}
	return nil
}

func hashToken(token string) string {
	hasher := sha256.New()
	hasher.Write([]byte(token))
	return hex.EncodeToString(hasher.Sum(nil))
}

func firstID(id *uint) uint {
	if id == nil {
		return 0
	}
	return *id
}