PrimeChatServer/friend/model.go

package friend

import (
	"fmt"
	. "prime/basic"
	auth "prime/user"

	"gorm.io/gorm"
)

type Application struct {
	gorm.Model
	SenderID   uint
	Sender     auth.User
	ReceiverID uint
	Receiver   auth.User
	Status     uint8
	Message    []byte
}

func Migrate() {
	db := DB()
	db.AutoMigrate(&Application{})
	// db.Create(&User{Username: "foo", Password: "bar"})
}

func GetFriends(uid uint) ([]auth.Profile, error) {
	db := DB()
	var friends []auth.User
	var profiles []auth.Profile
	var user auth.User
	res := db.Find(&user, uid)
	if res.RowsAffected == 1 {
		if err := db.Model(&user).Association("Friends").Find(&friends); err != nil {
			return nil, err
		}
		for _, friend := range friends {
			var profile auth.Profile
			res := db.Find(&profile, "user_id = ?", friend.ID)
			if res.Error != nil {
				return nil, res.Error
			}

			profiles = append(profiles, profile)
		}
		return profiles, nil
	}
	return nil, res.Error
}

func IsFriend(a uint, b uint) (bool, error) {
	db := DB()
	var alice, bob auth.User
	res := db.Find(&alice, a)
	if res.RowsAffected == 1 {
		err := db.Model(&alice).Association("Friends").Find(&bob, b)
		return bob.ID == b, err
	}
	return false, res.Error
}

func AddFriend(a uint, b uint) error {
	res, err := IsFriend(a, b)
	if err != nil {
		fmt.Printf("isfriend err %s\n", err.Error())
		return err
	}
	if res {
		return ErrFriendExist
	}
	var alice, bob auth.User
	db := DB()
	resA := db.Take(&alice, a)
	resB := db.Take(&bob, b)
	if resA.RowsAffected == 1 && resB.RowsAffected == 1 {
		alice.Friends = append(alice.Friends, &bob)
		bob.Friends = append(bob.Friends, &alice)
		db.Save(alice)
		db.Save(bob)
		return nil
	}
	return gorm.ErrRecordNotFound
}

func ApplyFriend(from uint, to uint, message []byte) (uint, error) {
	res, err := IsFriend(from, to)
	if err != nil {
		return 0, err
	}
	if res {
		return 0, ErrFriendExist
	}
	db := DB()
	var sender, receiver auth.User
	resA := db.Take(&sender, from)
	resB := db.Take(&receiver, to)
	if resA.RowsAffected != 1 || resB.RowsAffected != 1 {
		return 0, gorm.ErrRecordNotFound
	}
	app := Application{
		Sender: sender, Receiver: receiver, Status: 0, Message: message,
	}
	db.Create(&app)
	// TODO: add a notice ?
	return app.ID, nil
}

func ReplyApplication(uid uint, appID uint, accepted bool) error {
	db := DB()
	var app Application
	res := db.Find(&app, appID)
	fmt.Printf("%+v\n%d\n", app, res.RowsAffected)
	if res.RowsAffected != 1 {
		return res.Error
	}
	if app.ReceiverID != uid {
		return ErrForbidden
	}
	if app.Status != 0 {
		return ErrApplicationAlreadyReplied
	}
	if accepted {
		app.Status = 1
		err := AddFriend(app.SenderID, uid)
		if err != nil {
			fmt.Printf("addfriend error %s\n", err.Error())
			return err
		}
	} else {
		app.Status = 2
	}
	res = db.Save(&app)
	return res.Error
}

func GetApplications(uid uint) ([]Application, error) {
	db := DB()
	var user auth.User
	res := db.Take(&user, uid)
	var apps []Application
	if res.RowsAffected == 1 {
		db.Find(&apps, "sender_id = ? OR receiver_id = ?", uid, uid)
		return apps, nil
	}
	return apps, res.Error
}