package common

import (
	"fmt"
	"time"

	retry "github.com/avast/retry-go/v4"

	"github.com/tidbcloud/dataflow-service-e2e/internal/log"
	. "github.com/tidbcloud/dataflow-service-e2e/internal/model"
	"github.com/tidbcloud/dataflow-service-e2e/internal/tidbcloud/premium/consoleapi"
	dm_model "github.com/tidbcloud/dataflow-service-e2e/internal/tidbcloud/premium/consoleapi/dm/model"
)

// PremiumDMValidator provides validation helper methods for Premium DM E2E tests.
type PremiumDMValidator struct {
	lgr log.SimpleLogger
	dmW *consoleapi.PremiumDMClientWrapper
}

// NewPremiumDMValidator creates a new PremiumDMValidator instance.
func NewPremiumDMValidator(lgr log.SimpleLogger) *PremiumDMValidator {
	return &PremiumDMValidator{
		lgr: lgr,
		dmW: consoleapi.NewPremiumDMClientWrapper(lgr),
	}
}

// Client returns the underlying PremiumDMClientWrapper for direct API access.
func (v *PremiumDMValidator) Client() *consoleapi.PremiumDMClientWrapper {
	return v.dmW
}

// getSubTaskStage returns the stage of the first sub-task of a premium DM migration.
func (v *PremiumDMValidator) getSubTaskStage(tidbCluster *TiDBPremiumEntity, migrationID string) (dm_model.PremiumDMSubTaskStage, error) {
	migration, err := v.dmW.GetMigration(tidbCluster, migrationID)
	if err != nil {
		return "", err
	}
	if len(migration.SubTasks) == 0 || migration.SubTasks[0].Stage == nil {
		return dm_model.PremiumDMSubTaskStageSTAGEUNSPECIFIED, nil
	}
	return *migration.SubTasks[0].Stage, nil
}

// IsMigrationInSubTaskStage checks if a premium DM migration's first sub-task is in the expected stage.
// It returns retry.Unrecoverable if the migration or sub-task is in FAILED state,
// so callers using Expect can fail fast instead of retrying until timeout.
func (v *PremiumDMValidator) IsMigrationInSubTaskStage(tidbCluster *TiDBPremiumEntity, migrationID string, expected dm_model.PremiumDMSubTaskStage) error {
	migration, err := v.dmW.GetMigration(tidbCluster, migrationID)
	if err != nil {
		return err
	}

	// Check migration-level state for FAILED first
	if migration.State != nil && *migration.State == dm_model.PremiumMigrationStateFAILED {
		errMsg := v.extractSubTaskErrors(migration)
		return retry.Unrecoverable(fmt.Errorf("migration %s is in FAILED state: %s", migrationID, errMsg))
	}

	stage := dm_model.PremiumDMSubTaskStageSTAGEUNSPECIFIED
	if len(migration.SubTasks) > 0 && migration.SubTasks[0].Stage != nil {
		stage = *migration.SubTasks[0].Stage
	}

	// Check sub-task level FAILED
	if stage == dm_model.PremiumDMSubTaskStageFAILED && expected != dm_model.PremiumDMSubTaskStageFAILED {
		errMsg := v.extractSubTaskErrors(migration)
		return retry.Unrecoverable(fmt.Errorf("migration %s sub-task FAILED: %s", migrationID, errMsg))
	}

	if stage != expected {
		var stateStr string
		if migration.State != nil {
			stateStr = string(*migration.State)
		}
		return fmt.Errorf("migration %s sub-task stage is %s, expected %s (state: %s)",
			migrationID, stage, expected, stateStr)
	}
	v.lgr.Info("migration %s is in expected sub-task stage: %s", migrationID, expected)
	return nil
}

// extractSubTaskErrors extracts error messages from sub-tasks of a migration.
func (v *PremiumDMValidator) extractSubTaskErrors(migration *dm_model.V1beta2PremiumMigration) string {
	for _, st := range migration.SubTasks {
		if st.ErrorMsg != "" {
			return st.ErrorMsg
		}
	}
	return "unknown error"
}

// IsMigrationRunning checks if a premium DM migration's sub-task is in Running stage.
func (v *PremiumDMValidator) IsMigrationRunning(tidbCluster *TiDBPremiumEntity, migrationID string) error {
	return v.IsMigrationInSubTaskStage(tidbCluster, migrationID, dm_model.PremiumDMSubTaskStageRUNNING)
}

// IsMigrationPaused checks if a premium DM migration's sub-task is in Paused stage.
func (v *PremiumDMValidator) IsMigrationPaused(tidbCluster *TiDBPremiumEntity, migrationID string) error {
	return v.IsMigrationInSubTaskStage(tidbCluster, migrationID, dm_model.PremiumDMSubTaskStagePAUSED)
}

