package common

import (
	"fmt"
	"strings"
	"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/dedicated/consoleapi"
	dm_model "github.com/tidbcloud/dataflow-service-e2e/internal/tidbcloud/dedicated/consoleapi/dm/model"
)

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

// NewDMValidator creates a new DMValidator instance.
func NewDMValidator(lgr log.SimpleLogger) *DMValidator {
	return &DMValidator{
		lgr: lgr,
		dmW: consoleapi.NewDMClientWrapper(lgr),
	}
}

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

// GetFirstJobSpec retrieves the first available DM job spec.
func (v *DMValidator) GetFirstJobSpec() (*dm_model.DmJobSpec, error) {
	resp, err := v.dmW.ListJobSpecs()
	if err != nil {
		return nil, err
	}
	if len(resp.Payload.Specs) == 0 {
		return nil, fmt.Errorf("no DM job specs available")
	}
	spec := resp.Payload.Specs[0]
	v.lgr.Info("using DM job spec: id=%s, name=%s, rcu=%.0f", spec.ID, spec.DisplayName, spec.CapacityUnit)
	return spec, nil
}

// GetSecondJobSpec retrieves the second available DM job spec (for scaling tests).
// Returns an error if fewer than 2 specs are available.
func (v *DMValidator) GetSecondJobSpec() (*dm_model.DmJobSpec, error) {
	resp, err := v.dmW.ListJobSpecs()
	if err != nil {
		return nil, err
	}
	if len(resp.Payload.Specs) < 2 {
		return nil, fmt.Errorf("need at least 2 DM job specs for scaling, got %d", len(resp.Payload.Specs))
	}
	spec := resp.Payload.Specs[1]
	v.lgr.Info("using second DM job spec: id=%s, name=%s, rcu=%.0f", spec.ID, spec.DisplayName, spec.CapacityUnit)
	return spec, nil
}

// getJobStage returns the aggregated stage of a DM job based on its sub-tasks.
// DM jobs expose status via sub-tasks; a job's overall stage is determined by its sub-tasks.
func (v *DMValidator) getJobStage(tidbCluster *TiDBDedicatedEntity, jobID string) (dm_model.DmSubTaskStage, error) {
	job, err := v.dmW.GetJob(tidbCluster, jobID)
	if err != nil {
		return "", err
	}
	if len(job.SubTasks) == 0 {
		return dm_model.DmSubTaskStageCreating, nil
	}
	// Return the stage of the first sub-task (single-source is most common)
	for _, st := range job.SubTasks {
		if st.Stage != nil {
			return *st.Stage, nil
		}
		return dm_model.DmSubTaskStageUnknown, nil
	}
	return dm_model.DmSubTaskStageUnknown, nil
}

// IsJobInStage checks if a DM job is in the expected stage.
func (v *DMValidator) IsJobInStage(tidbCluster *TiDBDedicatedEntity, jobID string, expected dm_model.DmSubTaskStage) error {
	stage, err := v.getJobStage(tidbCluster, jobID)
	if err != nil {
		return err
	}
	if stage == dm_model.DmSubTaskStageFailed && expected != dm_model.DmSubTaskStageFailed {
		return retry.Unrecoverable(fmt.Errorf("job %s failed: %s", jobID, v.extractJobErrorMessage(tidbCluster, jobID)))
	}
	if stage != expected {
		return fmt.Errorf("job %s stage is %s, expected %s", jobID, stage, expected)
	}
	v.lgr.Info("job %s is in expected stage: %s", jobID, expected)
	return nil
}

// IsJobRunning checks if a DM job is in Running stage.
func (v *DMValidator) IsJobRunning(tidbCluster *TiDBDedicatedEntity, jobID string) error {
	return v.IsJobInStage(tidbCluster, jobID, dm_model.DmSubTaskStageRunning)
}

// IsJobPaused checks if a DM job is in Paused stage.
func (v *DMValidator) IsJobPaused(tidbCluster *TiDBDedicatedEntity, jobID string) error {
	return v.IsJobInStage(tidbCluster, jobID, dm_model.DmSubTaskStagePaused)
}

