package msgtest

import (
	"context"
	"flag"
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/openimsdk/openim-sdk-core/v3/msgtest/module"
	"github.com/openimsdk/openim-sdk-core/v3/sdk_struct"

	"github.com/openimsdk/tools/log"
)

const (
	TenThousandGroupUserNum = 10000
	ThousandGroupUserNum    = 1000
	HundredGroupUserNum     = 100
	FiftyGroupUserNum       = 50

	TenThousandGroupNum = 2
	ThousandGroupNum    = 5
	HundredGroupNum     = 50
	FiftyGroupNum       = 100

	FastenedUserPrefix = "fastened_user_prefix"
	RecvMsgPrefix      = "recv_msg_prefix"
)

var (
	totalOnlineUserNum    int // 总在线用户数
	friendMsgSenderNum    int // 好友消息发送者数
	NotFriendMsgSenderNum int // 非好友消息发送者数
	groupMsgSenderNum     int // 群消息发送者数
	msgSenderNumEvreyUser int // 每个用户的消息数
	fastenedUserNum       int // 固定用户数

	recvMsgUserNum int // 消息接收者数, 抽样账号

)

func InitWithFlag() {
	flag.IntVar(&totalOnlineUserNum, "t", 100000, "total online user num")
	flag.IntVar(&friendMsgSenderNum, "f", 100, "friend msg sender num")
	flag.IntVar(&NotFriendMsgSenderNum, "n", 100, "not friend msg sender num")
	flag.IntVar(&groupMsgSenderNum, "g", 100, "group msg sender num")
	flag.IntVar(&msgSenderNumEvreyUser, "m", 100, "msg sender num evrey user")

	flag.IntVar(&recvMsgUserNum, "r", 20, "recv msg user num")
	flag.IntVar(&fastenedUserNum, "u", 300, "fastened user num")
}

func init() {

	InitWithFlag()

	if err := log.InitFromConfig("sdk.log", "sdk", 4,
		true, false, "./chat_log", 2, 24, "v1.0.0"); err != nil {
		panic(err)
	}
}

func Test_PressureFull(t *testing.T) {
	flag.Parse()
	if friendMsgSenderNum+NotFriendMsgSenderNum+groupMsgSenderNum > totalOnlineUserNum {
		t.Fatal("sender num > total online user num")
	}

	p := NewPressureTester()
	// gen userIDs
	userIDs, fastenedUserIDs, recvMsgUserIDs := p.genUserIDs()

	// register
	if err := p.registerUsers(userIDs, fastenedUserIDs, recvMsgUserIDs); err != nil {
		t.Fatalf("register users failed, err: %v", err)
	}
	// init users
	p.initUserConns(userIDs, fastenedUserIDs)

	// create groups
	err := p.createTestGroups(userIDs, fastenedUserIDs, recvMsgUserIDs)
	if err != nil {
		t.Fatal(err)
	}

	// import friends
	if err := p.importFriends(p.friendSenderUserIDs, fastenedUserIDs); err != nil {
		t.Fatal(err)
	}

	p.pressureSendMsg()
	// send msg test
}

func Test_InitUserConn(t *testing.T) {
	flag.Parse()
	p := NewPressureTester()
	userNum := 10
	// gen userIDs
	userIDs := p.userManager.GenUserIDs(userNum)
	// register
	if err := p.registerUsers(userIDs, nil, nil); err != nil {
		t.Fatalf("register users failed, err: %v", err)
	}
	// init users
	p.initUserConns(userIDs, nil)
}

type PressureTester struct {
	friendManager  *module.TestFriendManager
	userManager    *module.TestUserManager
	groupManager   *module.TestGroupManager
	msgSender      map[string]*module.SendMsgUser
	groupMsgSender map[string]*module.SendMsgUser

	groupSenderUserIDs, friendSenderUserIDs, notfriendSenderUserIDs []string
	recvMsgUserIDs                                                  []string

	tenThousandGroupIDs, thousandGroupIDs, hundredGroupUserIDs, fiftyGroupUserIDs []string
}

