package service

import (
	"context"
	"net"
	"testing"
	"time"

	externalworkloadpb "github.com/pingcap/kvproto/pkg/externalworkloadpb"
	"github.com/stretchr/testify/require"
	scalerpb "github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/proto/scaler"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/test/bufconn"
)

type recordingScalerServer struct {
	scalerpb.UnimplementedScalerServer
	registerGCV2       *scalerpb.RegisterGCV2Request
	recycleBgTask      *scalerpb.RecycleBgTaskRequest
	updateBgTaskExecID *scalerpb.UpdateBgTaskExecIDRequest
	recycleTTLTask     *scalerpb.RecycleTTLTaskRequest
}

func (s *recordingScalerServer) Ping(
	context.Context,
	*scalerpb.PingRequest,
) (*scalerpb.Response, error) {
	return &scalerpb.Response{Error: &scalerpb.Error{
		Type:    scalerpb.ErrorType_PAUSED,
		Message: "paused",
	}}, nil
}

func (s *recordingScalerServer) RegisterGCV2(
	_ context.Context,
	request *scalerpb.RegisterGCV2Request,
) (*scalerpb.Response, error) {
	s.registerGCV2 = request
	return &scalerpb.Response{}, nil
}

func (s *recordingScalerServer) RecycleBgTask(
	_ context.Context,
	request *scalerpb.RecycleBgTaskRequest,
) (*scalerpb.Response, error) {
	s.recycleBgTask = request
	return &scalerpb.Response{}, nil
}

func (s *recordingScalerServer) UpdateBgTaskExecID(
	_ context.Context,
	request *scalerpb.UpdateBgTaskExecIDRequest,
) (*scalerpb.Response, error) {
	s.updateBgTaskExecID = request
	return &scalerpb.Response{}, nil
}

func (s *recordingScalerServer) RecycleTTLTask(
	_ context.Context,
	request *scalerpb.RecycleTTLTaskRequest,
) (*scalerpb.Response, error) {
	s.recycleTTLTask = request
	return &scalerpb.Response{}, nil
}

func TestExternalWorkloadServiceCompatibility(t *testing.T) {
	stub := &recordingScalerServer{}
	listener := bufconn.Listen(1024 * 1024)
	grpcServer := grpc.NewServer()
	scalerpb.RegisterScalerServer(grpcServer, stub)
	externalworkloadpb.RegisterExternalWorkloadControllerServer(
		grpcServer,
		NewExternalWorkloadService(stub),
	)
	go func() { _ = grpcServer.Serve(listener) }()

	conn, err := grpc.NewClient(
		"passthrough:///bufnet",
		grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) {
			return listener.Dial()
		}),
		grpc.WithTransportCredentials(insecure.NewCredentials()),
	)
	require.NoError(t, err)
	t.Cleanup(func() {
		require.NoError(t, conn.Close())
		grpcServer.Stop()
		require.NoError(t, listener.Close())
	})

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	t.Cleanup(cancel)

	legacyResponse, err := scalerpb.NewScalerClient(conn).Ping(ctx, &scalerpb.PingRequest{})
	require.NoError(t, err)
	require.Equal(t, scalerpb.ErrorType_PAUSED, legacyResponse.GetError().GetType())

	externalClient := externalworkloadpb.NewExternalWorkloadControllerClient(conn)
	externalResponse, err := externalClient.Ping(ctx, &externalworkloadpb.PingRequest{})
	require.NoError(t, err)
	require.Equal(t, externalworkloadpb.ErrorType_PAUSED, externalResponse.GetError().GetType())
	require.Equal(t, "paused", externalResponse.GetError().GetMessage())

	header := &externalworkloadpb.RequestHeader{
		KeyspaceId:   42,
		KeyspaceName: "starter-keyspace",
		TidbPool:     "starter-pool",
	}
	_, err = externalClient.RegisterGCV2(ctx, &externalworkloadpb.RegisterGCV2Request{
		Header:     header,
		SafePoint:  123,
		GcLifeTime: 600,
	})
	require.NoError(t, err)
	require.Equal(t, uint32(42), stub.registerGCV2.GetHeader().GetKeyspaceId())
	require.Equal(t, "starter-keyspace", stub.registerGCV2.GetHeader().GetKeyspaceName())
	require.Equal(t, "starter-pool", stub.registerGCV2.GetHeader().GetTidbPool())
	require.Equal(t, uint64(123), stub.registerGCV2.GetSafePoint())
	require.Equal(t, int64(600), stub.registerGCV2.GetGcLifeTime())

	_, err = externalClient.RecycleBgTask(ctx, &externalworkloadpb.RecycleBgTaskRequest{
		Header:       header,
		GlobalTaskId: 11,
		SubTaskId:    12,
	})
	require.NoError(t, err)
	require.Equal(t, int64(11), stub.recycleBgTask.GetGlobalTaskId())
	require.Equal(t, int64(12), stub.recycleBgTask.GetSubTaskId())

	_, err = externalClient.UpdateBgTaskExecID(ctx, &externalworkloadpb.UpdateBgTaskExecIDRequest{
		Header:       header,
		GlobalTaskId: 11,
		Assignments: []*externalworkloadpb.SubtaskExecIDAssignment{
			{SubTaskId: 21, ExecId: "worker-1"},
			{SubTaskId: 22, ExecId: "worker-2"},
		},
	})
	require.NoError(t, err)
	require.Equal(t, []int64{21, 22}, stub.updateBgTaskExecID.GetSubTaskIds())
	require.Equal(t, []string{"worker-1", "worker-2"}, stub.updateBgTaskExecID.GetExecIds())

	_, err = externalClient.RecycleTTLTask(ctx, &externalworkloadpb.RecycleTTLTaskRequest{
		Header:                 header,
		CompletedJobCreateTime: 456,
	})
	require.NoError(t, err)
	require.Equal(t, uint64(456), stub.recycleTTLTask.GetFinishTime())
}