package common

import (
	"fmt"

	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"
	import_model "github.com/tidbcloud/dataflow-service-e2e/internal/tidbcloud/premium/consoleapi/import/model"
)

// PremiumImportValidator provides validation helper methods for Premium Import E2E tests.
type PremiumImportValidator struct {
	lgr     log.SimpleLogger
	importW *consoleapi.PremiumImportClientWrapper
}

// NewPremiumImportValidator creates a new PremiumImportValidator instance.
func NewPremiumImportValidator(lgr log.SimpleLogger) *PremiumImportValidator {
	return &PremiumImportValidator{
		lgr:     lgr,
		importW: consoleapi.NewPremiumImportClientWrapper(lgr),
	}
}

// Client returns the underlying PremiumImportClientWrapper for direct API access.
func (v *PremiumImportValidator) Client() *consoleapi.PremiumImportClientWrapper {
	return v.importW
}

// IsImportInList checks if the import task appears in the import list.
func (v *PremiumImportValidator) IsImportInList(tidbCluster *TiDBPremiumEntity, importID string) error {
	pageToken := ""
	seenPageTokens := map[string]struct{}{}
	pagesChecked := 0
	totalSize := int32(0)

	for {
		var pageTokenPtr *string
		if pageToken != "" {
			pageTokenPtr = &pageToken
		}
		resp, err := v.importW.ListImportsWithPageToken(tidbCluster, 100, pageTokenPtr)
		if err != nil {
			return err
		}
		pagesChecked++
		totalSize = resp.TotalSize
		for _, imp := range resp.Imports {
			if imp == nil || imp.ImportID != importID {
				continue
			}
			var stateStr string
			if imp.State != nil {
				stateStr = string(*imp.State)
			}
			v.lgr.Info("import %s found in list, state: %s", importID, stateStr)
			return nil
		}

		if resp.NextPageToken == "" {
			break
		}
		if _, exists := seenPageTokens[resp.NextPageToken]; exists {
			return fmt.Errorf("detected repeated next_page_token while listing imports: %s", resp.NextPageToken)
		}
		seenPageTokens[resp.NextPageToken] = struct{}{}
		pageToken = resp.NextPageToken
	}

	return fmt.Errorf("import %s not found in list after %d pages (total: %d)", importID, pagesChecked, totalSize)
}

// IsImportCompleted checks if the import task has reached ENUM_COMPLETED state.
func (v *PremiumImportValidator) IsImportCompleted(tidbCluster *TiDBPremiumEntity, importID string) error {
	imp, err := v.importW.GetImport(tidbCluster, importID)
	if err != nil {
		return err
	}
	if _, err := evaluateImportState(importID, imp.State, imp.Message); err != nil {
		return err
	}
	v.lgr.Info("import %s completed, totalSize: %s", importID, imp.TotalSize)
	return nil
}

func evaluateImportState(importID string, state *import_model.V1beta2ImportStateEnum, message string) (bool, error) {
	if state == nil {
		return false, fmt.Errorf("import %s has nil state", importID)
	}
	s := *state
	switch s {
	case import_model.V1beta2ImportStateEnumENUMCOMPLETED:
		return true, nil
	case import_model.V1beta2ImportStateEnumENUMFAILED:
		return false, retry.Unrecoverable(fmt.Errorf("import %s failed: %s", importID, message))
	case import_model.V1beta2ImportStateEnumENUMCANCELED:
		return false, retry.Unrecoverable(fmt.Errorf("import %s was canceled", importID))
	default:
		return false, fmt.Errorf("import %s status: %s, waiting for %s",
			importID, s, import_model.V1beta2ImportStateEnumENUMCOMPLETED)
	}
}