package bufferedsocket

import (
	"errors"

	"github.com/gorilla/websocket"
)

type BufferedSocket struct {
	conn       *websocket.Conn
	bufferSize int
	buffer     []byte
	bufferLen  int
}

func NewBufferedSocket(conn *websocket.Conn, bufferSize int) *BufferedSocket {
	return &BufferedSocket{
		conn:       conn,
		bufferSize: bufferSize,
		buffer:     make([]byte, 0, bufferSize),
	}
}

func (bs *BufferedSocket) WriteMessage(message []byte) error {
	if len(message) > bs.bufferSize {
		return errors.New("message larger than buffer size")
	}

	if bs.bufferLen+len(message) > bs.bufferSize {
		err := bs.Flush()
		if err != nil {
			return err
		}
	}
	bs.buffer = append(bs.buffer, message...)
	bs.bufferLen += len(message)
	return nil
}

func (bs *BufferedSocket) Flush() error {
	if bs.bufferLen > 0 {
		err := bs.conn.WriteMessage(websocket.TextMessage, bs.buffer[:bs.bufferLen])
		if err != nil {
			return err
		}
		bs.buffer = make([]byte, 0, bs.bufferSize)
		bs.bufferLen = 0
	}
	return nil
}

func (bs *BufferedSocket) Close() error {
	err := bs.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
	if err != nil {
		return err
	}
	bs.conn.Close()
	return nil
}