package base

import (
	"context"
	"errors"
	"sync"
	"testing"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/stretchr/testify/assert"
	"github.com/tidbcloud/dataflow-service/app/repo/entity"
)

// mockIdAndVersioning implements entity.IdAndVersioning for testing
type mockIdAndVersioning struct {
	id     entity.ID
	status entity.BackupExportStatus
}

func (m mockIdAndVersioning) GetID() entity.ID   { return m.id }
func (m mockIdAndVersioning) GetVersioning() int { return 1 }

func TestMetricsCollector_Collect(t *testing.T) {
	type fields struct {
		unfinished map[entity.ID]mockIdAndVersioning
	}
	type args struct {
		unfinishedObjs []mockIdAndVersioning
		finishedObjs   []mockIdAndVersioning
	}
	var mu sync.Mutex
	tests := map[string]struct {
		fields            fields
		args              args
		wantCollected     map[entity.ID]int
		wantDeleted       map[entity.ID]int
		wantCalled        int
		wantPanic         bool
		listUnfinishedErr error
		listByIDsErr      error
	}{
		"unfinished_and_finished": {
			fields: fields{
				unfinished: map[entity.ID]mockIdAndVersioning{
					entity.ID(1): {id: entity.ID(1)},
					entity.ID(3): {id: entity.ID(3)},
				},
			},
			args: args{
				unfinishedObjs: []mockIdAndVersioning{
					{id: entity.ID(1)},
					{id: entity.ID(2)},
				},
				finishedObjs: []mockIdAndVersioning{
					{id: entity.ID(3)},
				},
			},
			wantCollected: map[entity.ID]int{
				entity.ID(1): 1,
				entity.ID(2): 1,
				entity.ID(3): 1,
			},
			wantDeleted: map[entity.ID]int{},
		},
		"list_unfinished_error": {
			fields:            fields{},
			args:              args{},
			listUnfinishedErr: errors.New("fail"),
			wantCollected:     map[entity.ID]int{},
			wantDeleted:       map[entity.ID]int{},
		},
		"list_by_ids_error": {
			fields: fields{
				unfinished: map[entity.ID]mockIdAndVersioning{
					entity.ID(2): {id: entity.ID(2)},
				},
			},
			args: args{
				unfinishedObjs: []mockIdAndVersioning{
					{id: entity.ID(1)},
				},
			},
			listByIDsErr: errors.New("fail"),
			wantCollected: map[entity.ID]int{ // the unfinished ones are still collected
				entity.ID(1): 1,
			},
			wantDeleted: map[entity.ID]int{},
		},
		"no_finished": {
			fields: fields{},
			args: args{
				unfinishedObjs: []mockIdAndVersioning{
					{id: entity.ID(1)},
				},
			},
			wantCalled: 1,
			wantCollected: map[entity.ID]int{
				entity.ID(1): 1,
			},
			wantDeleted: map[entity.ID]int{},
		},
	}

	for name, tt := range tests {
		t.Run(name, func(t *testing.T) {
			collected := make(map[entity.ID]int)
			deletions := make(map[entity.ID]int)
			var called int
			target := &Target[mockIdAndVersioning]{
				Name: "test",
				StatusCollector: prometheus.NewDesc(
					"test_status",
					"desc",
					nil, nil,
				),
				ListUnfinished: func(ctx context.Context) ([]mockIdAndVersioning, error) {
					if tt.listUnfinishedErr != nil {
						return nil, tt.listUnfinishedErr
					}
					return tt.args.unfinishedObjs, nil
				},
				ListByIDs: func(ctx context.Context, ids []entity.ID) ([]mockIdAndVersioning, error) {
					if tt.listByIDsErr != nil {
						return nil, tt.listByIDsErr
					}
					// Only for "unfinished_and_finished" case
					if len(tt.args.finishedObjs) > 0 {
						assert.ElementsMatch(t, []entity.ID{entity.ID(3)}, ids)
						return tt.args.finishedObjs, nil
					}
					return nil, nil
				},
				UpdateStatusMetrics: func(desc *prometheus.Desc, obj mockIdAndVersioning) {
					mu.Lock()
					collected[obj.GetID()]++
					mu.Unlock()
					called++
				},
				DeleteStatusMetrics: func(obj mockIdAndVersioning) {
					mu.Lock()
					deletions[obj.GetID()]++
					mu.Unlock()
				},
			}
			unfinished := tt.fields.unfinished
			if unfinished == nil {
				unfinished = make(map[entity.ID]mockIdAndVersioning)
			}
			c := &metricsCollector[mockIdAndVersioning]{
				target:     target,
				unfinished: unfinished,
			}
			ch := make(chan prometheus.Metric, 10)
			c.Collect(ch)
			close(ch)
			assert.Equal(t, tt.wantCollected, collected)
			assert.Equal(t, tt.wantDeleted, deletions)
			if tt.wantCalled > 0 {
				assert.Equal(t, tt.wantCalled, called)
			}
		})
	}
}

func TestMetricsCollector_CollectFinishedBackupExportStatusesDeletesMetrics(t *testing.T) {
	statuses := []entity.BackupExportStatus{
		entity.BackupExportStatusFailed,
		entity.BackupExportStatusSuccess,
		entity.BackupExportStatusCanceled,
	}
	collected := make(map[entity.ID]int)
	deleted := make(map[entity.ID]entity.BackupExportStatus)

	target := &Target[mockIdAndVersioning]{
		Name:            "backupExport",
		StatusCollector: prometheus.NewDesc("test_status", "desc", nil, nil),
		ListUnfinished: func(context.Context) ([]mockIdAndVersioning, error) {
			return nil, nil
		},
		ListByIDs: func(context.Context, []entity.ID) ([]mockIdAndVersioning, error) {
			return []mockIdAndVersioning{
				{id: entity.ID(1), status: statuses[0]},
				{id: entity.ID(2), status: statuses[1]},
				{id: entity.ID(3), status: statuses[2]},
			}, nil
		},
		UpdateStatusMetrics: func(_ *prometheus.Desc, obj mockIdAndVersioning) {
			collected[obj.GetID()]++
		},
		DeleteStatusMetrics: func(obj mockIdAndVersioning) {
			deleted[obj.GetID()] = obj.status
		},
	}

	collector := &metricsCollector[mockIdAndVersioning]{
		target: target,
		unfinished: map[entity.ID]mockIdAndVersioning{
			entity.ID(1): {id: entity.ID(1), status: statuses[0]},
			entity.ID(2): {id: entity.ID(2), status: statuses[1]},
			entity.ID(3): {id: entity.ID(3), status: statuses[2]},
		},
	}

	collector.Collect(make(chan prometheus.Metric))

	assert.Equal(t, map[entity.ID]int{
		entity.ID(1): 1,
		entity.ID(2): 1,
		entity.ID(3): 1,
	}, collected)
	assert.Empty(t, deleted)
}