package iowrappers

import "github.com/gorilla/websocket"

type WebSocketConn struct {
	conn *websocket.Conn
	buf  []byte
}

func NewWebSocketConn(conn *websocket.Conn) *WebSocketConn {
	return &WebSocketConn{conn: conn}
}

func (websocketConn *WebSocketConn) Read(p []byte) (n int, err error) {
	if len(websocketConn.buf) == 0 {
		_, message, err := websocketConn.conn.ReadMessage()
		if err != nil {
			return 0, err
		}
		websocketConn.buf = message
	}

	n = copy(p, websocketConn.buf)
	websocketConn.buf = websocketConn.buf[n:]

	return n, err
}

func (websocketConn *WebSocketConn) Write(p []byte) (n int, err error) {
	err = websocketConn.conn.WriteMessage(websocket.BinaryMessage, p)
	if err == nil {
		n = len(p)
	}
	return n, err
}

func (websocketConn *WebSocketConn) Close() error {
	return websocketConn.conn.Close()
}