package catalog

import (
	"fmt"
	"sort"
	"strings"
	"sync/atomic"

	"aigw/internal/config"
	"aigw/internal/domain"
)

type Catalog struct {
	state atomic.Pointer[snapshot]
}

type snapshot struct {
	models map[string]domain.Model
	list   []domain.Model
}

func New(cfg config.Config) *Catalog {
	providers := make(map[string]domain.Provider, len(cfg.Providers))
	for _, provider := range cfg.Providers {
		providers[provider.ID] = domain.Provider{
			ID:       provider.ID,
			Protocol: provider.Protocol,
			BaseURL:  strings.TrimRight(provider.BaseURL, "/"),
			APIKey:   provider.APIKey,
		}
	}

	models := make([]domain.Model, 0, len(cfg.Models))
	for _, modelCfg := range cfg.Models {
		model := domain.Model{ID: modelCfg.ID, OwnedBy: modelCfg.OwnedBy}
		for _, route := range modelCfg.Routes {
			model.Routes = append(model.Routes, domain.Route{
				Provider:      providers[route.Provider],
				UpstreamModel: route.UpstreamModel,
				Priority:      route.Priority,
				Weight:        route.Weight,
			})
		}
		models = append(models, model)
	}
	return NewModels(models)
}

func NewModels(models []domain.Model) *Catalog {
	catalog := &Catalog{}
	catalog.Replace(models)
	return catalog
}

func (c *Catalog) Replace(source []domain.Model) {
	models := make(map[string]domain.Model, len(source))
	list := make([]domain.Model, 0, len(source))
	for _, sourceModel := range source {
		model := sourceModel
		model.Routes = append([]domain.Route(nil), sourceModel.Routes...)
		models[model.ID] = model
		list = append(list, model)
	}
	sort.Slice(list, func(i, j int) bool { return list[i].ID < list[j].ID })
	c.state.Store(&snapshot{models: models, list: list})
}

func (c *Catalog) Model(id string) (domain.Model, error) {
	current := c.state.Load()
	if current == nil {
		return domain.Model{}, fmt.Errorf("model %q not found", id)
	}
	model, ok := current.models[id]
	if !ok {
		return domain.Model{}, fmt.Errorf("model %q not found", id)
	}
	return model, nil
}

func (c *Catalog) Models(protocol domain.Protocol) []domain.Model {
	current := c.state.Load()
	if current == nil {
		return nil
	}
	result := make([]domain.Model, 0, len(current.list))
	for _, model := range current.list {
		for _, route := range model.Routes {
			if route.Provider.Protocol == protocol {
				result = append(result, model)
				break
			}
		}
	}
	return result
}

func (c *Catalog) Count() int {
	current := c.state.Load()
	if current == nil {
		return 0
	}
	return len(current.list)
}