package az

import (
	"fmt"
	"slices"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	apicommon "github.com/tidbcloud/infra-provider/api/tidbcloud/common"
	"github.com/tidbcloud/infra-provider/pkg/namer"
)

// Group abstracts a TiDB group (default or ExtraGroup) for Offline/Recover AZ ops.
// Each AZ operation takes a Group and is idempotent for that group's resources.
type Group interface {
	ID() string
	IsDefault() bool
	// Components returns component types owned by this group.
	Components() []apicommon.ComponentType
	// ClusterLabel is the tidbcloud-cluster label value for this group's nodes.
	ClusterLabel() string
	// Namer selects infra resources (e.g. ASGs) belonging to this group.
	Namer() *namer.Namer
	// MatchesNode reports whether a node/pod with the given component and
	// cluster label belongs to this group.
	MatchesNode(ct apicommon.ComponentType, clusterLabel string) bool
}

type defaultGroup struct {
	id           string
	clusterLabel string
	namer        *namer.Namer
}

func (g *defaultGroup) ID() string           { return g.id }
func (g *defaultGroup) IsDefault() bool      { return true }
func (g *defaultGroup) ClusterLabel() string { return g.clusterLabel }
func (g *defaultGroup) Namer() *namer.Namer  { return g.namer }
func (g *defaultGroup) Components() []apicommon.ComponentType {
	return allTargetComponentTypes
}
func (g *defaultGroup) MatchesNode(ct apicommon.ComponentType, clusterLabel string) bool {
	switch ct {
	case apicommon.ComponentType_PD, apicommon.ComponentType_TIKV, apicommon.ComponentType_TIFLASH:
		return true
	case apicommon.ComponentType_TIDB, apicommon.ComponentType_TIPROXY:
		return clusterLabel == g.clusterLabel
	default:
		return false
	}
}

type extraGroup struct {
	id           string
	clusterLabel string
	namer        *namer.Namer
}

func (g *extraGroup) ID() string           { return g.id }
func (g *extraGroup) IsDefault() bool      { return false }
func (g *extraGroup) ClusterLabel() string { return g.clusterLabel }
func (g *extraGroup) Namer() *namer.Namer  { return g.namer }
func (g *extraGroup) Components() []apicommon.ComponentType {
	return sqlPlaneComponentTypes
}
func (g *extraGroup) MatchesNode(ct apicommon.ComponentType, clusterLabel string) bool {
	switch ct {
	case apicommon.ComponentType_TIDB, apicommon.ComponentType_TIPROXY:
		return clusterLabel == g.clusterLabel
	default:
		return false
	}
}

// resolveGroups builds the ordered group list for Offline/Recover.
// Empty groupIDs means default group first, then non-deleted ExtraGroups in
// cluster.Spec.ExtraGroups order.
func resolveGroups(cluster *cloudv1.Cluster, groupIDs []string) ([]Group, error) {
	defaultID := cluster.GetDefaultGroupID()
	known := map[string]Group{}

	defaultG := &defaultGroup{
		id:           defaultID,
		clusterLabel: cluster.Name,
		namer:        namer.NewNamer().WithClusterName(cluster.Name).WithTiDBNamespace(cluster.PrefixedId()),
	}
	known[defaultID] = defaultG

	extras := make([]Group, 0, len(cluster.Spec.ExtraGroups))
	for i := range cluster.Spec.ExtraGroups {
		eg := &cluster.Spec.ExtraGroups[i]
		if eg.Deleted() {
			continue
		}
		g := &extraGroup{
			id:           eg.Name,
			clusterLabel: fmt.Sprintf("%s-%s", cluster.Name, eg.Name),
			namer: namer.NewNamer().
				WithTiDBNamespace(cluster.PrefixedId()).
				ForTiDBGroup(cluster, eg, ""),
		}
		known[eg.Name] = g
		extras = append(extras, g)
	}

	if len(groupIDs) == 0 {
		out := make([]Group, 0, 1+len(extras))
		out = append(out, defaultG)
		out = append(out, extras...)
		if len(out) == 0 {
			return nil, cloudv1.NewWFNonRetriableError(fmt.Errorf("no groups available on cluster"))
		}
		return out, nil
	}

	out := make([]Group, 0, len(groupIDs))
	seen := map[string]struct{}{}
	for _, id := range groupIDs {
		if id == "" {
			return nil, cloudv1.NewWFNonRetriableError(fmt.Errorf("groupIds must not contain empty string"))
		}
		if _, dup := seen[id]; dup {
			continue
		}
		g, ok := known[id]
		if !ok {
			return nil, cloudv1.NewWFNonRetriableError(fmt.Errorf("unknown group id %q", id))
		}
		seen[id] = struct{}{}
		out = append(out, g)
	}
	return out, nil
}

func findDefaultGroup(groups []Group) Group {
	for _, g := range groups {
		if g.IsDefault() {
			return g
		}
	}
	return nil
}

func groupIDsOf(groups []Group) []string {
	ids := make([]string, 0, len(groups))
	for _, g := range groups {
		ids = append(ids, g.ID())
	}
	return ids
}

func componentLabelAllowed(comp string, g Group) bool {
	var ct apicommon.ComponentType
	switch comp {
	case "pd":
		ct = apicommon.ComponentType_PD
	case "tikv":
		ct = apicommon.ComponentType_TIKV
	case "tidb":
		ct = apicommon.ComponentType_TIDB
	case "tiflash":
		ct = apicommon.ComponentType_TIFLASH
	case "tiproxy":
		ct = apicommon.ComponentType_TIPROXY
	default:
		return false
	}
	return slices.Contains(g.Components(), ct)
}