airobot/lib/robot.go

package lib

import (
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gorilla/websocket"
	"github.com/sirupsen/logrus"
	"google.golang.org/protobuf/proto"
	"legu.airobot/pb"
	"legu.airobot/storage"
)

type IRobot interface {
	// 发送消息
	SendMsg(mainType, subType string, req proto.Message, rsp proto.Message) pb.ErrorCode
	// 存储数据
	Store(key string, data interface{})
	// 获取数据
	Get(key string) interface{}
	// 停止运行
	Stop() bool
}

type Robot struct {
	IStore
	account string
	sid     string
	// Deprecated
	scene      *scene
	conn       *websocket.Conn
	data       map[string]interface{} //机器人缓存数据
	status     uint32                 //状态
	lock       sync.Mutex             //
	sceneQueue *Queue[IScene]         //场景队列
	config     *storage.Config
}

func NewRobot(config *storage.Config) *Robot {
	robot := &Robot{
		data:       make(map[string]interface{}),
		sceneQueue: NewQueue[IScene](),
		config:     config,
	}
	robot.Store("sid", config.Global.SId)
	return robot
}

//存数据
func (a *Robot) Store(key string, data interface{}) {
	defer a.lock.Unlock()
	a.lock.Lock()
	a.data[key] = data
}

//取数据
func (a *Robot) Get(key string) interface{} {
	defer a.lock.Unlock()
	a.lock.Lock()
	return a.data[key]
}

// 发送消息
func (r *Robot) SendMsg(mainType, subType string, req proto.Message, rsp proto.Message) pb.ErrorCode {
	start := time.Now()
	defer time.Since(start)

	head := &pb.UserMessage{MainType: mainType, SubType: subType}
	if mainType == "user" && subType == "login" {
		loginReq := req.(*pb.UserLoginReq)
		head.Sec = BuildSecStr(loginReq.Sid, loginReq.Account)
	} else {
		head.Sec = BuildSecStr(r.sid, r.account)
	}

	if ProtoMarshal(req, head) {
		data, _ := proto.Marshal(head)
		if err := r.conn.WriteMessage(websocket.BinaryMessage, data); err != nil {
			logrus.WithField("err", err).Error("写数据异常")
			return pb.ErrorCode_SystemError
		}

		logrus.WithFields(logrus.Fields{"MainType": mainType, "SubType": subType, "req": req}).Debug("发送消息")
		for {
			_, data, err := r.conn.ReadMessage()
			if err != nil {
				logrus.WithField("err", err).Error("读数据异常")
				break
			}

			msg := &pb.UserMessage{}
			if err := proto.Unmarshal(data, msg); err != nil {
				logrus.Error("pb解析失败")
				break
			}

			if msg.MainType == mainType && msg.SubType == subType {
				if !ProtoUnmarshal(msg, rsp) {
					break
				}
				logrus.WithFields(logrus.Fields{"MainType": mainType, "SubType": subType, "rsp": rsp}).Debug("接收消息")
				return pb.ErrorCode_Success
			}
		}
	}

	return pb.ErrorCode_Success
}

// 设置场景队列
func (a *Robot) SetScenes(scenes []IScene) {
	for _, conf := range a.config.Scenes {
		for _, v := range scenes {
			info := v.Info()
			if conf.Name == info.Name {
				a.sceneQueue.Add(v)
				continue
			}
		}
	}
}

// Deprecated
func (m *Robot) SelScene(scene *scene) {
	m.scene = scene
}

// Deprecated
func (m *Robot) GetCurrentScene() *scene {
	return m.scene
}

func (m *Robot) Start() bool {
	if len(m.sceneQueue.List()) == 0 {
		logrus.Warn("没有设置场景队列")
		return false
	}

	// 创建链接
	dialer := &websocket.Dialer{
		HandshakeTimeout: 2 * time.Second,
	}
	conn, _, err := dialer.Dial(m.config.Global.WsAddr, nil)
	if err != nil {
		logrus.Error(err)
		return false
	}

	m.conn = conn

	//检查具备启动的状态
	if !atomic.CompareAndSwapUint32(
		&m.status, STATUS_ORIGINAL, STATUS_STARTING) {
		if !atomic.CompareAndSwapUint32(&m.status, STATUS_STOPPED, STATUS_STARTING) {
			return false
		}
	}

	//设置状态为已启动
	atomic.StoreUint32(&m.status, STATUS_STARTED)

	m.syncCall()

	logrus.Debug("机器人运行了")
	return true
}

func (m *Robot) Stop() bool {
	return false
}

func (m *Robot) syncCall() {
	for {
		scene, err := m.sceneQueue.Pop()
		if err != nil {
			logrus.WithField("err", err).Error("所有场景执行结束")
			return
		}

		start := time.Now()
		//这里执行会花很长时间
		if err := scene.Run(m); err != nil {
			logrus.WithField("err", err).Error("执行业务时发生错误")
			break
		}
		elapsedTime := time.Since(start)
		info := scene.Info()
		logrus.WithField("耗时", elapsedTime.String()).Debug("场景【" + info.Name + "】执行完毕统计")
	}

}

func (m *Robot) checkResp(data []byte, rsp proto.Message) bool {
	msg := &pb.UserMessage{}
	if err := proto.Unmarshal(data, msg); err != nil {
		logrus.Error("pb解析失败")
		return false
	}
	methodStr := msg.Data.TypeUrl
	methodName := SubStr(methodStr, 20, len(methodStr))

	if strings.HasSuffix(methodName, "Push") {
		if methodName == "NotifyErrorNotifyPush" {
			push := &pb.NotifyErrorNotifyPush{}
			if !ProtoUnmarshal(msg, push) {
				logrus.Error("pb解析失败")
				return false
			}
			logrus.WithField("methodName", methodName).WithField("code", push.Code).Debug("收到错误推送")
		}
		return false
	} else {
		if !ProtoUnmarshal(msg, rsp) {
			return false
		}
	}
	return true
}