package manager

import (
	"context"
	"sync"

	"golang.org/x/sync/errgroup"

	"github.com/milvus-io/milvus/internal/util/streamingutil/service/balancer/picker"
	"github.com/milvus-io/milvus/internal/util/streamingutil/service/contextutil"
	"github.com/milvus-io/milvus/internal/util/streamingutil/service/discoverer"
	"github.com/milvus-io/milvus/internal/util/streamingutil/service/lazygrpc"
	"github.com/milvus-io/milvus/internal/util/streamingutil/service/resolver"
	"github.com/milvus-io/milvus/internal/util/streamingutil/status"
	"github.com/milvus-io/milvus/pkg/v3/common"
	"github.com/milvus-io/milvus/pkg/v3/mlog"
	"github.com/milvus-io/milvus/pkg/v3/proto/streamingpb"
	"github.com/milvus-io/milvus/pkg/v3/streaming/util/types"
	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)

var _ ManagerClient = (*managerClientImpl)(nil)

// managerClientImpl implements ManagerClient.
type managerClientImpl struct {
	lifetime *typeutil.Lifetime
	stopped  chan struct{}

	rb      resolver.Builder
	service lazygrpc.Service[streamingpb.StreamingNodeManagerServiceClient]
}

func (c *managerClientImpl) WatchNodeChanged(ctx context.Context) (<-chan struct{}, error) {
	if !c.lifetime.Add(typeutil.LifetimeStateWorking) {
		return nil, status.NewOnShutdownError("manager client is closing")
	}
	defer c.lifetime.Done()

	resultCh := make(chan struct{}, 1)
	go func() {
		defer close(resultCh)
		c.rb.Resolver().Watch(ctx, func(state resolver.VersionedState) error {
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-c.stopped:
				return status.NewOnShutdownError("manager client is closing")
			case resultCh <- struct{}{}:
			}
			return nil
		})
	}()
	return resultCh, nil
}

// GetAllStreamingNodes fetches all streaming node info with resource group.
func (c *managerClientImpl) GetAllStreamingNodes(ctx context.Context) (map[int64]*types.StreamingNodeInfoWithResourceGroup, error) {
	if !c.lifetime.Add(typeutil.LifetimeStateWorking) {
		return nil, status.NewOnShutdownError("manager client is closing")
	}
	defer c.lifetime.Done()

	// Get all discovered streamingnode.
	state, err := c.rb.Resolver().GetLatestState(ctx)
	if err != nil {
		return nil, err
	}
	result := make(map[int64]*types.StreamingNodeInfoWithResourceGroup, len(state.State.Addresses))
	for serverID, session := range state.Sessions() {
		rg := session.GetResourceGroupName()
		if rg == "" {
			rg = common.DefaultResourceGroupName
		}
		result[serverID] = &types.StreamingNodeInfoWithResourceGroup{
			StreamingNodeInfo: types.StreamingNodeInfo{
				ServerID: serverID,
				Address:  session.Address,
			},
			ResourceGroup: rg,
		}
	}
	return result, nil
}

// CollectAllStatus collects status in underlying streamingnodes.
// If resourceGroupHint has discovered nodes, only nodes in that resource group are collected.
// Otherwise, it falls back to another discovered resource group.
func (c *managerClientImpl) CollectAllStatus(ctx context.Context, resourceGroupHint string) (map[int64]*types.StreamingNodeStatus, error) {
	if !c.lifetime.Add(typeutil.LifetimeStateWorking) {
		return nil, status.NewOnShutdownError("manager client is closing")
	}
	defer c.lifetime.Done()

	// Get all discovered streamingnode.
	state, err := c.rb.Resolver().GetLatestState(ctx)
	if err != nil {
		return nil, err
	}
	if len(state.State.Addresses) == 0 {
		return make(map[int64]*types.StreamingNodeStatus), nil
	}
	resourceGroupHint = filterStreamingNodeStatusByResourceGroupHint(state, resourceGroupHint)

	// Collect status from the selected resource group only. The hint fallback
	// is decided from service discovery before RPC to avoid broadcasting to
	// resource groups that will not be used for WAL balance.
	result, err := c.getAllStreamingNodeStatus(ctx, state, resourceGroupHint)
	if err != nil {
		return nil, err
	}

	// Collect status may cost some time, so we need to check the lifetime again.
	newState, err := c.rb.Resolver().GetLatestState(ctx)
	if err != nil {
		return nil, err
	}
	if newState.Version.GT(state.Version) {
		newSession := newState.Sessions()
		for serverID := range result {
			if session, ok := newSession[serverID]; !ok {
				result[serverID].Err = types.ErrNotAlive
			} else if session.Stopping {
				result[serverID].Err = types.ErrStopping
			}
		}
	}
	return result, nil
}

func filterStreamingNodeStatusByResourceGroupHint(state discoverer.VersionedState, resourceGroupHint string) string {
	resourceGroupHint = effectiveResourceGroupHintForCollectAllStatus(state, resourceGroupHint)
	if resourceGroupHint == "" {
		return ""
	}
	statsByRG := groupStreamingNodeSessionStatsByResourceGroup(state)
	if _, ok := statsByRG[resourceGroupHint]; ok {
		return resourceGroupHint
	}
	if fallbackRG, ok := chooseFallbackResourceGroup(statsByRG); ok {
		return fallbackRG
	}
	return ""
}

