package grpcutil

import (
	"context"
	"sync"
	"sync/atomic"

	"github.com/pingcap/log"
	"github.com/tidbcloud/ffs/config"
	"google.golang.org/grpc"
	"google.golang.org/grpc/resolver"
)

type ServiceClient struct {
	client   interface{}
	ctx      context.Context
	cancel   context.CancelFunc
	conn     *grpc.ClientConn
	resolver grpcResolver
	stopped  int32

	Schema    string
	Target    string
	Endpoints []string
	Options   []grpc.DialOption
	NewClient func(*grpc.ClientConn) interface{}
}

func (c *ServiceClient) Start() error {
	c.ctx, c.cancel = context.WithCancel(context.Background())
	if len(c.Target) == 0 {
		c.resolver.schema = c.Schema
		c.resolver.update(c.Endpoints...)
		resolver.Register(&c.resolver)
		c.Target = c.resolver.Scheme() + ":///service/"
		c.Options = []grpc.DialOption{
			grpc.WithBalancerName("round_robin"),
			grpc.WithInsecure(),
			grpc.WithDefaultCallOptions(
				grpc.MaxCallRecvMsgSize(config.GrpcMaxMsgSize),
				grpc.MaxCallSendMsgSize(config.GrpcMaxMsgSize),
			)}
	}
	return c.dial()
}

func (c *ServiceClient) dial() error {
	conn, err := grpc.Dial(
		c.Target,
		c.Options...,
	)
	if err != nil {
		return err
	}
	c.conn = conn
	c.client = c.NewClient(c.conn)
	return nil
}
func (c *ServiceClient) Stop() {
	atomic.StoreInt32(&c.stopped, 1)
	c.cancel()
	if conn := c.conn; conn != nil {
		conn.Close()
	}
}

func (c *ServiceClient) Client() interface{} {
	return c.client
}

func (c *ServiceClient) Update(endpoints ...string) {
	c.resolver.update(endpoints...)
}

type grpcResolver struct {
	schema string
	mu     sync.Mutex
	cc     resolver.ClientConn
	state  resolver.State
}

func (r *grpcResolver) Scheme() string {
	return r.schema
}

func (r *grpcResolver) ResolveNow(rn resolver.ResolveNowOptions) {

}

func (r *grpcResolver) Close() {

}

func (r *grpcResolver) Build(target resolver.Target, clientConn resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
	r.mu.Lock()
	r.cc = clientConn
	r.cc.UpdateState(r.state)
	r.mu.Unlock()
	log.S().Infof("resolver build target %v", target)
	return r, nil
}

func (r *grpcResolver) update(endpoints ...string) {
	r.mu.Lock()
	r.state = resolver.State{}
	for i := range endpoints {
		r.state.Addresses = append(r.state.Addresses, resolver.Address{Addr: endpoints[i]})
	}
	if r.cc != nil {
		r.cc.UpdateState(r.state)
	}
	r.mu.Unlock()
}