Foundry-Scrapping-API/internal/foundry/websocket.go

package foundry

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"strconv"
	"strings"
	"sync"
	"time"

	db_model "gitea.local.lab/Lbenedar/foundry_helper_service/internal/foundry/models/db"
	json_model "gitea.local.lab/Lbenedar/foundry_helper_service/internal/foundry/models/json"
	"gitea.local.lab/Lbenedar/foundry_helper_service/internal/foundry/requests"
	"gitea.local.lab/Lbenedar/foundry_helper_service/internal/foundry/types"
	"github.com/gorilla/websocket"
)

//TODO: refactor all code

type TransportCode int
type FoundryMode int

const (
	WriterCode = TransportCode(0)
	ReaderCode = TransportCode(1)
)

// const (
// 	SetupMode = FoundryMode(0)
// 	WorldMode = FoundryMode(1)
// )

// var

const (
	RespSessionData      = "0"
	RespPingCode         = "2"
	RespSessionId        = "40"
	RespServerChangeCode = "42"
	RespDataCode         = "43"

	ReqPongCode          = "3"
	ReqCreateSessionCode = "40"
	ReqDataCode          = "42"
)

var RequestCodes = []string{
	RespSessionData,
	RespPingCode,
	RespSessionId,
	RespServerChangeCode,
	RespDataCode,
}

var CodesRespToReq = map[string]string{
	RespSessionData:      ReqCreateSessionCode,
	RespPingCode:         ReqPongCode,
	RespServerChangeCode: ReqDataCode,
	//ReqDataCode:  RespServerChangeCode,
}

type webSocketUtil struct {
	wsConn      *websocket.Conn
	currWsId    int
	isReadReady bool
	logger      *slog.Logger
	config      requests.Config

	chanMutex    sync.Mutex
	readChan     types.ReadChannels
	exchangeChan types.ExchangeChannels
	models       db_model.Models
}

func NewWebSocketUtil() *webSocketUtil {
	return &webSocketUtil{
		currWsId:    0,
		isReadReady: false,
		exchangeChan: types.ExchangeChannels{
			Msgs:        make(map[int]chan []byte),
			ProgressMsg: map[string]chan struct{}{},
		},
		config: requests.Config{SessionID: ""}}
}

func parseCode(msg *string) string {
	for j := range RequestCodes {
		if !strings.HasPrefix(*msg, RequestCodes[j]) {
			continue
		}
		return RequestCodes[j]
	}
	return ""
}

func parseId(msg *string, start int) (int, int) {
	msgLen := len(*msg)
	j := start
	for ; j < msgLen; j++ {
		if (*msg)[j] < '0' || (*msg)[j] > '9' {
			break
		}
	}
	if j >= msgLen || j <= 0 {
		return 0, 0
	}

	msgId, err := strconv.Atoi((*msg)[start:j])
	if err != nil {
		return 0, 0
	}
	return msgId, j
}

func parseWsRespMessage(msg string) (*types.WsMessage, error) {
	data := &types.WsMessage{}

	data.Code = parseCode(&msg)

	i := len(data.Code)
	if i == 0 {
		return nil, ErrorMsgNotHaveNumber
	}

	data.Id, i = parseId(&msg, i)
	if i == 0 {
		i = len(data.Code)
	}
	// if i == 0 {
	// 	return nil, ErrorMsgNotHaveNumber
	// }

	data.MsgJson = msg[i:]
	return data, nil
}

// TODO: Lookup timeout
func (ws *webSocketUtil) ReceiveMessage(id int) ([]byte, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	for {
		select {
		case msg, ok := <-ws.exchangeChan.Msgs[id]:
			if !ok {
				time.Sleep(5 * time.Microsecond)
				continue
			}
			ws.closeMsgChannel(id, 0)
			return msg, nil
		case <-ctx.Done():
			err := ctx.Err()
			if err != nil {
				return nil, err
			}
			return nil, ErrorTimeout
		default:
			time.Sleep(5 * time.Microsecond)
		}
	}
}

func (ws *webSocketUtil) closeMsgChannel(id int, timeout time.Duration) bool {
	time.Sleep(timeout)

	ws.chanMutex.Lock()
	defer ws.chanMutex.Unlock()

	ok := true
	if _, ok = ws.exchangeChan.Msgs[id]; !ok {
		return false
	}
	select {
	case _, ok = <-ws.exchangeChan.Msgs[id]:
		if ok {
			close(ws.exchangeChan.Msgs[id])
			delete(ws.exchangeChan.Msgs, id)
		}
	default:
		close(ws.exchangeChan.Msgs[id])
		delete(ws.exchangeChan.Msgs, id)
	}

	ws.logger.Debug("WS: Channel has been closed\n", "id", id, "timeout", timeout.String())
	return ok
}

