package resolver

import "sort"

type ResolverArtifacts struct {
	IPs           []string
	DirectIPs     []string
	WildcardIPs   []string
	IPMap         [][2]string
	DirectIPMap   [][2]string
	WildcardIPMap [][2]string
}

func BuildResolverArtifacts(resolved map[string][]string, staticLabels map[string][]string, isWildcardHost func(string) bool) ResolverArtifacts {
	ipSetAll := map[string]struct{}{}
	ipSetDirect := map[string]struct{}{}
	ipSetWildcard := map[string]struct{}{}

	ipMapAll := map[string]map[string]struct{}{}
	ipMapDirect := map[string]map[string]struct{}{}
	ipMapWildcard := map[string]map[string]struct{}{}

	add := func(set map[string]struct{}, labels map[string]map[string]struct{}, ip, label string) {
		if ip == "" {
			return
		}
		set[ip] = struct{}{}
		m := labels[ip]
		if m == nil {
			m = map[string]struct{}{}
			labels[ip] = m
		}
		m[label] = struct{}{}
	}

	for host, ips := range resolved {
		wildcardHost := false
		if isWildcardHost != nil {
			wildcardHost = isWildcardHost(host)
		}
		for _, ip := range ips {
			add(ipSetAll, ipMapAll, ip, host)
			if wildcardHost {
				add(ipSetWildcard, ipMapWildcard, ip, host)
			} else {
				add(ipSetDirect, ipMapDirect, ip, host)
			}
		}
	}
	for ipEntry, labels := range staticLabels {
		for _, lbl := range labels {
			add(ipSetAll, ipMapAll, ipEntry, lbl)
			add(ipSetDirect, ipMapDirect, ipEntry, lbl)
		}
	}

	var out ResolverArtifacts

	appendMapPairs := func(dst *[][2]string, labelsByIP map[string]map[string]struct{}) {
		for ip := range labelsByIP {
			labels := labelsByIP[ip]
			for lbl := range labels {
				*dst = append(*dst, [2]string{ip, lbl})
			}
		}
		sort.Slice(*dst, func(i, j int) bool {
			if (*dst)[i][0] == (*dst)[j][0] {
				return (*dst)[i][1] < (*dst)[j][1]
			}
			return (*dst)[i][0] < (*dst)[j][0]
		})
	}
	appendIPs := func(dst *[]string, set map[string]struct{}) {
		for ip := range set {
			*dst = append(*dst, ip)
		}
		sort.Strings(*dst)
	}

	appendMapPairs(&out.IPMap, ipMapAll)
	appendMapPairs(&out.DirectIPMap, ipMapDirect)
	appendMapPairs(&out.WildcardIPMap, ipMapWildcard)
	appendIPs(&out.IPs, ipSetAll)
	appendIPs(&out.DirectIPs, ipSetDirect)
	appendIPs(&out.WildcardIPs, ipSetWildcard)

	return out
}