// IsJobFinished checks if a DM job is in Finished stage.
func (v *DMValidator) IsJobFinished(tidbCluster *TiDBDedicatedEntity, jobID string) error {
	return v.IsJobInStage(tidbCluster, jobID, dm_model.DmSubTaskStageFinished)
}

// IsJobDeleted checks if a DM job no longer exists in the job list.
func (v *DMValidator) IsJobDeleted(tidbCluster *TiDBDedicatedEntity, jobID string) error {
	resp, err := v.dmW.ListJobs(tidbCluster, 1, 100)
	if err != nil {
		return err
	}
	for _, job := range resp.Jobs {
		if job.ID == jobID {
			return fmt.Errorf("job %s still exists (stage may be in sub-tasks)", jobID)
		}
	}
	v.lgr.Info("job %s is deleted (not found in list)", jobID)
	return nil
}

// WaitForJobStage waits for a DM job to reach the expected stage with polling.
func (v *DMValidator) WaitForJobStage(
	tidbCluster *TiDBDedicatedEntity, jobID string, expected dm_model.DmSubTaskStage, timeout time.Duration,
) error {
	deadline := time.Now().Add(timeout)
	interval := 30 * time.Second

	for time.Now().Before(deadline) {
		stage, err := v.getJobStage(tidbCluster, jobID)
		if err != nil {
			v.lgr.Warn("failed to get job stage: %s", err.Error())
			time.Sleep(interval)
			continue
		}

		if stage == expected {
			v.lgr.Info("job %s reached expected stage: %s", jobID, expected)
			return nil
		}

		if stage == dm_model.DmSubTaskStageFailed {
			return fmt.Errorf("job %s failed: %s", jobID, v.extractJobErrorMessage(tidbCluster, jobID))
		}

		v.lgr.Info("job %s stage: %s, waiting for: %s, remaining: %s",
			jobID, stage, expected, time.Until(deadline).Round(time.Second))
		time.Sleep(interval)
	}

	return fmt.Errorf("timeout waiting for job %s to reach stage %s", jobID, expected)
}

func (v *DMValidator) extractJobErrorMessage(tidbCluster *TiDBDedicatedEntity, jobID string) string {
	job, err := v.dmW.GetJob(tidbCluster, jobID)
	if err != nil || job == nil {
		return "unknown error"
	}
	errMsg := "unknown error"
	for _, st := range job.SubTasks {
		if st.ProcessError != nil && st.ProcessError.Message != "" {
			errMsg = st.ProcessError.Message
		}
	}
	return errMsg
}

// IsPrecheckDone checks if a precheck has completed successfully.
func (v *DMValidator) IsPrecheckDone(tidbCluster *TiDBDedicatedEntity, precheckID string) error {
	resp, err := v.dmW.GetPrecheck(tidbCluster, precheckID)
	if err != nil {
		return err
	}
	if err := validatePrecheckResult(precheckID, resp); err != nil {
		return err
	}
	v.lgr.Info("precheck %s done, status: %s, total: %d, success: %d, failed: %d, warn: %d",
		precheckID, resp.Status, resp.Total, resp.SuccessCnt, resp.FailedCnt, resp.WarnCnt)
	return nil
}

func validatePrecheckResult(precheckID string, resp *dm_model.DmGetPrecheckJobResp) error {
	if resp == nil {
		return fmt.Errorf("precheck %s returned an empty response", precheckID)
	}

	summary := fmt.Sprintf("total=%d, success=%d, failed=%d, warn=%d",
		resp.Total, resp.SuccessCnt, resp.FailedCnt, resp.WarnCnt)
	switch strings.ToLower(resp.Status) {
	case "finished":
		if resp.FailedCnt > 0 {
			return retry.Unrecoverable(fmt.Errorf("precheck %s finished with failed checks: %s", precheckID, summary))
		}
		return nil
	case "failed":
		return retry.Unrecoverable(fmt.Errorf("precheck %s failed: %s", precheckID, summary))
	case "running":
		return fmt.Errorf("precheck %s still running", precheckID)
	default:
		return fmt.Errorf("precheck %s not finished yet, current status: %q", precheckID, resp.Status)
	}
}
