/*
 * Copyright 2021 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package lbcache

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/golang/mock/gomock"

	mocksloadbalance "github.com/cloudwego/kitex/internal/mocks/loadbalance"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/discovery"
	"github.com/cloudwego/kitex/pkg/loadbalance"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
)

var defaultOptions = Options{
	RefreshInterval: defaultRefreshInterval,
	ExpireInterval:  defaultExpireInterval,
}

func TestBuilder(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	ins := discovery.NewInstance("tcp", "127.0.0.1:8888", 10, nil)
	r := &discovery.SynthesizedResolver{
		ResolveFunc: func(ctx context.Context, key string) (discovery.Result, error) {
			return discovery.Result{Cacheable: true, CacheKey: key, Instances: []discovery.Instance{ins}}, nil
		},
		TargetFunc: func(ctx context.Context, target rpcinfo.EndpointInfo) string {
			return "mockRoute"
		},
		NameFunc: func() string { return t.Name() },
	}
	lb := mocksloadbalance.NewMockLoadbalancer(ctrl)
	lb.EXPECT().GetPicker(gomock.Any()).DoAndReturn(func(res discovery.Result) loadbalance.Picker {
		test.Assert(t, res.Cacheable)
		test.Assert(t, res.CacheKey == t.Name()+":mockRoute", res.CacheKey)
		test.Assert(t, len(res.Instances) == 1)
		test.Assert(t, res.Instances[0].Address().String() == "127.0.0.1:8888")
		picker := mocksloadbalance.NewMockPicker(ctrl)
		return picker
	}).AnyTimes()
	lb.EXPECT().Name().Return("Synthesized").AnyTimes()
	NewBalancerFactory(r, lb, Options{Cacheable: true})
	b, ok := balancerFactories.Load(cacheKey(t.Name(), "Synthesized", defaultOptions))
	test.Assert(t, ok)
	test.Assert(t, b != nil)
	bl, err := b.(*BalancerFactory).Get(context.Background(), nil)
	test.Assert(t, err == nil)
	test.Assert(t, bl.GetPicker() != nil)
	dump := Dump()
	dumpJson, err := json.Marshal(dump)
	test.Assert(t, err == nil)
	test.Assert(t, string(dumpJson) == `{"TestBuilder|Synthesized|{5s 15s}":{"mockRoute":[{"Address":"tcp://127.0.0.1:8888","Weight":10}]}}`)
}

func TestCacheKey(t *testing.T) {
	uniqueKey := cacheKey("hello", "world", Options{RefreshInterval: 15 * time.Second, ExpireInterval: 5 * time.Minute})
	test.Assert(t, uniqueKey == "hello|world|{15s 5m0s}")
}

func TestBalancerCache(t *testing.T) {
	count := 10
	inss := []discovery.Instance{}
	for i := 0; i < count; i++ {
		inss = append(inss, discovery.NewInstance("tcp", fmt.Sprint(i), 10, nil))
	}
	r := &discovery.SynthesizedResolver{
		TargetFunc: func(ctx context.Context, target rpcinfo.EndpointInfo) string {
			return target.ServiceName()
		},
		ResolveFunc: func(ctx context.Context, key string) (discovery.Result, error) {
			return discovery.Result{Cacheable: true, CacheKey: "svc", Instances: inss}, nil
		},
		NameFunc: func() string { return t.Name() },
	}
	lb := loadbalance.NewWeightedBalancer()
	for i := 0; i < count; i++ {
		blf := NewBalancerFactory(r, lb, Options{})
		info := rpcinfo.NewEndpointInfo("svc", "", nil, nil)
		b, err := blf.Get(context.Background(), info)
		test.Assert(t, err == nil)
		p := b.GetPicker()
		for a := 0; a < count; a++ {
			addr := p.Next(context.Background(), nil).Address().String()
			t.Logf("count: %d addr: %s\n", i, addr)
		}
	}
}

type mockRebalancer struct {
	rebalanceFunc func(ch discovery.Change)
	deleteFunc    func(ch discovery.Change)
}

// Rebalance implements the Rebalancer interface.
func (m *mockRebalancer) Rebalance(ch discovery.Change) {
	if m.rebalanceFunc != nil {
		m.rebalanceFunc(ch)
	}
}

// Delete implements the Rebalancer interface.
func (m *mockRebalancer) Delete(ch discovery.Change) {
	if m.deleteFunc != nil {
		m.deleteFunc(ch)
	}
}

type mockResolver struct{}

func (m *mockResolver) Target(ctx context.Context, target rpcinfo.EndpointInfo) (description string) {
	return "target"
}

func (m *mockResolver) Resolve(ctx context.Context, desc string) (discovery.Result, error) {
	return discovery.Result{}, nil
}

func (m *mockResolver) Diff(cacheKey string, prev, next discovery.Result) (discovery.Change, bool) {
	return discovery.Change{}, false
}

func (m *mockResolver) Name() string {
	return "name"
}

var _ discovery.Resolver = &mockResolver{}

func TestConcurrentGet(t *testing.T) {
	cacheOpts := Options{Cacheable: false, RefreshInterval: time.Second, ExpireInterval: 5 * time.Second}
	bf := newBalancerFactory(&mockResolver{}, loadbalance.NewWeightedBalancer(), cacheOpts)
	m := sync.Map{}
	wg := sync.WaitGroup{}

	// concurrent get
	for i := 0; i < 100; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			b, _ := bf.Get(context.Background(), nil)
			m.Store(b, b)
		}()
	}
	wg.Wait()

	// check if length == 1, (target -> balancer, 1:1)
	cnt := 0
	m.Range(func(key, value any) bool {
		cnt++
		return true
	})
	test.Assert(t, cnt == 1, cnt)
}
