package impl

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

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

	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/require"
)

func TestList(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(5)

	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(5)

	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","schema_name":"s1","sum_latency":10,"max_latency":7,"min_latency":3,"exec_count":2,"plan_digest":"p1","table_names":"n1.t1,n1.t2","stmt_type":"Select","normalized_sql":"sql1","first_seen":"2020-09-13T12:39:10+00:00","last_seen":"2020-09-13T12:39:20+00:00","auth_users":{"root":{}}}
{"begin":1600000700,"end":1600000800,"digest":"d2","schema_name":"s2","sum_latency":200,"max_latency":17,"min_latency":13,"exec_count":12,"plan_digest":"p2","table_names":"n2.t1,n3.t2","stmt_type":"Insert","normalized_sql":"sql2","first_seen":"2020-09-13T12:39:10+00:00","last_seen":"2020-09-13T12:39:20+00:00","auth_users":{"root":{}}}`))),
		}, 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":"d1","schema_name":"s1","sum_latency":390,"max_latency":77,"min_latency":13,"exec_count":22,"plan_digest":"g1","table_names":"n1.t1,n1.t2","stmt_type":"Select","normalized_sql":"sql1","first_seen":"2020-09-13T12:56:40+00:00","last_seen":"2020-09-13T12:56:50+00:00","auth_users":{"root":{}}}
{"begin":1600001700,"end":1600002000,"digest":"d2","schema_name":"s2","sum_latency":230,"max_latency":26,"min_latency":11,"exec_count":10,"plan_digest":"p2","table_names":"n2.t1,n3.t2","stmt_type":"Insert","normalized_sql":"sql2","first_seen":"2020-09-13T12:56:40+00:00","last_seen":"2020-09-13T12:56:50+00:00","auth_users":{"root":{}}}`))),
		}, nil).Times(1)

	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":"d1","schema_name":"s1","sum_latency":244,"max_latency":77,"min_latency":13,"exec_count":22,"plan_digest":"p1","table_names":"n1.t1,n1.t2","stmt_type":"Select","normalized_sql":"sql1","first_seen":"2020-09-13T13:02:30+00:00","last_seen":"2020-09-13T13:02:40+00:00","auth_users":{"root":{}}}
{"begin":1600002400,"end":1600002800,"digest":"d2","schema_name":"s2","sum_latency":330,"max_latency":20,"min_latency":17,"exec_count":20,"plan_digest":"p2","table_names":"n2.t1,n3.t2","stmt_type":"Insert","normalized_sql":"sql2","first_seen":"2020-09-13T13:07:30+00:00","last_seen":"2020-09-13T13:07:40+00:00","auth_users":{"root":{}}}`))),
		}, nil).Times(1)

	impl := NewStatementImpl("prod", sf)

	// No more filters
	param := &api.ListParam{
		Provider:     "aws",
		Region:       "us-east-1",
		OrgID:        7777,
		ProjectID:    8888,
		ClusterID:    "9999",
		SharedTier:   true,
		StartUnixSec: 1600000000,
		EndUnixSec:   1600003000,
		Fields: []string{
			"digest_text",
			"sum_latency",
			"avg_latency",
			"max_latency",
			"min_latency",
			"exec_count",
			"plan_count",
		},
	}
	res, err := impl.List(context.Background(), param)
	require.NoError(t, err)

	require.Equal(t, res, []*api.Model{
		{BeginTime: 1600000700, EndTime: 1600002800, Digest: "d2", ExecCount: 42, DigestText: "sql2", SumLatency: 760, MaxLatency: 26, MinLatency: 11, AvgLatency: 18, PlanCount: 1, SchemaName: "s2"},
		{BeginTime: 1600000700, EndTime: 1600002200, Digest: "d1", ExecCount: 46, DigestText: "sql1", SumLatency: 644, MaxLatency: 77, MinLatency: 3, AvgLatency: 14, PlanCount: 2, SchemaName: "s1"},
	})

	// narrow time range
	param = &api.ListParam{
		Provider:     "aws",
		Region:       "us-east-1",
		OrgID:        7777,
		ProjectID:    8888,
		ClusterID:    "9999",
		SharedTier:   true,
		StartUnixSec: 1600001500,
		EndUnixSec:   1600002500,
		Fields: []string{
			"digest_text",
			"sum_latency",
			"avg_latency",
			"max_latency",
			"min_latency",
			"exec_count",
			"plan_count",
		},
	}
	res, err = impl.List(context.Background(), param)
	require.NoError(t, err)

	require.Equal(t, res, []*api.Model{
		{BeginTime: 1600001700, EndTime: 1600002000, Digest: "d1", ExecCount: 22, DigestText: "sql1", SumLatency: 390, MaxLatency: 77, MinLatency: 13, AvgLatency: 17, PlanCount: 1, SchemaName: "s1"},
		{BeginTime: 1600001700, EndTime: 1600002000, Digest: "d2", ExecCount: 10, DigestText: "sql2", SumLatency: 230, MaxLatency: 26, MinLatency: 11, AvgLatency: 23, PlanCount: 1, SchemaName: "s2"},
	})

	// filter by schemas
	param = &api.ListParam{
		Provider:     "aws",
		Region:       "us-east-1",
		OrgID:        7777,
		ProjectID:    8888,
		ClusterID:    "9999",
		SharedTier:   true,
		StartUnixSec: 1600000000,
		EndUnixSec:   1600003000,
		Schemas:      []string{"n1"},
		Fields: []string{
			"digest_text",
			"sum_latency",
			"avg_latency",
			"max_latency",
			"min_latency",
			"exec_count",
			"plan_count",
		},
	}
	res, err = impl.List(context.Background(), param)
	require.NoError(t, err)

	require.Equal(t, res, []*api.Model{
		{BeginTime: 1600000700, EndTime: 1600002200, Digest: "d1", ExecCount: 46, DigestText: "sql1", SumLatency: 644, MaxLatency: 77, MinLatency: 3, AvgLatency: 14, PlanCount: 2, SchemaName: "s1"},
	})

	// filter by stmt types
	param = &api.ListParam{
		Provider:     "aws",
		Region:       "us-east-1",
		OrgID:        7777,
		ProjectID:    8888,
		ClusterID:    "9999",
		SharedTier:   true,
		StartUnixSec: 1600000000,
		EndUnixSec:   1600003000,
		StmtTypes:    []string{"Insert"},
		Fields: []string{
			"digest_text",
			"sum_latency",
			"avg_latency",
			"max_latency",
			"min_latency",
			"exec_count",
			"plan_count",
		},
	}
	res, err = impl.List(context.Background(), param)
	require.NoError(t, err)

	require.Equal(t, res, []*api.Model{
		{BeginTime: 1600000700, EndTime: 1600002800, Digest: "d2", ExecCount: 42, DigestText: "sql2", SumLatency: 760, MaxLatency: 26, MinLatency: 11, AvgLatency: 18, PlanCount: 1, SchemaName: "s2"},
	})

	// search sql
	param = &api.ListParam{
		Provider:     "aws",
		Region:       "us-east-1",
		OrgID:        7777,
		ProjectID:    8888,
		ClusterID:    "9999",
		SharedTier:   true,
		StartUnixSec: 1600000000,
		EndUnixSec:   1600003000,
		Text:         "sql2",
		Fields: []string{
			"digest_text",
			"sum_latency",
			"avg_latency",
			"max_latency",
			"min_latency",
			"exec_count",
			"plan_count",
		},
	}
	res, err = impl.List(context.Background(), param)
	require.NoError(t, err)

	require.Equal(t, res, []*api.Model{
		{BeginTime: 1600000700, EndTime: 1600002800, Digest: "d2", ExecCount: 42, DigestText: "sql2", SumLatency: 760, MaxLatency: 26, MinLatency: 11, AvgLatency: 18, PlanCount: 1, SchemaName: "s2"},
	})
}

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

	sf := mock.NewMockStorageFactory(ctrl)

	sf.EXPECT().
		Storage(gomock.Any(),
			gomock.Eq("aws"),
			gomock.Eq("us-east-1")).
		Return(nil, errors.New("no such storage")).Times(1)

	impl := NewStatementImpl("prod", sf)

	param := &api.ListParam{
		Provider:     "aws",
		Region:       "us-east-1",
		OrgID:        7777,
		ProjectID:    8888,
		ClusterID:    "9999",
		SharedTier:   true,
		StartUnixSec: 1600000000,
		EndUnixSec:   1600003000,
		Fields: []string{
			"digest_text",
			"sum_latency",
			"avg_latency",
			"max_latency",
			"min_latency",
			"exec_count",
			"plan_count",
		},
	}
	_, err := impl.List(context.Background(), param)
	require.ErrorContains(t, err, "load failed")
}

func TestListObjectError(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(1)

	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(1)

	store.EXPECT().
		Object(gomock.Any(),
			gomock.Eq("o11y-prod-shared-us-east-1"),
			gomock.Eq("1600001000-stmt.log")).
		Return(nil, errors.New("network error")).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":"d1","schema_name":"s1","sum_latency":390,"max_latency":77,"min_latency":13,"exec_count":22,"plan_digest":"g1","table_names":"n1.t1,n1.t2","stmt_type":"Select","normalized_sql":"sql1","first_seen":"2020-09-13T12:56:40+00:00","last_seen":"2020-09-13T12:56:50+00:00","auth_users":{"root":{}}}
{"begin":1600001700,"end":1600002000,"digest":"d2","schema_name":"s2","sum_latency":230,"max_latency":26,"min_latency":11,"exec_count":10,"plan_digest":"p2","table_names":"n2.t1,n3.t2","stmt_type":"Insert","normalized_sql":"sql2","first_seen":"2020-09-13T12:56:40+00:00","last_seen":"2020-09-13T12:56:50+00:00","auth_users":{"root":{}}}`))),
		}, nil).Times(1)

	store.EXPECT().
		Object(gomock.Any(),
			gomock.Eq("o11y-prod-shared-us-east-1"),
			gomock.Eq("1600003000-stmt.log")).
		Return(nil, errors.New("auth error")).Times(1)

	impl := NewStatementImpl("prod", sf)

	param := &api.ListParam{
		Provider:     "aws",
		Region:       "us-east-1",
		OrgID:        7777,
		ProjectID:    8888,
		ClusterID:    "9999",
		SharedTier:   true,
		StartUnixSec: 1600000000,
		EndUnixSec:   1600003000,
		Fields: []string{
			"digest_text",
			"sum_latency",
			"avg_latency",
			"max_latency",
			"min_latency",
			"exec_count",
			"plan_count",
		},
	}
	res, err := impl.List(context.Background(), param)
	require.NoError(t, err)
	require.Equal(t, res, []*api.Model{
		{BeginTime: 1600001700, EndTime: 1600002000, Digest: "d1", ExecCount: 22, DigestText: "sql1", SumLatency: 390, MaxLatency: 77, MinLatency: 13, AvgLatency: 17, PlanCount: 1, SchemaName: "s1"},
		{BeginTime: 1600001700, EndTime: 1600002000, Digest: "d2", ExecCount: 10, DigestText: "sql2", SumLatency: 230, MaxLatency: 26, MinLatency: 11, AvgLatency: 23, PlanCount: 1, SchemaName: "s2"},
	})
}