func NewPressureTester() *PressureTester {
	metaManager := module.NewMetaManager(APIADDR, SECRET, MANAGERUSERID)
	return &PressureTester{friendManager: metaManager.NewFriendManager(), userManager: metaManager.NewUserManager(), groupManager: metaManager.NewGroupMananger(),
		msgSender: make(map[string]*module.SendMsgUser), groupMsgSender: make(map[string]*module.SendMsgUser)}
}

func (p *PressureTester) genUserIDs() (userIDs, fastenedUserIDs, recvMsgUserIDs []string) {
	userIDs = p.userManager.GenUserIDs(totalOnlineUserNum - fastenedUserNum)                  // 在线用户
	fastenedUserIDs = p.userManager.GenUserIDsWithPrefix(fastenedUserNum, FastenedUserPrefix) // 指定300用户
	recvMsgUserIDs = p.userManager.GenUserIDsWithPrefix(recvMsgUserNum, RecvMsgPrefix)        // 抽样用户完整SDK
	return
}

func (p *PressureTester) registerUsers(userIDs []string, fastenedUserIDs []string, recvMsgUserIDs []string) error {
	for i := 0; i < len(userIDs); i += 1000 {
		end := i + 1000
		if end > len(userIDs) {
			end = len(userIDs)
		}
		userIDsSlice := userIDs[i:end]
		if err := p.userManager.RegisterUsers(userIDsSlice...); err != nil {
			return err
		}
		if len(userIDsSlice) < 1000 {
			break
		}
	}
	if len(fastenedUserIDs) != 0 {
		if err := p.userManager.RegisterUsers(fastenedUserIDs...); err != nil {
			return err
		}
	}
	if len(recvMsgUserIDs) != 0 {
		if err := p.userManager.RegisterUsers(recvMsgUserIDs...); err != nil {
			return err
		}
	}
	return nil
}

func (p *PressureTester) initUserConns(userIDs []string, fastenedUserIDs []string) {
	for i, userID := range userIDs {
		token, err := p.userManager.GetToken(userID, int32(PLATFORMID))
		if err != nil {
			log.ZError(context.Background(), "get token failed", err, "userID", userID, "platformID", PLATFORMID)
			continue
		}
		user := module.NewUser(userID, token, sdk_struct.IMConfig{WsAddr: WSADDR, ApiAddr: APIADDR, PlatformID: int32(PLATFORMID)})
		if 0 <= i && i < friendMsgSenderNum {
			p.msgSender[userID] = user
			p.friendSenderUserIDs = append(p.friendSenderUserIDs, userID)
		} else if friendMsgSenderNum <= i && i < friendMsgSenderNum+NotFriendMsgSenderNum {
			p.msgSender[userID] = user
			p.notfriendSenderUserIDs = append(p.notfriendSenderUserIDs, userID)
		}
	}
	if len(fastenedUserIDs) != 0 {
		for _, userID := range fastenedUserIDs {
			token, err := p.userManager.GetToken(userID, int32(PLATFORMID))
			if err != nil {
				log.ZError(context.Background(), "get token failed", err, "userID", userID, "platformID", PLATFORMID)
				continue
			}
			user := module.NewUser(userID, token, sdk_struct.IMConfig{WsAddr: WSADDR, ApiAddr: APIADDR, PlatformID: int32(PLATFORMID)})
			p.msgSender[userID] = user
			p.groupSenderUserIDs = append(p.groupSenderUserIDs, userID)
		}
	}
}

