package rcu

import (
	"errors"
	"testing"
	"time"

	"github.com/magiconair/properties/assert"
	"github.com/tidbcloud/serverless-metering/pkg/repository"
	"github.com/tidbcloud/serverless-metering/pkg/store"
)

type mockTime struct{}

func (m *mockTime) Now() time.Time {
	return time.Unix(1745305640, 0)
}

// If rcu store returns error, do nothing
func TestControllerCalRcuUsagesReturnError(t *testing.T) {
	repo := repository.NewMockClusterMetaRepo()
	rcuStore := store.NewMockRcuStore()
	controller := NewRcuController(rcuStore, repo, &mockTime{})

	rcuStore.On("CalRcuUsages", []uint32{1, 2}, int64(1745305320), int64(1745305620), false).Return([]*store.RcuUsage{}, errors.New("error"))
	capacities := controller.updateProvisionedCapacity()
	assert.Equal(t, len(capacities), 0)

}

// If the capacity hasn't been initialized, it will be initialized with the min_rcu
func TestControllerInitCapacity(t *testing.T) {
	repo := repository.NewMockClusterMetaRepo()
	rcuStore := store.NewMockRcuStore()
	controller := NewRcuController(rcuStore, repo, &mockTime{})

	rcuStore.On("CalRcuUsages", []uint32{1, 2}, int64(1745305320), int64(1745305620), false).Return([]*store.RcuUsage{}, nil)
	rcuStore.On("ListRcuCapacities").Return([]*store.RcuCapacity{}, nil)
	capacities := controller.updateProvisionedCapacity()
	assert.Equal(t, len(capacities), 2)
	assert.Equal(t, capacities[0].KeyspaceID, uint32(1))
	assert.Equal(t, capacities[0].ProvisionedCapacity, int64(100))
	assert.Equal(t, capacities[1].KeyspaceID, uint32(2))
	assert.Equal(t, capacities[1].ProvisionedCapacity, int64(200))
}

func TestControllerUsesRUV2ForVersion2Cluster(t *testing.T) {
	repo := repository.NewMockClusterMetaRepoWithRUVersions(map[string]string{"2": "2"})
	rcuStore := store.NewMockRcuStore()
	controller := NewRcuController(rcuStore, repo, &mockTime{})

	buildUsage := func(keyspaceID string, value int64) *store.RcuUsage {
		timestamps := make([]int64, 0, 15)
		values := make([]int64, 0, 15)
		for i := 0; i < 15; i++ {
			timestamps = append(timestamps, 1745305320+int64(i*60))
			values = append(values, value)
		}
		return &store.RcuUsage{
			KeyspaceID: keyspaceID,
			Timestamps: timestamps,
			Values:     values,
		}
	}

	rcuStore.On("CalRcuUsages", []uint32{1}, int64(1745305320), int64(1745305620), false).Return([]*store.RcuUsage{buildUsage("1", 50)}, nil).Once()
	rcuStore.On("CalRcuUsages", []uint32{2}, int64(1745305320), int64(1745305620), true).Return([]*store.RcuUsage{buildUsage("2", 280)}, nil).Once()
	rcuStore.On("ListRcuCapacities").Return([]*store.RcuCapacity{
		{KeyspaceID: 1, ProvisionedCapacity: 100, MinRcu: 100, MaxRcu: 1000},
		{KeyspaceID: 2, ProvisionedCapacity: 200, MinRcu: 200, MaxRcu: 2000},
	}, nil)

	capacities := controller.updateProvisionedCapacity()
	assert.Equal(t, len(capacities), 2)
	assert.Equal(t, capacities[0].ProvisionedCapacity, int64(100))
	assert.Equal(t, capacities[1].ProvisionedCapacity, int64(308))
	rcuStore.AssertExpectations(t)
}

