tecamino-driver-artNet/websocket/client.go

package websocket

import (
	"encoding/json"
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	json_dataModels "github.com/tecamino/tecamino-json_data/models"
)

type Client struct {
	conn      *websocket.Conn
	writeMu   sync.Mutex
	OnMessage func(data []byte)
	OnOpen    func()
	OnClose   func(code int, reason string)
	OnError   func(err error)
	OnPing    func()
	OnPong    func()
	timeout   time.Duration
}

// Connect to websocket server
// ip: ip address of server
func NewClient(ip, id string, port uint) (*Client, error) {
	url := fmt.Sprintf("ws://%s:%d/ws?id=%s", ip, port, id)
	c := &Client{}

	dialer := websocket.DefaultDialer
	conn, resp, err := dialer.Dial(url, nil)
	if err != nil {
		if c.OnError != nil {
			c.OnError(err)
		}
		return nil, fmt.Errorf("dial error %v (status %v)", err, resp)
	}
	c.conn = conn

	// Setup control handlers
	c.conn.SetPingHandler(func(appData string) error {
		if c.OnPing != nil {
			c.OnPing()
		}
		return c.conn.WriteMessage(websocket.PongMessage, nil)
	})
	c.conn.SetPongHandler(func(appData string) error {
		c.conn.SetReadDeadline(time.Now().Add(c.timeout))
		if c.OnPong != nil {
			c.OnPong()
		}
		return nil
	})
	c.conn.SetCloseHandler(func(code int, text string) error {
		if c.OnClose != nil {
			c.OnClose(code, text)
		}
		return nil
	})

	if c.OnOpen != nil {
		c.OnOpen()
	}
	return c, nil
}

func (c *Client) Connect(timeout uint) {
	if timeout > 0 {
		fmt.Println(1234, timeout)
		c.timeout = time.Duration(timeout) * time.Second
	}

	go c.pingLoop()

	c.conn.SetReadDeadline(time.Now().Add(c.timeout))
	go func() {

		for {
			msgType, msg, err := c.conn.ReadMessage()
			if err != nil {
				if websocket.IsCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
					log.Println("WebSocket closed:", err)
				}
				if c.OnError != nil {
					c.OnError(fmt.Errorf("read error: %w", err))
				}
				return
			}
			switch msgType {
			case websocket.TextMessage, websocket.BinaryMessage:
				if c.OnMessage != nil {
					c.OnMessage(msg)
				}
			default:
				log.Printf("Unhandled message type: %d", msgType)
			}
		}
	}()
}

func (c *Client) pingLoop() {
	interval := c.timeout / 2
	if interval <= 0 {
		interval = 5 * time.Second
	}
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for range ticker.C {
		if err := c.Write(websocket.PingMessage, nil); err != nil {
			if c.OnError != nil {
				c.OnError(fmt.Errorf("ping error: %w", err))
			}
			return
		}
		if c.OnPing != nil {
			c.OnPing()
		}
	}
}

func (c *Client) Write(msgType int, data []byte) error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()

	c.conn.SetWriteDeadline(time.Now().Add(c.timeout))

	if err := c.conn.WriteMessage(msgType, data); err != nil {
		if c.OnError != nil {
			c.OnError(err)
		}
		return err
	}
	return nil
}

// Close connection to websocket server
func (c *Client) Close(code int, reason string) {
	if c.conn != nil {
		if c.OnClose != nil {
			c.OnClose(code, reason)
		}
		c.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(code, reason))
		c.conn.Close()
	}
}

func (c *Client) SendData(data any) error {
	c.conn.SetWriteDeadline(time.Now().Add(c.timeout))
	if err := c.conn.WriteJSON(data); err != nil {
		if c.OnError != nil {
			c.OnError(err)
		}
		return err
	}
	return nil
}

func (c *Client) ReadJsonData(data []byte) (json_dataModels.Response, error) {
	var resp json_dataModels.Response
	err := json.Unmarshal(data, &resp)
	return resp, err
}