opencatd-open/internal/channel/channel.go

package channel

import (
	"context"
	"fmt"
	"log"
	"math/rand"
	"opencatd-open/internal/dao"
	"opencatd-open/internal/store"
	"opencatd-open/internal/pkg/crypto"
	"sync"
	"time"
)

type Service struct {
	channelDAO *dao.ChannelDAO
	modelDAO   *dao.ModelDAO

	// Health tracking
	mu           sync.RWMutex
	healthStatus map[uint64]*channelHealth
}

type channelHealth struct {
	status      string
	consecutive int
	lastCheck   time.Time
	cooldown    time.Time
}

func NewService(channelDAO *dao.ChannelDAO, modelDAO *dao.ModelDAO) *Service {
	return &Service{
		channelDAO:  channelDAO,
		modelDAO:    modelDAO,
		healthStatus: make(map[uint64]*channelHealth),
	}
}

// SelectChannel selects the best channel for a given model using weighted random selection
func (s *Service) SelectChannel(ctx context.Context, modelName string) (*store.Channel, error) {
	channels, err := s.channelDAO.GetEnabledChannelsByModel(modelName)
	if err != nil {
		return nil, fmt.Errorf("failed to get channels for model %s: %w", modelName, err)
	}
	if len(channels) == 0 {
		return nil, fmt.Errorf("no enabled channels for model: %s", modelName)
	}

	// Filter out unhealthy channels
	candidates := s.filterHealthy(channels)
	if len(candidates) == 0 {
		// If all channels are unhealthy, try the first one anyway
		candidates = channels[:1]
	}

	// Weighted random selection
	totalWeight := 0
	for _, ch := range candidates {
		totalWeight += ch.Weight
	}
	if totalWeight == 0 {
		return candidates[0], nil
	}

	r := rand.Intn(totalWeight)
	for _, ch := range candidates {
		r -= ch.Weight
		if r < 0 {
			return ch, nil
		}
	}

	return candidates[0], nil
}

// GetChannelByKeyID decrypts the API key for a channel
func (s *Service) GetChannelByKeyID(ctx context.Context, channelID uint64) (*store.Channel, error) {
	ch, err := s.channelDAO.GetByID(channelID)
	if err != nil {
		return nil, err
	}
	return ch, nil
}

// GetAPIKey decrypts the channel's API key
func (s *Service) GetAPIKey(ch *store.Channel) (string, error) {
	return crypto.Decrypt(ch.APIKeyEnc)
}

// RecordSuccess records a successful request to a channel
func (s *Service) RecordSuccess(channelID uint64) {
	s.mu.Lock()
	defer s.mu.Unlock()

	h := s.getOrCreateHealth(channelID)
	h.consecutive = 0
	h.status = store.ChannelHealthHealthy
	h.lastCheck = time.Now()
}

// RecordFailure records a failed request to a channel
func (s *Service) RecordFailure(channelID uint64) {
	s.mu.Lock()
	defer s.mu.Unlock()

	h := s.getOrCreateHealth(channelID)
	h.consecutive++
	h.lastCheck = time.Now()

	if h.consecutive >= 3 {
		h.status = store.ChannelHealthDegraded
		h.cooldown = time.Now().Add(5 * time.Minute)
	}
	if h.consecutive >= 5 {
		h.status = store.ChannelHealthCooldown
		h.cooldown = time.Now().Add(15 * time.Minute)
	}
}

// RecordTimeout records a timeout to a channel
func (s *Service) RecordTimeout(channelID uint64) {
	s.RecordFailure(channelID)
}

func (s *Service) getOrCreateHealth(channelID uint64) *channelHealth {
	h, ok := s.healthStatus[channelID]
	if !ok {
		h = &channelHealth{
			status: store.ChannelHealthHealthy,
		}
		s.healthStatus[channelID] = h
	}
	return h
}

func (s *Service) filterHealthy(channels []*store.Channel) []*store.Channel {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var healthy []*store.Channel
	now := time.Now()

	for _, ch := range channels {
		h, ok := s.healthStatus[ch.ID]
		if !ok {
			healthy = append(healthy, ch)
			continue
		}

		// Check if cooldown has expired
		if now.After(h.cooldown) && h.cooldown.IsZero() == false {
			h.consecutive = 0
			h.status = store.ChannelHealthHealthy
			healthy = append(healthy, ch)
			continue
		}

		if h.status == store.ChannelHealthHealthy || h.status == store.ChannelHealthDegraded {
			healthy = append(healthy, ch)
		}
	}

	return healthy
}

// GetHealthStatus returns the health status of a channel
func (s *Service) GetHealthStatus(channelID uint64) string {
	s.mu.RLock()
	defer s.mu.RUnlock()

	h, ok := s.healthStatus[channelID]
	if !ok {
		return store.ChannelHealthHealthy
	}
	return h.status
}

// ChannelCandidate represents a channel with its resolved API key
type ChannelCandidate struct {
	Channel *store.Channel
	APIKey  string
	Format  string
}

// SelectCandidates returns candidates for a model, sorted by priority
func (s *Service) SelectCandidates(ctx context.Context, modelName string, preferredFormat string) ([]ChannelCandidate, error) {
	channels, err := s.channelDAO.GetEnabledChannelsByModel(modelName)
	if err != nil {
		return nil, err
	}

	var candidates []ChannelCandidate
	for _, ch := range channels {
		// Check if channel supports the preferred format
		formats := ch.FormatsEffective()
		supported := false
		for _, f := range formats {
			if f == preferredFormat || preferredFormat == "" {
				supported = true
				break
			}
		}
		if !supported {
			continue
		}

		apiKey, err := crypto.Decrypt(ch.APIKeyEnc)
		if err != nil {
			log.Printf("Failed to decrypt API key for channel %s: %v", ch.Name, err)
			continue
		}

		candidates = append(candidates, ChannelCandidate{
			Channel: ch,
			APIKey:  apiKey,
			Format:  preferredFormat,
		})
	}

	return candidates, nil
}