package eventsub

import (
	"errors"
	"log"
	"net/url"
	"time"

	"github.com/gorilla/websocket"
)

const endpoint = "websocket-edge-alb-1716504659.us-west-2.elb.amazonaws.com"

const closeTimeout = 30 * time.Second

type WebSocketClient struct {
	OnReceive       func([]byte) error
	OnReceiveError  func(error)
	OnInternalError func(error)

	done chan struct{}
	conn *websocket.Conn
}

type WebSocketClientOption func(*WebSocketClient)

func OnReceive(f func([]byte) error) WebSocketClientOption {
	return func(c *WebSocketClient) {
		c.OnReceive = f
	}
}

func OnReceiveError(f func(error)) WebSocketClientOption {
	return func(c *WebSocketClient) {
		c.OnReceiveError = f
	}
}

func OnInternalError(f func(error)) WebSocketClientOption {
	return func(c *WebSocketClient) {
		c.OnInternalError = f
	}
}

func NewWebSocketClient(opts ...WebSocketClientOption) *WebSocketClient {
	ws := &WebSocketClient{
		OnReceive:       defaultOnReceive,
		OnReceiveError:  defaultOnReceiveError,
		OnInternalError: defaultOnInternalError,
		done:            make(chan struct{}),
	}

	for _, opt := range opts {
		opt(ws)
	}

	return ws
}

func (ws *WebSocketClient) Start() error {
	u := url.URL{Scheme: "ws", Host: endpoint, Path: "/ws"}

	c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
	if err != nil {
		return err
	}
	ws.conn = c

	go func() {
		defer close(ws.done)
		for {
			_, message, err := c.ReadMessage()
			if _, ok := err.(*websocket.CloseError); ok {
				return
			} else if err != nil {
				ws.OnInternalError(err)
			} else {
				err := ws.OnReceive(message)
				if err != nil {
					ws.OnReceiveError(err)
				}
			}
		}
	}()

	return nil
}

func (ws *WebSocketClient) Stop() error {
	defer ws.conn.Close()

	err := ws.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
	if err != nil {
		return err
	}

	// Wait up to closeTimeout for the backlog of messages to be processed
	select {
	case <-ws.done:
		return nil
	case <-time.After(closeTimeout):
		return errors.New("timeout waiting for graceful stop")
	}
}

func defaultOnReceive(b []byte) error {
	log.Println("Received: ", string(b))
	return nil
}

func defaultOnReceiveError(err error) {
	log.Println("Receive callback error: ", err.Error())
}

func defaultOnInternalError(err error) {
	log.Println("Receive internal error: ", err.Error())
}