func effectiveResourceGroupHintForCollectAllStatus(state discoverer.VersionedState, resourceGroupHint string) string {
	if resourceGroupHint != "" {
		return resourceGroupHint
	}
	for _, session := range state.Sessions() {
		rg := session.GetResourceGroupName()
		if rg == "" || rg == common.DefaultResourceGroupName {
			return common.DefaultResourceGroupName
		}
	}
	return ""
}

type resourceGroupSessionStats struct {
	count       int
	maxServerID int64
}

func groupStreamingNodeSessionStatsByResourceGroup(state discoverer.VersionedState) map[string]resourceGroupSessionStats {
	statsByRG := make(map[string]resourceGroupSessionStats)
	for serverID, session := range state.Sessions() {
		rg := session.GetResourceGroupName()
		if rg == "" {
			rg = common.DefaultResourceGroupName
		}
		stats := statsByRG[rg]
		stats.count++
		if serverID > stats.maxServerID {
			stats.maxServerID = serverID
		}
		statsByRG[rg] = stats
	}
	return statsByRG
}

func chooseFallbackResourceGroup(statsByRG map[string]resourceGroupSessionStats) (string, bool) {
	var selectedRG string
	var selectedCount int
	var selectedMaxServerID int64
	for rg, stats := range statsByRG {
		if stats.count == 0 {
			continue
		}
		if selectedRG == "" || stats.count > selectedCount || (stats.count == selectedCount && stats.maxServerID > selectedMaxServerID) {
			selectedRG = rg
			selectedCount = stats.count
			selectedMaxServerID = stats.maxServerID
		}
	}
	return selectedRG, selectedRG != ""
}

func (c *managerClientImpl) getAllStreamingNodeStatus(ctx context.Context, state discoverer.VersionedState, resourceGroup string) (map[int64]*types.StreamingNodeStatus, error) {
	// wait for manager service ready.
	manager, err := c.service.GetService(ctx)
	if err != nil {
		return nil, err
	}

	g, _ := errgroup.WithContext(ctx)
	g.SetLimit(16)
	var mu sync.Mutex
	result := make(map[int64]*types.StreamingNodeStatus, len(state.Sessions()))
	for serverID, session := range state.Sessions() {
		serverID := serverID
		address := session.Address
		rg := session.GetResourceGroupName()
		if rg == "" {
			rg = common.DefaultResourceGroupName
		}
		if resourceGroup != "" && rg != resourceGroup {
			continue
		}
		g.Go(func() error {
			ctx := contextutil.WithPickServerID(ctx, serverID)
			resp, err := manager.CollectStatus(ctx, &streamingpb.StreamingNodeManagerCollectStatusRequest{})
			mu.Lock()
			defer mu.Unlock()

			if err != nil {
				mlog.Warn(ctx, "collect status failed, skip", mlog.Int64("serverID", serverID), mlog.Err(err))
				return err
			}
			result[serverID] = &types.StreamingNodeStatus{
				StreamingNodeInfo: types.StreamingNodeInfo{
					ServerID: serverID,
					Address:  address,
				},
				Metrics: types.NewStreamingNodeBalanceAttrsFromProto(resp.Metrics),
				Err:     err,
			}
			mlog.Debug(ctx, "collect status success", mlog.Int64("serverID", serverID), mlog.Any("status", resp))
			return nil
		})
	}
	g.Wait()
	return result, nil
}

// Assign a wal instance for the channel on log node of given server id.
func (c *managerClientImpl) Assign(ctx context.Context, pchannel types.PChannelInfoAssigned) error {
	if !c.lifetime.Add(typeutil.LifetimeStateWorking) {
		return status.NewOnShutdownError("manager client is closing")
	}
	defer c.lifetime.Done()

	// wait for manager service ready.
	manager, err := c.service.GetService(ctx)
	if err != nil {
		return err
	}

	// Select a log node to assign the wal instance.
	ctx = contextutil.WithPickServerID(ctx, pchannel.Node.ServerID)
	_, err = manager.Assign(ctx, &streamingpb.StreamingNodeManagerAssignRequest{
		Pchannel: types.NewProtoFromPChannelInfo(pchannel.Channel),
	})
	return err
}

// Remove the wal instance for the channel on log node of given server id.
func (c *managerClientImpl) Remove(ctx context.Context, pchannel types.PChannelInfoAssigned) error {
	if !c.lifetime.Add(typeutil.LifetimeStateWorking) {
		return status.NewOnShutdownError("manager client is closing")
	}
	defer c.lifetime.Done()

	// wait for manager service ready.
	manager, err := c.service.GetService(ctx)
	if err != nil {
		return err
	}

	// Select a streaming node to remove the wal instance.
	ctx = contextutil.WithPickServerID(ctx, pchannel.Node.ServerID)
	_, err = manager.Remove(ctx, &streamingpb.StreamingNodeManagerRemoveRequest{
		Pchannel: types.NewProtoFromPChannelInfo(pchannel.Channel),
	})
	// The following error can be treated as success.
	// 1. err is nil, a real remove operation at streaming node has been happened.
	// 2. err is ErrSubConnNoExist, the streaming node is not alive at view of session, so the wal on it is already removed.
	// 3. err is SkippedOperation, the streaming node is not the owner of the wal, so the wal on it is already removed.
	if err == nil || picker.IsErrSubConnNoExist(err) {
		return nil
	}
	statusErr := status.AsStreamingError(err)
	if statusErr == nil || statusErr.IsSkippedOperation() {
		return nil
	}
	return err
}

// Close closes the manager client.
func (c *managerClientImpl) Close() {
	c.lifetime.SetState(typeutil.LifetimeStateStopped)
	close(c.stopped)
	c.lifetime.Wait()

	c.service.Close()
	c.rb.Close()
}