// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package grpcquerynode

import (
	"context"
	"os"
	"testing"

	"github.com/cockroachdb/errors"

	"github.com/milvus-io/milvus/internal/types"
	"github.com/milvus-io/milvus/internal/util/paramtable"

	"github.com/stretchr/testify/assert"
	clientv3 "go.etcd.io/etcd/client/v3"

	"github.com/milvus-io/milvus-proto/go-api/commonpb"
	"github.com/milvus-io/milvus-proto/go-api/milvuspb"
	"github.com/milvus-io/milvus/internal/proto/internalpb"
	"github.com/milvus-io/milvus/internal/proto/querypb"
)

// /////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
type MockQueryNode struct {
	states     *milvuspb.ComponentStates
	status     *commonpb.Status
	err        error
	initErr    error
	startErr   error
	regErr     error
	stopErr    error
	strResp    *milvuspb.StringResponse
	infoResp   *querypb.GetSegmentInfoResponse
	metricResp *milvuspb.GetMetricsResponse
	configResp *internalpb.ShowConfigurationsResponse
	StatsResp  *internalpb.GetStatisticsResponse
	searchResp *internalpb.SearchResults
	queryResp  *internalpb.RetrieveResults
	distResp   *querypb.GetDataDistributionResponse
}

func (m *MockQueryNode) Init() error {
	return m.initErr
}

func (m *MockQueryNode) Start() error {
	return m.startErr
}

func (m *MockQueryNode) Stop() error {
	return m.stopErr
}

func (m *MockQueryNode) Register() error {
	return m.regErr
}

func (m *MockQueryNode) GetComponentStates(ctx context.Context) (*milvuspb.ComponentStates, error) {
	return m.states, m.err
}

func (m *MockQueryNode) GetStatisticsChannel(ctx context.Context) (*milvuspb.StringResponse, error) {
	return m.strResp, m.err
}

func (m *MockQueryNode) GetTimeTickChannel(ctx context.Context) (*milvuspb.StringResponse, error) {
	return m.strResp, m.err
}

func (m *MockQueryNode) WatchDmChannels(ctx context.Context, req *querypb.WatchDmChannelsRequest) (*commonpb.Status, error) {
	return m.status, m.err
}

func (m *MockQueryNode) LoadSegments(ctx context.Context, req *querypb.LoadSegmentsRequest) (*commonpb.Status, error) {
	return m.status, m.err
}

func (m *MockQueryNode) ReleaseCollection(ctx context.Context, req *querypb.ReleaseCollectionRequest) (*commonpb.Status, error) {
	return m.status, m.err
}

func (m *MockQueryNode) LoadPartitions(ctx context.Context, req *querypb.LoadPartitionsRequest) (*commonpb.Status, error) {
	return m.status, m.err
}

func (m *MockQueryNode) ReleasePartitions(ctx context.Context, req *querypb.ReleasePartitionsRequest) (*commonpb.Status, error) {
	return m.status, m.err
}

func (m *MockQueryNode) ReleaseSegments(ctx context.Context, req *querypb.ReleaseSegmentsRequest) (*commonpb.Status, error) {
	return m.status, m.err
}

func (m *MockQueryNode) GetSegmentInfo(ctx context.Context, req *querypb.GetSegmentInfoRequest) (*querypb.GetSegmentInfoResponse, error) {
	return m.infoResp, m.err
}

func (m *MockQueryNode) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {
	return m.metricResp, m.err
}

func (m *MockQueryNode) GetStatistics(ctx context.Context, req *querypb.GetStatisticsRequest) (*internalpb.GetStatisticsResponse, error) {
	return m.StatsResp, m.err
}

func (m *MockQueryNode) Search(ctx context.Context, req *querypb.SearchRequest) (*internalpb.SearchResults, error) {
	return m.searchResp, m.err
}

func (m *MockQueryNode) Query(ctx context.Context, req *querypb.QueryRequest) (*internalpb.RetrieveResults, error) {
	return m.queryResp, m.err
}

func (m *MockQueryNode) SyncReplicaSegments(ctx context.Context, req *querypb.SyncReplicaSegmentsRequest) (*commonpb.Status, error) {
	return m.status, m.err
}

func (m *MockQueryNode) SetAddress(address string) {
}

func (m *MockQueryNode) GetAddress() string {
	return ""
}

func (m *MockQueryNode) SetEtcdClient(client *clientv3.Client) {
}

func (m *MockQueryNode) UpdateStateCode(code commonpb.StateCode) {
}

func (m *MockQueryNode) SetRootCoord(rc types.RootCoord) error {
	return m.err
}

func (m *MockQueryNode) ShowConfigurations(ctx context.Context, req *internalpb.ShowConfigurationsRequest) (*internalpb.ShowConfigurationsResponse, error) {
	return m.configResp, m.err
}

func (m *MockQueryNode) UnsubDmChannel(ctx context.Context, req *querypb.UnsubDmChannelRequest) (*commonpb.Status, error) {
	return m.status, m.err
}
func (m *MockQueryNode) GetDataDistribution(context.Context, *querypb.GetDataDistributionRequest) (*querypb.GetDataDistributionResponse, error) {
	return m.distResp, m.err
}
func (m *MockQueryNode) SyncDistribution(context.Context, *querypb.SyncDistributionRequest) (*commonpb.Status, error) {
	return m.status, m.err
}

func (m *MockQueryNode) Delete(context.Context, *querypb.DeleteRequest) (*commonpb.Status, error) {
	return m.status, m.err
}