func (ws *webSocketUtil) HandleWebsocketRequest(msg *types.WsMessage) ([]byte, error) {
	if !ws.IsReadReady() {
		return nil, ErrorIsNotReady
	}

	err := ws.wsConn.WriteMessage(websocket.TextMessage, msg.ToByteSlice())
	if err != nil {
		return nil, err
	}

	data, err := ws.ReceiveMessage(msg.Id)
	if err != nil {
		return nil, err
	}
	return data, nil
}

func (ws *webSocketUtil) sendOnlyCodeRequest(code string) error {
	ws.logger.Debug("WS: Data has been send\n", "msg", CodesRespToReq[code])

	return ws.wsConn.WriteMessage(websocket.TextMessage, []byte(CodesRespToReq[code]))
}

func (ws *webSocketUtil) ListenWebSocket() {
	for {
		_, message, err := ws.wsConn.ReadMessage()
		if err != nil {
			ws.readChan.Err() <- &FoundryError{Type: ReaderCode, Err: err, IsFatal: true}
			return
		}
		data, err := parseWsRespMessage(string(message))
		if err != nil {
			ws.readChan.Err() <- &FoundryError{Type: ReaderCode, Err: err}
			continue
		}
		ws.readChan.Msg() <- data
	}
}

func (ws *webSocketUtil) ServeWebSocket() {
	for {
		select {
		case message := <-ws.readChan.Msg():
			ws.logger.Debug("WS: Data has been received\n", "msgCode", message.Code)

			switch message.Code {
			case RespPingCode, RespSessionData:
				err := ws.sendOnlyCodeRequest(message.Code)
				if err != nil {
					ws.isReadReady = false
					ws.readChan.Err() <- &FoundryError{Type: WriterCode, Err: err}
					continue
				}
			case RespServerChangeCode:
				ws.isReadReady = true
				dataType, data, err := ws.SplitToTypeAndData([]byte(message.MsgJson))
				if err != nil {
					ws.readChan.Err() <- &FoundryError{Type: WriterCode, Err: err}
					continue
				}
				// ws.logger.Debug("Parsed server change message", "dataType", dataType, "data", data)
				ws.ProcessAction(dataType, data)
			case RespDataCode:
				ws.exchangeChan.Msgs[message.Id] = make(chan []byte, 1)
				ws.exchangeChan.Msgs[message.Id] <- []byte(message.MsgJson)
				go ws.closeMsgChannel(message.Id, 5*time.Second)
			default:
			}
		case <-ws.readChan.Done():
			return
		}
	}
}

func (ws *webSocketUtil) IsReadReady() bool {
	return ws.isReadReady
}

func (ws *webSocketUtil) GetJsonData(stateType string) ([]byte, error) {
	msg := ws.CreateWSMessageByPage(stateType)

	return ws.HandleWebsocketRequest(msg)
}

func (ws *webSocketUtil) InsertJsonDataToDB(statePath string) error {
	_, ok1 := requests.PathToSetupState[statePath]
	_, ok2 := requests.PathToWorldState[statePath]
	if !ok1 && !ok2 {
		return ErrorStateTypeNotExist
	}

	msgJson, err := ws.GetJsonData(statePath)
	if err != nil {
		return err
	}

	foundryStateJson, err := json_model.ParseSetupModel(msgJson)
	if err != nil {
		return err
	}

	foundryStateDb := foundryStateJson.GetFoundryStateDB(requests.PathToSetupState[statePath])
	err = ws.models.FoundryState.Insert(foundryStateDb)
	if err != nil {
		return err
	}
	return nil
}

func (ws *webSocketUtil) GetInitialData() {
	err := ws.InitSetupData()
	if err != nil {
		ws.readChan.Err() <- &FoundryError{Type: WriterCode, Err: err}
		return
	}

	idState, err := ws.models.FoundryState.GetIdByType(db_model.SetupState)
	if err != nil {
		ws.readChan.Err() <- &FoundryError{Type: WriterCode, Err: err}
		return
	}
	ws.logger.Debug("Id state", "id", idState)

	err = ws.InitWorldsData(idState)
	if err != nil {
		ws.readChan.Err() <- &FoundryError{Type: WriterCode, Err: err}
		return
	}
}

