package pubsubclient

import (
	"fmt"
	"math"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

type connMaker func() (pubsubConn, error)

type client struct {
	newConn connMaker
	conn    pubsubConn

	connStateCond *sync.Cond
	connState     state
	// nReconnectFailures is the number of consecutive times we've
	// failed to connect. Used for the exponential backoff
	// calculation. Should only be accessed while the connStateCond lock
	// is held.
	nReconnectFailures   int
	maxReconnectSleep    time.Duration
	reconnectBackoffUnit time.Duration // exposed only for testing

	subscriptions *subscriptions

	messages   chan *Message
	abortReads chan struct{}

	wg sync.WaitGroup
}

var defaultDialer = &websocket.Dialer{}

// New connects to the pubsub server at url and returns a new client
// connected to that pubsub server. It uses the provided
// *websocket.Dialer, or uses a default dialer if the provided one is
// nil.
//
// The Client will automatically attempt to reconnect to the pubsub
// edge if it becomes disconnected.
func New(url string, dialer *websocket.Dialer) Client {
	if dialer == nil {
		dialer = defaultDialer
	}

	connMaker := func() (pubsubConn, error) {
		conn, _, err := dialer.Dial(url, nil)
		if err != nil {
			return nil, err
		}
		return newPubsubConn(conn), nil
	}

	return newClient(connMaker)
}

func newClient(newConn connMaker) *client {
	return &client{
		newConn:              newConn,
		connState:            waitingToConnect,
		connStateCond:        sync.NewCond(&sync.Mutex{}),
		maxReconnectSleep:    30 * time.Second,
		reconnectBackoffUnit: time.Second,
		subscriptions:        newSubscriptions(),
		messages:             make(chan *Message, 1000),
		abortReads:           make(chan struct{}),
	}
}

func (c *client) reconnect() {
	c.connStateCond.L.Lock()
	defer c.connStateCond.L.Unlock()
	if c.connState != connected {
		return
	}
	c.connect()
}

// Should be run with c.connStateCond.L locked. Returns whether the
// attempt succeeded.
func (c *client) connect() bool {
	defer c.connStateCond.Broadcast()

	c.connState = connecting

	conn, err := c.newConn()
	if err != nil {
		c.conn = nil
		c.connState = disconnected
		return false
	}

	for _, listen := range c.subscriptions.replay() {
		err := conn.send(listen)
		if err != nil {
			warnf("unable to re-establish topic subscriptions after reconnecting: %s", err)
			c.conn = nil
			c.connState = disconnected
			return false
		}
	}
	c.wg.Add(1)
	go c.monitorConnErrors(conn)

	if c.conn != nil {
		c.conn.forwardTo(nil)
	}
	conn.forwardTo(c.messages)
	if c.conn != nil {
		err := c.conn.close()
		if err != nil {
			warnf("dirty shut down while closing connection during a reconnect: %s", err)
		}
	}

	c.conn = conn
	c.connState = connected
	return true
}

func (c *client) monitorConnErrors(conn pubsubConn) {
	defer c.wg.Done()
	for err := range conn.errors() {
		if err == errReconnectRequested {
			// This is a polite request; jitter slightly.
			time.Sleep(randDuration(0, c.maxReconnectSleep))
		} else {
			warnf("low-level pubsub connection error: %q", err)
		}
		c.reconnect()
	}
}

// Execute a function only while a connection is established.
func (c *client) whileConnected(f func(pubsubConn) error) error {
	c.connStateCond.L.Lock()
	defer c.connStateCond.L.Unlock()

	// There are five states the client can be in:
	//
	//  - connected: Everything is normal. We can apply f to the connection.
	//
	//  - disconnected: We unexpectedly lost our connection. Alert the
	//  application of our error, but next time, we should try to
	//  connect.
	//
	//  - waitingToConnect: We don't have a connection, but we're ready
	//  to try to connect. Go for it!
	//
	//  - connecting: Some other goroutine is trying to connect - it
	//  must have hit "waitingToConnect" and is executing
	//  c.connect(). Wait for it to finishh.
	//
	//  - closed: We're permanently shut down.

	// If another goroutine is connecting, wait for them to finish;
	// afterwards, re-check our state and continue the execution.
	for c.connState == connecting {
		c.connStateCond.Wait()
	}

	if c.connState == closed {
		return ErrClosed
	}

	// We're ready to try connecting. Do it!
	if c.connState == waitingToConnect {
		// If we've been failing repeatedly, we should backoff to avoid
		// hammering the backend. Compute the appropriate backoff and
		// sleep.
		time.Sleep(c.preconnectSleep())
		succeeded := c.connect()
		if succeeded {
			c.nReconnectFailures = 0
		} else {
			c.nReconnectFailures++
			warnf("failed to connect %d consecutive times", c.nReconnectFailures)
		}
	}

	// We are disconnected. Immediately error, and mark the client as
	// waitingToConnect; this way, we'll alert the application of the
	// problem by giving them an error, but we'll try to reconnect on
	// their next try.
	if c.connState == disconnected {
		c.connState = waitingToConnect
		return ErrDisconnected
	}

	// The normal, happy case: we're connected.
	if c.connState == connected {
		return f(c.conn)
	}

	panic(fmt.Sprintf("unexpected connection state: %q", c.connState))
}

// Return the amount of time that the client should back off before
// attempting to connect.
func (c *client) preconnectSleep() time.Duration {
	if c.nReconnectFailures == 0 {
		return 0
	}

	// If nReconnectFailures gets even a little high, we run the risk of
	// overflowing, so do the math in terms of the number of attempts:
	// we'll compute the number of reconnect failures that _would_
	// trigger our maxReconnectSleep.
	maxReconnectFailureCount := math.Log2(float64(c.maxReconnectSleep / c.reconnectBackoffUnit))
	if c.nReconnectFailures > int(maxReconnectFailureCount) {
		return c.maxReconnectSleep
	}

	sleepFor := time.Duration(math.Pow(2, float64(c.nReconnectFailures-1))) * c.reconnectBackoffUnit
	if sleepFor > c.maxReconnectSleep {
		return c.maxReconnectSleep
	}
	return sleepFor
}

func (c *client) Subscribe(topic string, authToken string) error {
	msg := &listenMsg{
		msg:   msg{listen},
		Nonce: randString(32),
		Data: topicList{
			Topics:    []string{topic},
			AuthToken: authToken,
		},
	}
	err := c.whileConnected(func(conn pubsubConn) error {
		return conn.send(msg)
	})
	if err == nil {
		c.subscriptions.listen(msg)
	}
	return err
}

func (c *client) Unsubscribe(topic string) error {
	msg := &unlistenMsg{
		msg:   msg{unlisten},
		Nonce: randString(32),
		Data: topicList{
			Topics: []string{topic},
		},
	}
	err := c.whileConnected(func(conn pubsubConn) error {
		return conn.send(msg)
	})
	if err == nil {
		c.subscriptions.unlisten(msg)
	}
	return err
}

func (c *client) Close() error {
	close(c.abortReads)
	err := c.whileConnected(func(conn pubsubConn) error {
		err := conn.close()
		if err != nil {
			return err
		}
		c.connState = closed
		return nil
	})
	if err == nil {
		c.wg.Wait()
		return nil
	}
	return err
}

func (c *client) NextMessage() (*Message, error) {
	var msg *Message
	err := c.whileConnected(func(conn pubsubConn) error {
		select {
		case msg = <-c.messages:
		case <-c.abortReads:
			return ErrClosed
		}
		return nil
	})
	return msg, err
}