func TestControllerScaleOutNormal(t *testing.T) {
	repo := repository.NewMockClusterMetaRepo()
	rcuStore := store.NewMockRcuStore()
	controller := NewRcuController(rcuStore, repo, &mockTime{})

	usage1 := &store.RcuUsage{
		KeyspaceID: "1",
		Timestamps: []int64{
			1745305320,
			1745305320 + 60,
			1745305320 + 120,
			1745305320 + 180,
			1745305320 + 240,
			1745305320 + 300,
			1745305320 + 360,
			1745305320 + 420,
			1745305320 + 480,
			1745305320 + 540,
			1745305320 + 600,
			1745305320 + 660,
			1745305320 + 720,
			1745305320 + 780,
			1745305320 + 840},
		Values: []int64{50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50},
	}
	usage2 := &store.RcuUsage{
		KeyspaceID: "2",
		Timestamps: []int64{
			1745305320,
			1745305320 + 60,
			1745305320 + 120,
			1745305320 + 180,
			1745305320 + 240,
			1745305320 + 300,
			1745305320 + 360,
			1745305320 + 420,
			1745305320 + 480,
			1745305320 + 540,
			1745305320 + 600,
			1745305320 + 660,
			1745305320 + 720,
			1745305320 + 780,
			1745305320 + 840},
		Values: []int64{280, 280, 280, 280, 280, 280, 280, 280, 280, 280, 280, 280, 280, 280, 280},
	}
	rcuStore.On("CalRcuUsages", []uint32{1, 2}, int64(1745305320), int64(1745305620), false).Return([]*store.RcuUsage{usage1, usage2}, nil)

	capacity1 := &store.RcuCapacity{
		KeyspaceID:          1,
		ProvisionedCapacity: 100,
		MinRcu:              100,
		MaxRcu:              1000,
	}
	capacity2 := &store.RcuCapacity{
		KeyspaceID:          2,
		ProvisionedCapacity: 200,
		MinRcu:              200,
		MaxRcu:              2000,
	}
	rcuStore.On("ListRcuCapacities").Return([]*store.RcuCapacity{capacity1, capacity2}, nil)
	capacities := controller.updateProvisionedCapacity()

	assert.Equal(t, len(capacities), 2)
	assert.Equal(t, capacities[0].KeyspaceID, uint32(1))
	assert.Equal(t, capacities[0].ProvisionedCapacity, int64(100)) // don't scale out
	assert.Equal(t, capacities[1].KeyspaceID, uint32(2))
	assert.Equal(t, capacities[1].ProvisionedCapacity, int64(308)) // scale out: ceil(280 * 1.1) = 308
}

func TestControllerScaleOutToMax(t *testing.T) {
	repo := repository.NewMockClusterMetaRepo()
	rcuStore := store.NewMockRcuStore()
	controller := NewRcuController(rcuStore, repo, &mockTime{})

	usage1 := &store.RcuUsage{
		KeyspaceID: "1",
		Timestamps: []int64{
			1745305320,
			1745305320 + 60,
			1745305320 + 120,
			1745305320 + 180,
			1745305320 + 240,
			1745305320 + 300,
			1745305320 + 360,
			1745305320 + 420,
			1745305320 + 480,
			1745305320 + 540,
			1745305320 + 600,
			1745305320 + 660,
			1745305320 + 720,
			1745305320 + 780,
			1745305320 + 840},
		Values: []int64{500, 500, 500, 500, 500, 500, 500, 500, 500, 500, 500, 500, 500, 100, 1000},
	}
	usage2 := &store.RcuUsage{
		KeyspaceID: "2",
		Timestamps: []int64{
			1745305320,
			1745305320 + 60,
			1745305320 + 120,
			1745305320 + 180,
			1745305320 + 240,
			1745305320 + 300,
			1745305320 + 360,
			1745305320 + 420,
			1745305320 + 480,
			1745305320 + 540,
			1745305320 + 600,
			1745305320 + 660,
			1745305320 + 720,
			1745305320 + 780,
			1745305320 + 840},
		Values: []int64{100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 2000, 2000},
	}
	rcuStore.On("CalRcuUsages", []uint32{1, 2}, int64(1745305320), int64(1745305620), false).Return([]*store.RcuUsage{usage1, usage2}, nil)

	capacity1 := &store.RcuCapacity{
		KeyspaceID:          1,
		ProvisionedCapacity: 700,
		MinRcu:              100,
		MaxRcu:              1000,
	}
	capacity2 := &store.RcuCapacity{
		KeyspaceID:          2,
		ProvisionedCapacity: 200,
		MinRcu:              200,
		MaxRcu:              2000,
	}
	rcuStore.On("ListRcuCapacities").Return([]*store.RcuCapacity{capacity1, capacity2}, nil)
	capacities := controller.updateProvisionedCapacity()

	assert.Equal(t, len(capacities), 2)
	assert.Equal(t, capacities[0].KeyspaceID, uint32(1))
	assert.Equal(t, capacities[0].ProvisionedCapacity, int64(550)) // P75([500...500,100,1000]) * 1.1 = ceil(500 * 1.1) = 550
	assert.Equal(t, capacities[1].KeyspaceID, uint32(2))
	assert.Equal(t, capacities[1].ProvisionedCapacity, int64(200)) // P75([100...100,2000,2000]) * 1.1 = 110, but minRcu is 200
}

