package server

import (
	"context"
	"fmt"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	apiv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/cluster-server/v1"
)

// replayRouteStrategy builds the concrete task materials for a replay by matching
// source audit-log units to target endpoints.
type replayRouteStrategy interface {
	BuildTaskMaterials(ctx context.Context) ([]replayTaskMaterial, error)
}

// newReplayRouteStrategy returns the appropriate strategy for the given source/target
// cluster type combination, validating that the pair is supported.
// All currently-supported combinations use the uniform roundRobinReplayRouteStrategy;
// the factory can be extended to return specialised strategies per combination in the future.
func newReplayRouteStrategy(
	source replayClusterProvider,
	target replayClusterProvider,
) (replayRouteStrategy, error) {
	if source == nil {
		return nil, fmt.Errorf("source provider is not set")
	}
	if target == nil {
		return nil, fmt.Errorf("target provider is not set")
	}
	if err := validateReplayRouteCombination(source.GetClusterType(), target.GetClusterType()); err != nil {
		return nil, err
	}
	return &roundRobinReplayRouteStrategy{source: source, target: target}, nil
}

func validateReplayRouteCombination(srcType, tgtType apiv1.Replay_ClusterType) error {
	switch {
	case srcType == apiv1.Replay_CLUSTER_TYPE_DEDICATED && tgtType == apiv1.Replay_CLUSTER_TYPE_DEDICATED,
		srcType == apiv1.Replay_CLUSTER_TYPE_DEDICATED && tgtType == apiv1.Replay_CLUSTER_TYPE_PREMIUM,
		srcType == apiv1.Replay_CLUSTER_TYPE_PREMIUM && tgtType == apiv1.Replay_CLUSTER_TYPE_DEDICATED,
		srcType == apiv1.Replay_CLUSTER_TYPE_PREMIUM && tgtType == apiv1.Replay_CLUSTER_TYPE_PREMIUM:
		return nil
	default:
		return unsupportedReplayRouteStrategy(srcType, tgtType)
	}
}

// roundRobinReplayRouteStrategy distributes source audit-log units uniformly across
// the available target endpoints using index modulo:
//
//	source unit i  →  endpoints[i % len(validEndpoints)]
//
// Multiple source units that land on the same endpoint have their audit-log addresses
// merged into a single task. Empty endpoints are filtered out before distribution so
// that the result only contains reachable targets.
//
// Example: sources=[A,B,C], valid endpoints=[x,y]
//
//	A→x, B→y, C→x  →  [{x,[A,C]}, {y,[B]}]
//
// Naming:
//   - When a single source maps 1:1 to an endpoint and the source has a name, that name
//     is preserved (maintains backward compatibility with the dedicated→dedicated case).
//   - Otherwise the task is named "group-{endpointIndex}" where endpointIndex is the
//     position in the original (unfiltered) endpoint slice.
type roundRobinReplayRouteStrategy struct {
	source replayClusterProvider
	target replayClusterProvider
}

func (s *roundRobinReplayRouteStrategy) BuildTaskMaterials(ctx context.Context) ([]replayTaskMaterial, error) {
	// Load target endpoints first for fast-fail when no valid targets exist.
	endpoints, err := s.target.LoadTargetEndpoints(ctx)
	if err != nil {
		return nil, err
	}

	type validEP struct {
		idx int
		ep  string
	}
	valid := make([]validEP, 0, len(endpoints))
	for i, ep := range endpoints {
		if ep != "" {
			valid = append(valid, validEP{idx: i, ep: ep})
		}
	}
	if len(valid) == 0 {
		return nil, fmt.Errorf("no valid target endpoints")
	}

	// Load source material after confirming at least one endpoint exists.
	sourceMaterial, err := s.source.LoadSourceMaterial(ctx)
	if err != nil {
		return nil, err
	}
	if len(sourceMaterial.units) == 0 {
		return nil, errPremiumReplayMissingSourceAuditLog
	}

	// Round-robin assignment: source unit i → valid[i % len(valid)].
	groups := make([][]int, len(valid))
	for i := range sourceMaterial.units {
		groups[i%len(valid)] = append(groups[i%len(valid)], i)
	}

	materials := make([]replayTaskMaterial, 0, len(valid))
	for gi, srcIdxs := range groups {
		if len(srcIdxs) == 0 {
			continue
		}

		addrs := make([]cloudv1.AuditLogAddr, 0)
		for _, si := range srcIdxs {
			addrs = append(addrs, sourceMaterial.units[si].auditLogAddrs...)
		}
		if len(addrs) == 0 {
			return nil, errPremiumReplayMissingSourceAuditLog
		}

		// Preserve the source unit name when exactly one source maps to this endpoint
		// and the unit already has a non-empty name (dedicated→dedicated compatibility).
		name := fmt.Sprintf("group-%d", valid[gi].idx)
		if len(srcIdxs) == 1 && sourceMaterial.units[srcIdxs[0]].name != "" {
			name = sourceMaterial.units[srcIdxs[0]].name
		}

		materials = append(materials, replayTaskMaterial{
			name:           name,
			targetEndpoint: valid[gi].ep,
			auditLogAddrs:  addrs,
		})
	}
	return materials, nil
}

// ----------------------------------------
// Data transfer objects (shared by strategy and task builder)
// ----------------------------------------

type replaySourceMaterial struct {
	units []replaySourceTaskUnit
}

type replaySourceTaskUnit struct {
	name          string
	auditLogAddrs []cloudv1.AuditLogAddr
}

type replayTaskMaterial struct {
	name           string
	targetEndpoint string
	auditLogAddrs  []cloudv1.AuditLogAddr
}
