Atlas/internal/provider/router.go

package provider

import (
	"context"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/rose_cat707/Atlas/internal/domain"
	"github.com/rose_cat707/Atlas/internal/service/reqlog"
	"github.com/rose_cat707/Atlas/internal/store"
)

type Config struct {
	ID          int64
	Name        string
	Protocol    string
	BaseURL     string
	APIKey      string
	ServiceType string
	ExtraJSON   string
}

type Router struct {
	store        *store.Store
	client       *http.Client
	streamClient *http.Client
	logger       *reqlog.Logger
}

func NewRouter(st *store.Store, logger *reqlog.Logger) *Router {
	return &Router{
		store: st,
		client: &http.Client{
			Timeout: 120 * time.Second,
		},
		streamClient: &http.Client{},
		logger:       logger,
	}
}

func (r *Router) Resolve(ctx context.Context, serviceType string, providerID int64) (*Config, error) {
	if providerID > 0 {
		return r.resolveDirect(ctx, serviceType, providerID)
	}
	for _, st := range resolveServiceTypeCandidates(serviceType) {
		if cfg, err := r.resolveDirect(ctx, st, 0); err == nil {
			return cfg, nil
		}
	}
	return nil, fmt.Errorf("no provider for %s", serviceType)
}

func resolveServiceTypeCandidates(serviceType string) []string {
	if !domain.IsImageProviderServiceType(serviceType) {
		return []string{serviceType}
	}
	seen := map[string]struct{}{}
	var out []string
	for _, st := range domain.ImageProviderServiceTypes() {
		if _, ok := seen[st]; ok {
			continue
		}
		seen[st] = struct{}{}
		out = append(out, st)
	}
	return out
}

func (r *Router) resolveDirect(ctx context.Context, serviceType string, providerID int64) (*Config, error) {
	var p *store.Provider
	var err error
	if providerID > 0 {
		p, err = r.store.GetProvider(ctx, providerID)
	} else {
		p, err = r.store.DefaultProvider(ctx, serviceType)
	}
	if err != nil {
		return nil, fmt.Errorf("no provider for %s: %w", serviceType, err)
	}
	return r.providerToConfig(p)
}

func (r *Router) ResolveByProtocol(ctx context.Context, serviceType, protocol string, providerID int64) (*Config, error) {
	if providerID > 0 {
		p, err := r.store.GetProvider(ctx, providerID)
		if err != nil {
			return nil, err
		}
		return r.providerToConfig(p)
	}
	for _, st := range resolveServiceTypeCandidates(serviceType) {
		providers, err := r.store.ListProviders(ctx, st, true)
		if err != nil {
			continue
		}
		for _, p := range providers {
			if protocol == "" || p.Protocol == protocol {
				return r.providerToConfig(&p)
			}
		}
	}
	return nil, fmt.Errorf("no provider for %s/%s", serviceType, protocol)
}

func (r *Router) providerToConfig(p *store.Provider) (*Config, error) {
	if !p.Enabled {
		return nil, fmt.Errorf("provider %d disabled", p.ID)
	}
	return &Config{
		ID:          p.ID,
		Name:        p.Name,
		Protocol:    p.Protocol,
		BaseURL:     strings.TrimRight(p.BaseURL, "/"),
		APIKey:      p.APIKey,
		ServiceType: p.ServiceType,
		ExtraJSON:   p.ExtraJSON,
	}, nil
}

type ProxyResult struct {
	StatusCode int
	Body       []byte
	Headers    http.Header
}