package query

import (
	"context"
	"sort"
	"testing"
	"time"

	"github.com/tidbcloud/serverless-metering/pkg/config"
	"github.com/tidbcloud/serverless-metering/pkg/store"

	"github.com/tj/assert"
)

func TestDuckDBQueryEngine(t *testing.T) {
	ctx := context.TODO()
	config.InitMockConfig()
	storage := store.NewMockStorage()
	engine, err := NewDuckDBQueryEngine(ctx, storage, time.Now().Unix(), 1)
	assert.Nil(t, err)
	defer engine.Close()
	// test load data ru
	loadDataRu, err := engine.CalLoadDataRu(ctx)
	assert.Nil(t, err)
	usages := loadDataRu
	assert.Equal(t, 3, len(usages))
	sort.Slice(usages, func(i, j int) bool { return usages[i].KeyspaceId < usages[j].KeyspaceId })
	assert.Equal(t, MetricTypeLoadDataRU, usages[0].MetricType)
	assert.Equal(t, MetricTypeLoadDataRU, usages[1].MetricType)
	assert.Equal(t, MetricTypeLoadDataRU, usages[2].MetricType)
	assert.Equal(t, "1", usages[0].KeyspaceId)
	assert.Equal(t, "2", usages[1].KeyspaceId)
	assert.Equal(t, "3", usages[2].KeyspaceId)
	assert.Equal(t, int64(19800), usages[0].Value)
	assert.Equal(t, int64(9900), usages[1].Value)
	assert.Equal(t, int64(4900), usages[2].Value)
	// test public egress
	trafficRu, err := engine.CalPublicEgress(ctx)
	assert.Nil(t, err)
	usages = trafficRu
	assert.Equal(t, 2, len(usages))
	sort.Slice(usages, func(i, j int) bool { return usages[i].PrimaryId < usages[j].PrimaryId })
	assert.Equal(t, MetricTypePublicEgress, usages[0].MetricType)
	assert.Equal(t, MetricTypePublicEgress, usages[1].MetricType)
	assert.Equal(t, "1001", usages[0].PrimaryId)
	assert.Equal(t, int64(29700), usages[0].Value)
	assert.Equal(t, "1002", usages[1].PrimaryId)
	assert.Equal(t, int64(9800), usages[1].Value)
	// test privatelik egress
	privatelinkEgress, err := engine.CalPrivatelinkEgress(ctx)
	assert.Nil(t, err)
	usages = privatelinkEgress
	assert.Equal(t, 2, len(usages))
	sort.Slice(usages, func(i, j int) bool { return usages[i].PrimaryId < usages[j].PrimaryId })
	assert.Equal(t, MetricTypePrivatelinkEgress, usages[0].MetricType)
	assert.Equal(t, MetricTypePrivatelinkEgress, usages[1].MetricType)
	assert.Equal(t, "1001", usages[0].PrimaryId)
	assert.Equal(t, int64(29700), usages[0].Value)
	assert.Equal(t, "1002", usages[1].PrimaryId)
	assert.Equal(t, int64(9800), usages[1].Value)
	// test tidb ru
	tidbRu, err := engine.CalTidbRU(ctx)
	assert.Nil(t, err)
	usages = tidbRu
	assert.Equal(t, 1, len(usages))
	assert.Equal(t, MetricTypeTidbRU, usages[0].MetricType)
	assert.Equal(t, "1", usages[0].KeyspaceId)
	assert.Equal(t, int64(79200), usages[0].Value)
	// test tiflash ru
	tiflashRu, err := engine.CalTiflashRU(ctx)
	assert.Nil(t, err)
	usages = tiflashRu
	assert.Equal(t, 1, len(usages))
	assert.Equal(t, MetricTypeTiflashRU, usages[0].MetricType)
	assert.Equal(t, "1", usages[0].KeyspaceId)
	assert.Equal(t, int64(29700), usages[0].Value)
	// test tidb ru v2
	tidbRUV2, err := engine.CalTidbRUV2(ctx)
	assert.Nil(t, err)
	usages = tidbRUV2
	assert.Equal(t, 1, len(usages))
	assert.Equal(t, MetricTypeTidbRUV2, usages[0].MetricType)
	assert.Equal(t, "1", usages[0].KeyspaceId)
	assert.Equal(t, int64(9900), usages[0].Value)
	// test tiflash ru v2
	tiflashRUV2, err := engine.CalTiflashRUV2(ctx)
	assert.Nil(t, err)
	usages = tiflashRUV2
	assert.Equal(t, 1, len(usages))
	assert.Equal(t, MetricTypeTiflashRUV2, usages[0].MetricType)
	assert.Equal(t, "1", usages[0].KeyspaceId)
	assert.Equal(t, int64(9900), usages[0].Value)
	// test tikv storage medium
	tikvStorage, err := engine.CalTikvStorageMedium(ctx)
	assert.Nil(t, err)
	usages = tikvStorage
	assert.Equal(t, 2, len(usages))
	sort.Slice(usages, func(i, j int) bool {
		return usages[i].KeyspaceId < usages[j].KeyspaceId
	})
	assert.Equal(t, MetricTypeTikvStorage, usages[0].MetricType)
	assert.Equal(t, MetricTypeTikvStorage, usages[1].MetricType)
	assert.Equal(t, "1", usages[0].KeyspaceId)
	assert.Equal(t, "2", usages[1].KeyspaceId)
	assert.Equal(t, int64((60+210)*5*6*1024*1024), usages[0].Value)
	assert.Equal(t, int64((70+245)*5*6*1024*1024), usages[1].Value)
	// test tiflash storage medium
	tiflashStorage, err := engine.CalTiflashStorageMedium(ctx)
	assert.Nil(t, err)
	usages = tiflashStorage
	assert.Equal(t, 1, len(usages))
	assert.Equal(t, MetricTypeTiflashStorage, usages[0].MetricType)
	assert.Equal(t, "1", usages[0].KeyspaceId)
	assert.Equal(t, int64((60+210)*5*6*1024*1024), usages[0].Value)
	// test tikv storage large
	tikvStorage, err = engine.CalTikvStorageLarge(ctx)
	assert.Nil(t, err)
	usages = tikvStorage
	assert.Equal(t, 1, len(usages))
	assert.Equal(t, MetricTypeTikvStorage, usages[0].MetricType)
	assert.Equal(t, "3", usages[0].KeyspaceId)
	assert.Equal(t, int64((2000+7000)*5*6*1024*1024), usages[0].Value)
	// test tiflash storage medium
	tiflashStorage, err = engine.CalTiflashStorageLarge(ctx)
	assert.Nil(t, err)
	usages = tiflashStorage
	assert.Equal(t, 1, len(usages))
	assert.Equal(t, MetricTypeTiflashStorage, usages[0].MetricType)
	assert.Equal(t, "3", usages[0].KeyspaceId)
	assert.Equal(t, int64((2000+7000)*5*6*1024*1024), usages[0].Value)
	// test exec count
	execCount, err := engine.CalExecCount(ctx)
	assert.Nil(t, err)
	assert.NotNil(t, execCount)
	// Check cluster 1001 has 2 databases
	assert.Equal(t, 2, len(execCount["1001"]))
	// test_db should have 99*100 executions (counter from 0 to 9900, increments of 100)
	assert.Equal(t, int64(9900), execCount["1001"]["test_db"])
	// another_db should have 99*100 executions
	assert.Equal(t, int64(9900), execCount["1001"]["another_db"])
	// Check cluster 1002 has 1 database (with reset, so 49*100 before reset + 49*100 after reset = 9800)
	assert.Equal(t, 1, len(execCount["1002"]))
	assert.Equal(t, int64(9800), execCount["1002"]["test_db"])
	// Check cluster 1003 has 1 database
	assert.Equal(t, 1, len(execCount["1003"]))
	assert.Equal(t, int64(9900), execCount["1003"]["prod_db"])
}