package invoke

import (
	"fmt"
	"google.golang.org/grpc/resolver"
	"strings"
	"sync"
)

const (
	scheme = "bifrost"
)

type CustomBuilder struct{}

func (cb *CustomBuilder) Scheme() string {
	return scheme
}

func (cb *CustomBuilder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
	cr := &customResolver{
		cc:     cc,
		target: target,
	}

	cr.ResolveNow(resolver.ResolveNowOptions{})

	return cr, nil
}

type customResolver struct {
	sync.Mutex
	target resolver.Target
	cc     resolver.ClientConn
	ips    map[string]string
}

func (cr *customResolver) ResolveNow(o resolver.ResolveNowOptions) {
	var (
		addrs = make([]resolver.Address, 0)
		hp    []string
	)

	cr.Lock()
	defer cr.Unlock()

	if hp = strings.Split(cr.target.URL.Host, ":"); len(hp) >= 2 {
		if ip, ok := pool[hp[0]]; ok {
			addr := fmt.Sprintf("%s:%s", ip, hp[1])
			addrs = append(addrs, resolver.Address{Addr: addr})
		}
	}

	_ = cr.cc.UpdateState(resolver.State{Addresses: addrs})
}

func (cr *customResolver) Close() {}

var (
	cb   = &CustomBuilder{}
	pool = make(map[string]string)
)

func init() {
	resolver.Register(cb)
}

type CustomDomain struct {
	Domain string
	IP     string
}

func NewCustomBuilder(cds ...CustomDomain) resolver.Builder {
	locker.Lock()
	defer locker.Unlock()

	for _, cd := range cds {
		pool[cd.Domain] = cd.IP
	}

	return cb
}