package routing

import (
	"errors"
	"sort"
	"strconv"
	"sync"
	"sync/atomic"

	"aigw/internal/catalog"
	"aigw/internal/domain"
)

var ErrNoRoute = errors.New("no compatible upstream route")

type Router struct {
	catalog  *catalog.Catalog
	counters sync.Map
}

func New(catalog *catalog.Catalog) *Router {
	return &Router{catalog: catalog}
}

func (r *Router) Plan(modelID string, protocol domain.Protocol) ([]domain.Route, error) {
	model, err := r.catalog.Model(modelID)
	if err != nil {
		return nil, err
	}
	routes := make([]domain.Route, 0, len(model.Routes))
	for _, route := range model.Routes {
		if route.Provider.Protocol == protocol {
			routes = append(routes, route)
		}
	}
	if len(routes) == 0 {
		return nil, ErrNoRoute
	}

	sort.SliceStable(routes, func(i, j int) bool { return routes[i].Priority < routes[j].Priority })
	result := make([]domain.Route, 0, len(routes))
	for start := 0; start < len(routes); {
		end := start + 1
		for end < len(routes) && routes[end].Priority == routes[start].Priority {
			end++
		}
		result = append(result, r.rotate(modelID, protocol, routes[start:end])...)
		start = end
	}
	return result, nil
}

func (r *Router) rotate(modelID string, protocol domain.Protocol, routes []domain.Route) []domain.Route {
	if len(routes) < 2 {
		return append([]domain.Route(nil), routes...)
	}
	key := modelID + "\x00" + string(protocol) + "\x00" + strconv.Itoa(routes[0].Priority)
	counterValue, _ := r.counters.LoadOrStore(key, &atomic.Uint64{})
	counter := counterValue.(*atomic.Uint64).Add(1) - 1

	totalWeight := 0
	for _, route := range routes {
		totalWeight += route.Weight
	}
	position := int(counter % uint64(totalWeight))
	selected := 0
	for i, route := range routes {
		if position < route.Weight {
			selected = i
			break
		}
		position -= route.Weight
	}

	result := make([]domain.Route, 0, len(routes))
	result = append(result, routes[selected])
	for offset := 1; offset < len(routes); offset++ {
		result = append(result, routes[(selected+offset)%len(routes)])
	}
	return result
}