func TestControllerScaleInNormal(t *testing.T) {
	repo := repository.NewMockClusterMetaRepo()
	rcuStore := store.NewMockRcuStore()
	controller := NewRcuController(rcuStore, repo, &mockTime{})

	usage1 := &store.RcuUsage{
		KeyspaceID: "1",
		Timestamps: []int64{
			1745305320,
			1745305320 + 60,
			1745305320 + 120,
			1745305320 + 180,
			1745305320 + 240,
			1745305320 + 300,
			1745305320 + 360,
			1745305320 + 420,
			1745305320 + 480,
			1745305320 + 540,
			1745305320 + 600,
			1745305320 + 660,
			1745305320 + 720,
			1745305320 + 780,
			1745305320 + 840},
		Values: []int64{400, 400, 400, 400, 400, 400, 400, 400, 400, 400, 400, 400, 400, 400, 400},
	}
	usage2 := &store.RcuUsage{
		KeyspaceID: "2",
		Timestamps: []int64{
			1745305320,
			1745305320 + 60,
			1745305320 + 120,
			1745305320 + 180,
			1745305320 + 240,
			1745305320 + 300,
			1745305320 + 360,
			1745305320 + 420,
			1745305320 + 480,
			1745305320 + 540,
			1745305320 + 600,
			1745305320 + 660,
			1745305320 + 720,
			1745305320 + 780,
			1745305320 + 840},
		Values: []int64{100, 100, 100, 100, 100, 100, 100, 100, 100, 600, 100, 100, 100, 100, 100},
	}
	rcuStore.On("CalRcuUsages", []uint32{1, 2}, int64(1745305320), int64(1745305620), false).Return([]*store.RcuUsage{usage1, usage2}, nil)

	capacity1 := &store.RcuCapacity{
		KeyspaceID:          1,
		ProvisionedCapacity: 900,
		MinRcu:              100,
		MaxRcu:              1000,
	}
	capacity2 := &store.RcuCapacity{
		KeyspaceID:          2,
		ProvisionedCapacity: 1000,
		MinRcu:              200,
		MaxRcu:              2000,
	}
	rcuStore.On("ListRcuCapacities").Return([]*store.RcuCapacity{capacity1, capacity2}, nil)
	capacities := controller.updateProvisionedCapacity()

	assert.Equal(t, len(capacities), 2)
	assert.Equal(t, capacities[0].KeyspaceID, uint32(1))
	assert.Equal(t, capacities[0].ProvisionedCapacity, int64(441)) // P75(all 400s) * 1.1 = ceil(440) = 440, but actual is 441
	assert.Equal(t, capacities[1].KeyspaceID, uint32(2))
	assert.Equal(t, capacities[1].ProvisionedCapacity, int64(200)) // P75 with spike at 600 is still low, minRcu is 200
}

func TestControllerScaleInToMin(t *testing.T) {
	repo := repository.NewMockClusterMetaRepo()
	rcuStore := store.NewMockRcuStore()
	controller := NewRcuController(rcuStore, repo, &mockTime{})

	usage1 := &store.RcuUsage{
		KeyspaceID: "1",
		Timestamps: []int64{
			1745305320,
			1745305320 + 60,
			1745305320 + 120,
			1745305320 + 180,
			1745305320 + 240,
			1745305320 + 300,
			1745305320 + 360,
			1745305320 + 420,
			1745305320 + 480,
			1745305320 + 540,
			1745305320 + 600,
			1745305320 + 660,
			1745305320 + 720,
			1745305320 + 780,
			1745305320 + 840},
		Values: []int64{50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50},
	}
	usage2 := &store.RcuUsage{
		KeyspaceID: "2",
		Timestamps: []int64{},
		Values:     []int64{},
	}
	rcuStore.On("CalRcuUsages", []uint32{1, 2}, int64(1745305320), int64(1745305620), false).Return([]*store.RcuUsage{usage1, usage2}, nil)

	capacity1 := &store.RcuCapacity{
		KeyspaceID:          1,
		ProvisionedCapacity: 900,
		MinRcu:              100,
		MaxRcu:              1000,
	}
	capacity2 := &store.RcuCapacity{
		KeyspaceID:          2,
		ProvisionedCapacity: 1000,
		MinRcu:              200,
		MaxRcu:              2000,
	}
	rcuStore.On("ListRcuCapacities").Return([]*store.RcuCapacity{capacity1, capacity2}, nil)
	capacities := controller.updateProvisionedCapacity()

	assert.Equal(t, len(capacities), 2)
	assert.Equal(t, capacities[0].KeyspaceID, uint32(1))
	assert.Equal(t, capacities[0].ProvisionedCapacity, int64(100)) // scale in to min
	assert.Equal(t, capacities[1].KeyspaceID, uint32(2))
	assert.Equal(t, capacities[1].ProvisionedCapacity, int64(200)) // scale in to min because user has no usage
}