package dm

import (
	"context"
	"testing"
	"time"

	"github.com/pingcap/errors"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/testutil"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/dataflow-service/internal/db/entity"
	"github.com/tidbcloud/dataflow-service/internal/metrics"
	"github.com/tidbcloud/dataflow-service/pkg/log"
	"github.com/tidbcloud/eks-provider/api/v1alpha1"
	"go.uber.org/zap"
)

type mockResourceAPI struct {
	getDMJobFunc    func(ctx context.Context, projectID, dmJobID uint64) (*v1alpha1.DMJob, error)
	deleteDMJobFunc func(ctx context.Context, projectID, dmJobID uint64) error
	isNotFoundFunc  func(err error) bool
}

func (m *mockResourceAPI) CreateDMJobIgnoreExists(_ context.Context, _ *v1alpha1.DMJob) error {
	return nil
}
func (m *mockResourceAPI) UpdateDMJob(_ context.Context, _ *v1alpha1.DMJob) error { return nil }
func (m *mockResourceAPI) GetDMJob(ctx context.Context, projectID, dmJobID uint64) (*v1alpha1.DMJob, error) {
	if m.getDMJobFunc != nil {
		return m.getDMJobFunc(ctx, projectID, dmJobID)
	}
	return nil, errors.New("not found")
}

func (m *mockResourceAPI) DeleteDMJob(ctx context.Context, projectID, dmJobID uint64) error {
	if m.deleteDMJobFunc != nil {
		return m.deleteDMJobFunc(ctx, projectID, dmJobID)
	}
	return nil
}
func (m *mockResourceAPI) GetObjectNamespace(_ uint64) string { return "default" }
func (m *mockResourceAPI) IsNotFound(err error) bool {
	if m.isNotFoundFunc != nil {
		return m.isNotFoundFunc(err)
	}
	return false
}

func newTestTaskContext(resourceAPI ResourceAPI) *TaskContext {
	now := time.Now()
	ctx := log.WithLogger(context.Background(), zap.NewNop())
	dfJob := &entity.DFJob{
		TenantID:  entity.ID(10),
		ClusterID: entity.ID(20),
		ProjectID: entity.ID(30),
		Phase:     entity.DFJobPhaseCreated,
	}
	dfJob.ID = entity.ID(1)
	dfJob.UpdatedAt = &now

	dmJob := &entity.DMJob{
		TenantID:  entity.ID(10),
		ProjectID: entity.ID(30),
		ClusterID: entity.ID(20),
		Mode:      entity.DMJobModeAll,
	}
	dmJob.ID = entity.ID(1)

	return &TaskContext{
		ctx: ctx,
		job: &Job{
			DFJob: dfJob,
			DMJob: dmJob,
		},
		model: Model{
			resourceAPI: resourceAPI,
		},
	}
}

func replaceCreateStageCounter(t *testing.T) *prometheus.CounterVec {
	counter := prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: metrics.DataflowNamespace,
			Subsystem: "dm",
			Name:      "create_stage_failed_total_state_test",
			Help:      "test counter",
		}, []string{metrics.TenantLabel, metrics.ClusterIDLabel, "job_id", metrics.StageLabel, metrics.ReasonClassLabel})
	orig := metrics.DMCreateStageFailedTotal
	metrics.DMCreateStageFailedTotal = counter
	t.Cleanup(func() { metrics.DMCreateStageFailedTotal = orig })
	return counter
}

func replaceModifyCounter(t *testing.T) *prometheus.CounterVec {
	counter := prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: metrics.DataflowNamespace,
			Subsystem: "dm",
			Name:      "modify_failed_total_state_test",
			Help:      "test counter",
		}, []string{metrics.TenantLabel, metrics.ClusterIDLabel, "job_id", metrics.StepLabel, metrics.ReasonClassLabel})
	orig := metrics.DMModifyFailedTotal
	metrics.DMModifyFailedTotal = counter
	t.Cleanup(func() { metrics.DMModifyFailedTotal = orig })
	return counter
}

// ---- WS1: StateDump tests ----

func TestStateDump_GetDMJobSourcesFails_NoCreateStageMetric(t *testing.T) {
	counter := replaceCreateStageCounter(t)
	mock := &mockResourceAPI{}
	tc := newTestTaskContext(mock)

	s := &StateDump{}
	_, err := s.Exec(tc)
	// GetDMJobSources will fail because there's no DB, but this is NOT a
	// dump_create or dump_start failure, so no create-stage metric is emitted.
	require.Error(t, err)

	reg := prometheus.NewRegistry()
	reg.MustRegister(counter)
	count, gerr := testutil.GatherAndCount(reg)
	require.NoError(t, gerr)
	require.Equal(t, 0, count, "no create-stage metric should be emitted for source loading failure")
}

