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

package foundry

import (
	"context"
	"database/sql"
	"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"
)

type TransportCode int

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

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

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

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

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

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

	msgMap    map[int](chan []byte)
	chanMutex sync.Mutex
	channels  types.Channels
	models    db_model.Models
}

func NewWebSocketUtil() *webSocketUtil {
	return &webSocketUtil{currWsId: 0, isReadReady: false, msgMap: make(map[int]chan []byte)}
}

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)
	if data.Code != RespDataCode {
		return data, nil
	}

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

	data.Id, i = parseId(&msg, i)
	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.msgMap[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.msgMap[id]; !ok {
		return false
	}
	select {
	case _, ok = <-ws.msgMap[id]:
		if ok {
			close(ws.msgMap[id])
			delete(ws.msgMap, id)
		}
	default:
		close(ws.msgMap[id])
		delete(ws.msgMap, 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.channels.Err() <- &FoundryError{Type: ReaderCode, Err: err, IsFatal: true}
			return
		}
		data, err := parseWsRespMessage(string(message))
		if err != nil {
			ws.channels.Err() <- &FoundryError{Type: ReaderCode, Err: err}
			continue
		}
		ws.channels.Msg() <- data
	}
}

func (ws *webSocketUtil) ServeWebSocket() {
	for {
		select {
		case message := <-ws.channels.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.channels.Err() <- &FoundryError{Type: WriterCode, Err: err}
					continue
				}
			case RespCreateSessionCode:
				ws.isReadReady = true
				go ws.GetInitialData()
			case RespDataCode:
				ws.msgMap[message.Id] = make(chan []byte, 1)
				ws.msgMap[message.Id] <- []byte(message.MsgJson)
				go ws.closeMsgChannel(message.Id, 5*time.Second)
			default:
			}
		case <-ws.channels.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() {
	ws.models.FoundryState.DeleteAll()
	ws.models.FoundryState.DeleteAllSeq()
	for k := range requests.PathToSetupState {
		err := ws.InsertJsonDataToDB(k)
		if err != nil {
			ws.channels.Err() <- &FoundryError{Type: WriterCode, Err: err}
			break
		}
	}
	// for k := range requests.PathToWorldState {

	// }
}

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

func (ws *webSocketUtil) CreateWSMessageByPage(page string) *types.WsMessage {
	msgToSend := &types.WsMessage{Code: CodesRespToReq[RespCreateSessionCode], 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
}