// IsMigrationFinished checks if a premium DM migration's sub-task is in Finished stage.
func (v *PremiumDMValidator) IsMigrationFinished(tidbCluster *TiDBPremiumEntity, migrationID string) error {
	return v.IsMigrationInSubTaskStage(tidbCluster, migrationID, dm_model.PremiumDMSubTaskStageFINISHED)
}

// IsMigrationDeleted checks if a premium DM migration no longer exists in the migration list.
func (v *PremiumDMValidator) IsMigrationDeleted(tidbCluster *TiDBPremiumEntity, migrationID string) error {
	resp, err := v.dmW.ListMigrations(tidbCluster, 100)
	if err != nil {
		return err
	}
	for _, m := range resp.PremiumMigrations {
		if m.MigrationID == migrationID {
			stateStr := "<nil>"
			if m.State != nil {
				stateStr = string(*m.State)
			}
			return fmt.Errorf("migration %s still exists (state: %s)", migrationID, stateStr)
		}
	}
	v.lgr.Info("migration %s is deleted (not found in list)", migrationID)
	return nil
}

// WaitForSubTaskStage waits for a premium DM migration's sub-task to reach the expected stage with polling.
func (v *PremiumDMValidator) WaitForSubTaskStage(
	tidbCluster *TiDBPremiumEntity, migrationID string, expected dm_model.PremiumDMSubTaskStage, timeout time.Duration,
) error {
	deadline := time.Now().Add(timeout)
	interval := 30 * time.Second

	for time.Now().Before(deadline) {
		stage, err := v.getSubTaskStage(tidbCluster, migrationID)
		if err != nil {
			v.lgr.Warn("failed to get migration sub-task stage: %s", err.Error())
			time.Sleep(interval)
			continue
		}

		if stage == expected {
			v.lgr.Info("migration %s reached expected sub-task stage: %s", migrationID, expected)
			return nil
		}

		if stage == dm_model.PremiumDMSubTaskStageFAILED {
			migration, _ := v.dmW.GetMigration(tidbCluster, migrationID)
			errMsg := "unknown error"
			if migration != nil {
				for _, st := range migration.SubTasks {
					if st.ErrorMsg != "" {
						errMsg = st.ErrorMsg
					}
				}
			}
			return fmt.Errorf("migration %s failed: %s", migrationID, errMsg)
		}

		v.lgr.Info("migration %s sub-task stage: %s, waiting for: %s, remaining: %s",
			migrationID, stage, expected, time.Until(deadline).Round(time.Second))
		time.Sleep(interval)
	}

	return fmt.Errorf("timeout waiting for migration %s to reach sub-task stage %s", migrationID, expected)
}

// IsMigrationInSyncStep checks if a premium DM migration's first sub-task has reached the SYNC step.
// For ALL-mode (Full + Incremental) migrations, this indicates the full data migration (dump + load)
// is complete and incremental sync has started. The migration stays RUNNING indefinitely in SYNC step.
func (v *PremiumDMValidator) IsMigrationInSyncStep(tidbCluster *TiDBPremiumEntity, migrationID string) error {
	migration, err := v.dmW.GetMigration(tidbCluster, migrationID)
	if err != nil {
		return err
	}
	if migration.State != nil && *migration.State == dm_model.PremiumMigrationStateFAILED {
		errMsg := v.extractSubTaskErrors(migration)
		return retry.Unrecoverable(fmt.Errorf("migration %s is in FAILED state: %s", migrationID, errMsg))
	}
	if len(migration.SubTasks) == 0 {
		return fmt.Errorf("migration %s has no sub-tasks yet", migrationID)
	}
	step := migration.SubTasks[0].CurrentStep
	if step != nil && *step == dm_model.PremiumDMSubTaskStepSYNC {
		v.lgr.Info("migration %s full data migration complete, now in SYNC (incremental) step", migrationID)
		return nil
	}
	var stepStr string
	if step != nil {
		stepStr = string(*step)
	}
	return fmt.Errorf("migration %s current step is %s, waiting for SYNC", migrationID, stepStr)
}

// IsPrecheckDone checks if a precheck has completed (status is FINISHED or FAILED).
func (v *PremiumDMValidator) IsPrecheckDone(tidbCluster *TiDBPremiumEntity, precheckID string) error {
	resp, err := v.dmW.GetPrecheck(tidbCluster, precheckID)
	if err != nil {
		return err
	}
	if resp.Status == nil {
		return fmt.Errorf("precheck %s status is nil", precheckID)
	}
	status := *resp.Status
	switch status {
	case dm_model.V1beta2PremiumMigrationPrecheckStatusFINISHED, dm_model.V1beta2PremiumMigrationPrecheckStatusFAILED:
		v.lgr.Info("precheck %s done, status: %s, total: %d, success: %d, failed: %d, warn: %d",
			precheckID, status, resp.Total, resp.SuccessCnt, resp.FailedCnt, resp.WarnCnt)
		return nil
	default:
		return fmt.Errorf("precheck %s not finished yet, current status: %s", precheckID, status)
	}
}