package load

import (
	"bytes"
	"context"
	"io"
	"testing"

	"github.com/golang/mock/gomock"
	lru "github.com/hashicorp/golang-lru/v2"
	"github.com/stretchr/testify/require"

	"github.com/pingcap-inc/clinic/pkg/ngm/storage"
	"github.com/pingcap-inc/clinic/pkg/ngm/storage/mock"
)

func TestLoadFromStorage(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	sf := mock.NewMockStorageFactory(ctrl)
	store := mock.NewMockStorage(ctrl)

	sf.EXPECT().
		Storage(gomock.Any(),
			gomock.Eq("aws"),
			gomock.Eq("us-east-1")).
		Return(store, nil).Times(2)

	store.EXPECT().
		Objects(gomock.Any(),
			gomock.Eq("o11y-prod-shared-us-east-1"),
			gomock.Eq("0/prod-devtier/serverless/7777/8888/9999/logs/statement/160000")).
		Return([]storage.ObjectMeta{
			{Key: "1600001000-stmt.log"},
			{Key: "1600002000-stmt.log"},
			{Key: "1600003000-stmt.log"},
		}, nil).Times(2)

	store.EXPECT().
		Object(gomock.Any(),
			gomock.Eq("o11y-prod-shared-us-east-1"),
			gomock.Eq("1600001000-stmt.log")).
		Return(&storage.Object{
			Reader: io.NopCloser(bytes.NewReader([]byte(
				`{"begin":1600000700,"end":1600000800,"digest":"d1"}
{"begin":1600000700,"end":1600000800,"digest":"d2"}`))),
		}, nil).Times(1)

	store.EXPECT().
		Object(gomock.Any(),
			gomock.Eq("o11y-prod-shared-us-east-1"),
			gomock.Eq("1600002000-stmt.log")).
		Return(&storage.Object{
			Reader: io.NopCloser(bytes.NewReader([]byte(
				`{"begin":1600001700,"end":1600002000,"digest":"d3"}
{"begin":1600001700,"end":1600002000,"digest":"d4"}`))),
		}, nil).Times(2)

	store.EXPECT().
		Object(gomock.Any(),
			gomock.Eq("o11y-prod-shared-us-east-1"),
			gomock.Eq("1600003000-stmt.log")).
		Return(&storage.Object{
			Reader: io.NopCloser(bytes.NewReader([]byte(
				`{"begin":1600002100,"end":1600002200,"digest":"d5"}
{"begin":1600002400,"end":1600002800,"digest":"d6"}`))),
		}, nil).Times(1)

	// load all
	param := &LoadParam{
		Env:          "prod",
		Provider:     "aws",
		Region:       "us-east-1",
		OrgID:        7777,
		ProjectID:    8888,
		ClusterID:    "9999",
		SharedTier:   true,
		StartUnixSec: 1600000000,
		EndUnixSec:   1600003000,
	}
	cache, _ := lru.New[CacheKey, TimeRangeRecords](2000)

	ch, err := Load(context.Background(), param, sf, cache)
	require.NoError(t, err)

	var res []TimeRangeRecords
	for r := range ch {
		res = append(res, r)
	}

	require.Len(t, res, 3)
	require.Equal(t, cache.Len(), 3)

	// load part
	param.StartUnixSec = 1600001500
	param.EndUnixSec = 1600002500
	cache, _ = lru.New[CacheKey, TimeRangeRecords](2000)
	ch, err = Load(context.Background(), param, sf, cache)
	require.NoError(t, err)

	res = nil
	for r := range ch {
		res = append(res, r)
	}
	require.Len(t, res, 1)
	require.Equal(t, cache.Len(), 1)
}

func TestLoadFromCache(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	sf := mock.NewMockStorageFactory(ctrl)
	store := mock.NewMockStorage(ctrl)

	sf.EXPECT().
		Storage(gomock.Any(),
			gomock.Eq("aws"),
			gomock.Eq("us-east-1")).
		Return(store, nil).Times(2)

	store.EXPECT().
		Objects(gomock.Any(),
			gomock.Eq("o11y-prod-shared-us-east-1"),
			gomock.Eq("0/prod-devtier/serverless/7777/8888/9999/logs/statement/160000")).
		Return([]storage.ObjectMeta{
			{Key: "1600001000-stmt.log"},
			{Key: "1600002000-stmt.log"},
			{Key: "1600003000-stmt.log"},
		}, nil).Times(2)

	cache, _ := lru.New[CacheKey, TimeRangeRecords](2000)
	cache.Add(CacheKey{"o11y-prod-shared-us-east-1", "1600001000-stmt.log"},
		TimeRangeRecords{{Begin: 1600000700, End: 1600000800}: {
			"d1": {{Begin: 1600000700, End: 1600000800, Digest: "d1"}},
			"d2": {{Begin: 1600000700, End: 1600000800, Digest: "d2"}},
		}},
	)
	cache.Add(CacheKey{"o11y-prod-shared-us-east-1", "1600002000-stmt.log"},
		TimeRangeRecords{{Begin: 1600001700, End: 1600002000}: {
			"d3": {{Begin: 1600001700, End: 1600002000, Digest: "d3"}},
			"d4": {{Begin: 1600001700, End: 1600002000, Digest: "d4"}},
		}},
	)
	cache.Add(CacheKey{"o11y-prod-shared-us-east-1", "1600003000-stmt.log"},
		TimeRangeRecords{{Begin: 1600002100, End: 1600002200}: {
			"d5": {{Begin: 1600002100, End: 1600002200, Digest: "d5"}},
		}, {Begin: 1600002400, End: 1600002800}: {
			"d6": {{Begin: 1600002400, End: 1600002800, Digest: "d6"}},
		}},
	)

	// load all
	param := &LoadParam{
		Env:          "prod",
		Provider:     "aws",
		Region:       "us-east-1",
		OrgID:        7777,
		ProjectID:    8888,
		ClusterID:    "9999",
		SharedTier:   true,
		StartUnixSec: 1600000000,
		EndUnixSec:   1600003000,
	}

	ch, err := Load(context.Background(), param, sf, cache)
	require.NoError(t, err)

	var res []TimeRangeRecords
	for r := range ch {
		res = append(res, r)
	}

	require.Len(t, res, 3)

	// load part
	param.StartUnixSec = 1600001500
	param.EndUnixSec = 1600002500
	ch, err = Load(context.Background(), param, sf, cache)
	require.NoError(t, err)

	res = nil
	for r := range ch {
		res = append(res, r)
	}
	require.Len(t, res, 1)
}

func TestStorageInfoLocationOverridesDefaultPath(t *testing.T) {
	param := &LoadParam{
		Env:           "prod",
		Provider:      "aws",
		Region:        "us-east-1",
		OrgID:         7777,
		ProjectID:     8888,
		ClusterID:     "9999",
		SharedTier:    true,
		StorageBucket: "bucket-from-o11y",
		StoragePrefix: "path/from/o11y",
	}

	require.Equal(t, "bucket-from-o11y", bucket(param))
	require.Equal(t, "path/from/o11y/", keyPrefix(param))
}