package resolver

import (
	"common/db/etcd"
	"common/discover/common"
	"common/log"
	"context"
	"go.etcd.io/etcd/client/v3"
	"google.golang.org/grpc/resolver"
	"strings"
)

const etcdSchema = "etcd"

func init() {
	resolver.Register(&etcdBuilder{})
}

type etcdBuilder struct{}

func (*etcdBuilder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
	serviceName := strings.TrimPrefix(target.URL.Path, "/")
	r := &etcdResolver{
		serviceName: serviceName,
		cc:          cc,
	}
	r.start()
	return r, nil
}

func (*etcdBuilder) Scheme() string { return etcdSchema }

type etcdResolver struct {
	serviceName string
	cc          resolver.ClientConn
	ctx         context.Context
	cancel      context.CancelFunc
}

func (r *etcdResolver) start() {
	r.ctx, r.cancel = context.WithCancel(context.Background())
	r.update()
	go r.watch()
}

func (r *etcdResolver) getPrefix() string {
	return common.KeyDiscoverService + "/" + r.serviceName
}

func (r *etcdResolver) update() {
	resp, err := etcd.GetClient().Get(r.getPrefix(), clientv3.WithPrefix())
	if err != nil {
		log.Errorf("etcd resolver get error: %v", err)
		return
	}

	var addrArray []resolver.Address
	for _, kv := range resp.Kvs {
		addr := string(kv.Value)
		if addr != "" {
			addrArray = append(addrArray, resolver.Address{Addr: addr})
		}
	}

	_ = r.cc.UpdateState(resolver.State{Addresses: addrArray})
}

func (r *etcdResolver) watch() {
	watchCh := etcd.GetClient().Watch(r.getPrefix(), clientv3.WithPrefix())
	for w := range watchCh {
		if w.Err() != nil {
			continue
		}
		r.update()
	}
}

func (r *etcdResolver) ResolveNow(resolver.ResolveNowOptions) {
	r.update()
}

func (r *etcdResolver) Close() {
	r.cancel()
}