func (p *PressureTester) createTestGroups(userIDs, fastenedUserIDs, recvMsgUserIDs []string) (err error) {
	// create ten thousand group
	for i := 1; i <= TenThousandGroupNum; i++ {
		groupID := p.groupManager.GenGroupID(fmt.Sprintf("tenThousandGroup_%d", i))
		err = p.groupManager.CreateGroup(groupID, "tenThousandGroup", userIDs[0], append(userIDs[(i-1)*TenThousandGroupUserNum:i*TenThousandGroupUserNum-1], fastenedUserIDs...))
		if err != nil {
			return
		}
		p.tenThousandGroupIDs = append(p.tenThousandGroupIDs, groupID)
	}
	// create two thousand group
	exclude := TenThousandGroupNum * TenThousandGroupUserNum
	for i := 1; i <= ThousandGroupNum; i++ {
		groupID := p.groupManager.GenGroupID(fmt.Sprintf("thousandGroup_%d", i))
		err = p.groupManager.CreateGroup(groupID, "thousandGroup", userIDs[0], append(userIDs[exclude+(i-1)*ThousandGroupUserNum:exclude+i*ThousandGroupUserNum-1], fastenedUserIDs...))
		if err != nil {
			return
		}
		p.thousandGroupIDs = append(p.thousandGroupIDs, groupID)
	}
	// create five hundred group
	exclude += exclude + ThousandGroupNum*ThousandGroupUserNum
	for i := 1; i <= HundredGroupNum; i++ {
		groupID := p.groupManager.GenGroupID(fmt.Sprintf("hundredGroup_%d", i))
		err = p.groupManager.CreateGroup(groupID, "hundredGroup", userIDs[0], append(fastenedUserIDs[0:80], recvMsgUserIDs...))
		if err != nil {
			return
		}
		p.hundredGroupUserIDs = append(p.hundredGroupUserIDs, groupID)
	}
	// create fifty group
	exclude += exclude + HundredGroupNum*HundredGroupUserNum
	for i := 1; i <= FiftyGroupNum; i++ {
		groupID := p.groupManager.GenGroupID(fmt.Sprintf("fiftyGroup_%d", i))
		err = p.groupManager.CreateGroup(groupID, "fiftyGroup", userIDs[0], append(fastenedUserIDs[0:30], recvMsgUserIDs...))
		if err != nil {
			return
		}
		p.fiftyGroupUserIDs = append(p.fiftyGroupUserIDs, groupID)
	}
	return
}

func (p *PressureTester) sendMsgs2Users(senderIDs, recvIDs []string, num int, duration time.Duration) {
	var wg sync.WaitGroup
	for _, senderID := range senderIDs {
		for _, recvID := range recvIDs {
			wg.Add(1)
			go func(senderID, recvID string) {
				defer wg.Done()
				for i := 0; i < num; i++ {
					if user, ok := p.msgSender[senderID]; ok {
						user.SendMsgWithContext(recvID, i)
					}
					time.Sleep(duration)
				}
			}(senderID, recvID)
		}
	}
	wg.Wait()
}

func (p *PressureTester) sendMsgs2Groups(senderIDs, groupIDs []string, num int, duration time.Duration) {
	var wg sync.WaitGroup
	for _, senderID := range senderIDs {
		for _, groupID := range groupIDs {
			wg.Add(1)
			go func(senderID, groupID string) {
				defer wg.Done()
				for i := 0; i < num; i++ {
					if user, ok := p.groupMsgSender[senderID]; ok {
						user.SendGroupMsgWithContext(groupID, i)
					}
					time.Sleep(duration)
				}
			}(senderID, groupID)
		}
	}
	wg.Wait()
}

func (p *PressureTester) pressureSendMsg() {
	var wg sync.WaitGroup
	wg.Add(3)
	go func() {
		defer wg.Done()
		p.sendMsgs2Users(p.friendSenderUserIDs, p.recvMsgUserIDs, msgSenderNumEvreyUser, time.Second)
	}()
	go func() {
		defer wg.Done()
		p.sendMsgs2Users(p.notfriendSenderUserIDs, p.recvMsgUserIDs, msgSenderNumEvreyUser, time.Second)
	}()
	go func() {
		defer wg.Done()
		totalGroupIDs := append(append(p.tenThousandGroupIDs, p.thousandGroupIDs...), append(p.hundredGroupUserIDs, p.fiftyGroupUserIDs...)...)
		p.sendMsgs2Groups(p.groupSenderUserIDs, totalGroupIDs, msgSenderNumEvreyUser, time.Second)
	}()
	wg.Wait()
}

func (p *PressureTester) importFriends(friendSenderUserIDs, recvMsgUserIDs []string) error {
	for _, recvMsgUserID := range recvMsgUserIDs {
		if err := p.friendManager.ImportFriends(recvMsgUserID, friendSenderUserIDs); err != nil {
			return err
		}
	}
	return nil
}