func (ws *webSocketUtil) InitSetupData() error {
	ws.models.FoundryState.DeleteAll()
	ws.models.FoundryState.DeleteAllSeq()
	for k := range requests.PathToSetupState {
		err := ws.InsertJsonDataToDB(k)
		if err != nil {
			return err
		}
	}
	return nil
}

func (ws *webSocketUtil) InitWorldsData(idState int64) error {
	worlds, err := ws.models.FoundryState.GetWorlds(idState)
	if err != nil {
		return err
	}
	ws.logger.Debug("Worlds", "len", len(worlds))

	for i := range worlds {
		worldName := worlds[i].TextId
		ws.InitWorldData(worldName)
		_, err = ws.config.PostReturnToSetup()
		if err != nil {
			return err
		}
	}
	return nil
}

func (ws *webSocketUtil) InitWorldData(worldName string) error {
	ws.logger.Debug("World name", "name", worldName)
	err := ws.LaunchWorld(worldName)
	if err != nil {
		return err
	}
	for k := range requests.PathToWorldState {
		err := ws.InsertJsonDataToDB(k)
		if err != nil {
			return err
		}
	}
	return nil
}

func (ws *webSocketUtil) LaunchWorld(worldName string) error {
	resp, err := ws.config.PostLaunchWorld(worldName)
	if err != nil {
		return err
	}
	if resp.StatusCode != 200 {
		return errors.New("World is not found")
	}

	ws.exchangeChan.ProgressMsg[worldName] = make(chan struct{})
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()

	select {
	case <-ws.exchangeChan.ProgressMsg[worldName]:
		ws.logger.Debug("World has been started(from GetInitialData)", "world", worldName)
	case <-ctx.Done():
		close(ws.exchangeChan.ProgressMsg[worldName])
		err := ctx.Err()
		if err != nil {
			return err
		}
		return ErrorTimeout
	}
	return nil
}

func (ws *webSocketUtil) CreateNewDataModels(db *sql.DB) {
	ws.models = db_model.NewModels(db)
}

func (ws *webSocketUtil) SplitToTypeAndData(dataJson []byte) (string, any, error) {
	var rawItems []json.RawMessage
	err := json.Unmarshal(dataJson, &rawItems)
	if err != nil {
		return "", nil, err
	}

	var dataType string
	err = json.Unmarshal(rawItems[0], &dataType)
	if err != nil {
		return "", nil, err
	}

	var data map[string]any
	err = json.Unmarshal(rawItems[1], &data)
	if err != nil {
		return "", nil, err
	}

	actionDataPtr := types.GetWsMsgAction(dataType)
	if actionDataPtr == nil {
		return dataType, nil, nil
	}

	err = types.FillStruct(data, actionDataPtr)
	if err != nil {
		return "", nil, err
	}

	return dataType, actionDataPtr, nil
}

func (ws *webSocketUtil) ProcessAction(dataType string, data any) {
	switch tp := data.(type) {
	case *types.WsSessionMsg:
		ws.ProcessSessionAction(tp)
	case *types.WsProgressMsg:
		ws.ProcessProgressAction(tp)
	case *types.WsShutdownMsg:
		ws.ProcessShutdownAction(tp)
	default:
		ws.logger.Debug("default")
	}
}

func (ws *webSocketUtil) ProcessSessionAction(data *types.WsSessionMsg) {
	ws.logger.Debug("ProcessSessionAction", "data", data)
	go ws.GetInitialData()
}

func (ws *webSocketUtil) ProcessProgressAction(data *types.WsProgressMsg) {
	if !data.HasChanged {
		ws.logger.Debug("World has been launched", "world", data.Id)
		select {
		case <-ws.exchangeChan.ProgressMsg[data.Id]:
			ws.logger.Warn("Channel for the world is closed", "world", data.Id)
		default:
			close(ws.exchangeChan.ProgressMsg[data.Id])
		}
	}
}

func (ws *webSocketUtil) ProcessShutdownAction(data *types.WsShutdownMsg) {
	ws.logger.Debug("ProcessShutdownAction", "data", data)
}

func (ws *webSocketUtil) CreateWSMessageByPage(page string) *types.WsMessage {
	msgToSend := &types.WsMessage{Code: CodesRespToReq[RespServerChangeCode], Id: ws.currWsId, MsgJson: fmt.Sprintf("[\"%s\"]", requests.WsTypeData[page])}
	ws.currWsId++

	ws.logger.Debug("WS: Data to send\n", "msgToSend", msgToSend.ToString())
	return msgToSend
}