func TestStateDump_GetDMJobCRFails_NoCreateStageMetric(t *testing.T) {
	counter := replaceCreateStageCounter(t)
	mock := &mockResourceAPI{
		getDMJobFunc: func(_ context.Context, _, _ uint64) (*v1alpha1.DMJob, error) {
			return nil, errors.New("cr not found")
		},
	}
	tc := newTestTaskContext(mock)
	tc.job.DMJob.Mode = entity.DMJobModeIncr // skip dump stage to avoid DB call

	s := &StateDump{}
	next, err := s.Exec(tc)
	// needDumpStage returns false for incr mode, so it skips to StateLoad
	require.NoError(t, err)
	require.NotNil(t, next)

	reg := prometheus.NewRegistry()
	reg.MustRegister(counter)
	count, gerr := testutil.GatherAndCount(reg)
	require.NoError(t, gerr)
	require.Equal(t, 0, count, "no create-stage metric when dump stage is skipped")
}

// ---- WS2: StateScale tests ----

func TestStateScale_GetDMJobCRFails_EmitsModifyMetric(t *testing.T) {
	counter := replaceModifyCounter(t)
	mock := &mockResourceAPI{
		getDMJobFunc: func(_ context.Context, _, _ uint64) (*v1alpha1.DMJob, error) {
			return nil, errors.New("cr api error")
		},
	}
	tc := newTestTaskContext(mock)

	s := &StateScale{}
	_, err := s.Exec(tc)
	require.Error(t, err)

	val := testutil.ToFloat64(counter.WithLabelValues("10", "20", "1", "load_job", "get_cr"))
	require.InDelta(t, 1, val, 0.001, "modify metric should be emitted for CR load failure")
}

func TestStateScale_SuccessfulRetry_NoModifyMetric(t *testing.T) {
	counter := replaceModifyCounter(t)
	mock := &mockResourceAPI{
		getDMJobFunc: func(_ context.Context, _, _ uint64) (*v1alpha1.DMJob, error) {
			return &v1alpha1.DMJob{
				Status: v1alpha1.DMJobStatus{
					Phase:            v1alpha1.DMJobPhaseScaling,
					MasterServerName: "http://master:8261",
				},
			}, nil
		},
	}
	tc := newTestTaskContext(mock)

	// GetDMJobSources will fail (no DB), so it will hit load_job/get_sources.
	// That IS a metric emission. But the test verifies it does NOT emit on
	// the CR path since GetDMJob succeeded.
	s := &StateScale{}
	_, _ = s.Exec(tc)

	crVal := testutil.ToFloat64(counter.WithLabelValues("10", "20", "1", "load_job", "get_cr"))
	require.InDelta(t, 0, crVal, 0.001, "no CR failure metric when GetDMJob succeeds")
}

func TestCreateStageMetricHelper_DirectEmission(t *testing.T) {
	counter := replaceCreateStageCounter(t)
	dfJob := &entity.DFJob{
		TenantID:  entity.ID(10),
		ClusterID: entity.ID(20),
	}
	dfJob.ID = entity.ID(42)

	metrics.IncDMCreateStageFailure(dfJob, "dump_create", "create_task")
	val := testutil.ToFloat64(counter.WithLabelValues("10", "20", "42", "dump_create", "create_task"))
	require.InDelta(t, 1, val, 0.001)

	metrics.IncDMCreateStageFailure(dfJob, "dump_start", "start_task")
	val = testutil.ToFloat64(counter.WithLabelValues("10", "20", "42", "dump_start", "start_task"))
	require.InDelta(t, 1, val, 0.001)
}

func TestModifyMetricHelper_DirectEmission(t *testing.T) {
	counter := replaceModifyCounter(t)
	dfJob := &entity.DFJob{
		TenantID:  entity.ID(10),
		ClusterID: entity.ID(20),
	}
	dfJob.ID = entity.ID(42)

	allSteps := []struct {
		step        string
		reasonClass string
	}{
		{"load_job", "get_cr"},
		{"load_job", "get_sources"},
		{"load_spec", "get_modification"},
		{"load_spec", "get_spec"},
		{"scale_worker", "scale"},
		{"update_phase", "update"},
	}

	for _, s := range allSteps {
		metrics.IncDMModifyFailure(dfJob, s.step, s.reasonClass)
		val := testutil.ToFloat64(counter.WithLabelValues("10", "20", "42", s.step, s.reasonClass))
		require.InDelta(t, 1, val, 0.001, "step=%s reason_class=%s", s.step, s.reasonClass)
	}
}