package common

import (
	"github.com/go-logr/logr"
	opv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
)

// RestartPhaseID returns the phase for the annotation-trigger step of a restart phase.
func RestartPhaseID(ph RestartPhase) UpgradePhase {
	switch ph.Reason {
	case ReasonRestartingTiProxy:
		return PhaseRestartTiProxy
	case ReasonRestartingStorage:
		return PhaseRestartStorage
	case ReasonRestartingTiDB:
		return PhaseRestartTiDB
	default:
		return "restart-unknown"
	}
}

// RestartMigratePhaseID returns the phase for the pod-migration step of a restart phase.
func RestartMigratePhaseID(ph RestartPhase) UpgradePhase {
	switch ph.Reason {
	case ReasonRestartingTiProxy:
		return PhaseRestartTiProxyMigrate
	case ReasonRestartingStorage:
		return PhaseRestartStorageMigrate
	case ReasonRestartingTiDB:
		return PhaseRestartTiDBMigrate
	default:
		return "restart-unknown-migrate"
	}
}

// RestartReadyPhaseID returns the phase for the component-readiness step of a restart phase.
func RestartReadyPhaseID(ph RestartPhase) UpgradePhase {
	switch ph.Reason {
	case ReasonRestartingTiProxy:
		return PhaseRestartTiProxyReady
	case ReasonRestartingStorage:
		return PhaseRestartStorageReady
	case ReasonRestartingTiDB:
		return PhaseRestartTiDBReady
	default:
		return "restart-unknown-ready"
	}
}

// CountRestartPhaseReady returns how many (TC, component) pairs are restarted and ready.
func CountRestartPhaseReady(logger logr.Logger, allTCs []*opv1alpha1.TidbCluster, desiredVersion string, ph RestartPhase) (ready, total int) {
	for _, tc := range allTCs {
		for _, comp := range ph.Components {
			if TCComponentReplicas(tc, comp) == 0 {
				continue
			}
			total++
			if IsComponentRestartedAndReady(tc, desiredVersion, comp) {
				ready++
				logger.Info("component ready", "tc", tc.Name, "component", comp)
			} else {
				logger.Info("component not restarted and ready", "tc", tc.Name, "component", comp)
			}
		}
	}
	return ready, total
}

// CountTCsReady returns how many TiDB clusters report ready.
func CountTCsReady(allTCs []*opv1alpha1.TidbCluster, defaultTC *opv1alpha1.TidbCluster) (ready, total int) {
	total = len(allTCs)
	for _, tc := range allTCs {
		if !IsTCReady(tc) {
			continue
		}
		if defaultTC != nil && tc.Name == defaultTC.Name {
			if tc.Status.TiDB.StatefulSet == nil ||
				tc.Spec.TiDB.Replicas != tc.Status.TiDB.StatefulSet.ReadyReplicas {
				continue
			}
		}
		ready++
	}
	return ready, total
}

// ComponentMigrationCounter counts component pods migrated to new nodes for a single TC.
type ComponentMigrationCounter func(namespace, component string, tc *opv1alpha1.TidbCluster) (onNew, total int, err error)

// CountRestartPhaseMigration aggregates migration progress across all TCs for a restart phase.
// migratable is true when at least one non-TiDB component with replicas > 0 was checked.
// When migratable is true but total is 0, no matching pods were found yet (abnormal — caller should wait).
func CountRestartPhaseMigration(logger logr.Logger, namespace string, allTCs []*opv1alpha1.TidbCluster, ph RestartPhase, counter ComponentMigrationCounter) (done, total int, migratable bool, err error) {
	for _, tc := range allTCs {
		for _, comp := range ph.Components {
			if comp == cloudv1.ComponentTiDB {
				continue
			}
			if TCComponentReplicas(tc, comp) == 0 {
				continue
			}
			migratable = true
			onNew, compTotal, countErr := counter(namespace, string(comp), tc)
			if countErr != nil {
				logger.Error(countErr, "unable to count component migration", "tc", tc.Name, "component", comp)
				return done, total, migratable, countErr
			}
			if compTotal == 0 {
				logger.Info("no pods found for migratable component", "tc", tc.Name, "component", comp,
					"replicas", TCComponentReplicas(tc, comp))
				continue
			}
			done += onNew
			total += compTotal
			logger.Info("component migration progress", "tc", tc.Name, "component", comp, "onNew", onNew, "compTotal", compTotal)
		}
	}
	return done, total, migratable, nil
}