type MockRootCoord struct {
	types.RootCoord
	initErr  error
	startErr error
	regErr   error
	stopErr  error
	stateErr commonpb.ErrorCode
}

func (m *MockRootCoord) Init() error {
	return m.initErr
}

func (m *MockRootCoord) Start() error {
	return m.startErr
}

func (m *MockRootCoord) Stop() error {
	return m.stopErr
}

func (m *MockRootCoord) Register() error {
	return m.regErr
}

func (m *MockRootCoord) SetEtcdClient(client *clientv3.Client) {
}

func (m *MockRootCoord) GetComponentStates(ctx context.Context) (*milvuspb.ComponentStates, error) {
	return &milvuspb.ComponentStates{
		State:  &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Healthy},
		Status: &commonpb.Status{ErrorCode: m.stateErr},
	}, nil
}

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

func TestMain(m *testing.M) {
	paramtable.Init()
	os.Exit(m.Run())
}

func Test_NewServer(t *testing.T) {
	ctx := context.Background()
	server, err := NewServer(ctx, nil)
	assert.Nil(t, err)
	assert.NotNil(t, server)

	mqn := &MockQueryNode{
		states:     &milvuspb.ComponentStates{State: &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Healthy}},
		status:     &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
		err:        nil,
		strResp:    &milvuspb.StringResponse{Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}},
		infoResp:   &querypb.GetSegmentInfoResponse{Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}},
		metricResp: &milvuspb.GetMetricsResponse{Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}},
		configResp: &internalpb.ShowConfigurationsResponse{Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}},
	}
	server.querynode = mqn

	t.Run("Run", func(t *testing.T) {
		err = server.Run()
		assert.Nil(t, err)
	})

	t.Run("GetComponentStates", func(t *testing.T) {
		req := &milvuspb.GetComponentStatesRequest{}
		states, err := server.GetComponentStates(ctx, req)
		assert.Nil(t, err)
		assert.Equal(t, commonpb.StateCode_Healthy, states.State.StateCode)
	})

	t.Run("GetStatisticsChannel", func(t *testing.T) {
		req := &internalpb.GetStatisticsChannelRequest{}
		resp, err := server.GetStatisticsChannel(ctx, req)
		assert.Nil(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
	})

	t.Run("GetTimeTickChannel", func(t *testing.T) {
		req := &internalpb.GetTimeTickChannelRequest{}
		resp, err := server.GetTimeTickChannel(ctx, req)
		assert.Nil(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
	})

	t.Run("WatchDmChannels", func(t *testing.T) {
		req := &querypb.WatchDmChannelsRequest{}
		resp, err := server.WatchDmChannels(ctx, req)
		assert.Nil(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
	})

	t.Run("LoadSegments", func(t *testing.T) {
		req := &querypb.LoadSegmentsRequest{}
		resp, err := server.LoadSegments(ctx, req)
		assert.Nil(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
	})

	t.Run("ReleaseCollection", func(t *testing.T) {
		req := &querypb.ReleaseCollectionRequest{}
		resp, err := server.ReleaseCollection(ctx, req)
		assert.Nil(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
	})

	t.Run("LoadPartitions", func(t *testing.T) {
		req := &querypb.LoadPartitionsRequest{}
		resp, err := server.LoadPartitions(ctx, req)
		assert.Nil(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
	})

	t.Run("ReleasePartitions", func(t *testing.T) {
		req := &querypb.ReleasePartitionsRequest{}
		resp, err := server.ReleasePartitions(ctx, req)
		assert.Nil(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
	})

	t.Run("ReleaseSegments", func(t *testing.T) {
		req := &querypb.ReleaseSegmentsRequest{}
		resp, err := server.ReleaseSegments(ctx, req)
		assert.Nil(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
	})

	t.Run("GetSegmentInfo", func(t *testing.T) {
		req := &querypb.GetSegmentInfoRequest{}
		resp, err := server.GetSegmentInfo(ctx, req)
		assert.Nil(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
	})

	t.Run("GetMetrics", func(t *testing.T) {
		req := &milvuspb.GetMetricsRequest{
			Request: "",
		}
		resp, err := server.GetMetrics(ctx, req)
		assert.Nil(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
	})

	t.Run("Search", func(t *testing.T) {
		req := &querypb.SearchRequest{}
		resp, err := server.Search(ctx, req)
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})

	t.Run("Query", func(t *testing.T) {
		req := &querypb.QueryRequest{}
		resp, err := server.Query(ctx, req)
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})

	t.Run("SyncReplicaSegments", func(t *testing.T) {
		req := &querypb.SyncReplicaSegmentsRequest{}
		resp, err := server.SyncReplicaSegments(ctx, req)
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetErrorCode())
	})

	t.Run("ShowConfigurtaions", func(t *testing.T) {
		req := &internalpb.ShowConfigurationsRequest{
			Pattern: "Cache",
		}
		resp, err := server.ShowConfigurations(ctx, req)
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})

	err = server.Stop()
	assert.Nil(t, err)
}

func Test_Run(t *testing.T) {
	ctx := context.Background()
	server, err := NewServer(ctx, nil)
	assert.Nil(t, err)
	assert.NotNil(t, server)

	server.querynode = &MockQueryNode{startErr: errors.New("Failed")}
	err = server.Run()
	assert.Error(t, err)

	server.querynode = &MockQueryNode{regErr: errors.New("Failed")}
	err = server.Run()
	assert.Error(t, err)

	server.querynode = &MockQueryNode{stopErr: errors.New("Failed")}
	err = server.Stop()
	assert.Error(t, err)
}
