package nsqd

import (
    "encoding/binary"
    "fmt"
    "io"
    "net"
    "time"

    "github.com/nsqio/go-nsq"
)

// lookupPeer is a low-level type for connecting/reading/writing to nsqlookupd
//
// A lookupPeer instance is designed to connect lazily to nsqlookupd and reconnect
// gracefully (i.e. it is all handled by the library).  Clients can simply use the
// Command interface to perform a round-trip.
type lookupPeer struct {
    l               Logger
    addr            string
    conn            net.Conn
    state           int32
    connectCallback func(*lookupPeer)
    maxBodySize     int64
    Info            peerInfo
}

// peerInfo contains metadata for a lookupPeer instance (and is JSON marshalable)
type peerInfo struct {
    TCPPort          int    `json:"tcp_port"`
    HTTPPort         int    `json:"http_port"`
    Version          string `json:"version"`
    BroadcastAddress string `json:"broadcast_address"`
}

// newLookupPeer creates a new lookupPeer instance connecting to the supplied address.
//
// The supplied connectCallback will be called *every* time the instance connects.
func newLookupPeer(addr string, maxBodySize int64, l Logger, connectCallback func(*lookupPeer)) *lookupPeer {
    return &lookupPeer{
        l:               l,
        addr:            addr,
        state:           stateDisconnected,
        maxBodySize:     maxBodySize,
        connectCallback: connectCallback,
    }
}

// Connect will Dial the specified address, with timeouts
func (lp *lookupPeer) Connect() error {
    lp.l.Output(2, fmt.Sprintf("LOOKUP connecting to %s", lp.addr))
    conn, err := net.DialTimeout("tcp", lp.addr, time.Second)
    if err != nil {
        return err
    }
    lp.conn = conn
    return nil
}

// String returns the specified address
func (lp *lookupPeer) String() string {
    return lp.addr
}

// Read implements the io.Reader interface, adding deadlines
func (lp *lookupPeer) Read(data []byte) (int, error) {
    lp.conn.SetReadDeadline(time.Now().Add(time.Second))
    return lp.conn.Read(data)
}

// Write implements the io.Writer interface, adding deadlines
func (lp *lookupPeer) Write(data []byte) (int, error) {
    lp.conn.SetWriteDeadline(time.Now().Add(time.Second))
    return lp.conn.Write(data)
}

// Close implements the io.Closer interface
func (lp *lookupPeer) Close() error {
    lp.state = stateDisconnected
    if lp.conn != nil {
        return lp.conn.Close()
    }
    return nil
}

// Command performs a round-trip for the specified Command.
//
// It will lazily connect to nsqlookupd and gracefully handle
// reconnecting in the event of a failure.
//
// It returns the response from nsqlookupd as []byte
func (lp *lookupPeer) Command(cmd *nsq.Command) ([]byte, error) {
    initialState := lp.state
    if lp.state != stateConnected {
        err := lp.Connect()
        if err != nil {
            return nil, err
        }
        lp.state = stateConnected
        lp.Write(nsq.MagicV1)
        if initialState == stateDisconnected {
            lp.connectCallback(lp)
        }
    }
    if cmd == nil {
        return nil, nil
    }
    _, err := cmd.WriteTo(lp)
    if err != nil {
        lp.Close()
        return nil, err
    }
    resp, err := readResponseBounded(lp, lp.maxBodySize)
    if err != nil {
        lp.Close()
        return nil, err
    }
    return resp, nil
}

func readResponseBounded(r io.Reader, limit int64) ([]byte, error) {
    var msgSize int32

    // message size
    err := binary.Read(r, binary.BigEndian, &msgSize)
    if err != nil {
        return nil, err
    }

    if int64(msgSize) > limit {
        return nil, fmt.Errorf("response body size (%d) is greater than limit (%d)",
            msgSize, limit)
    }

    // message binary data
    buf := make([]byte, msgSize)
    _, err = io.ReadFull(r, buf)
    if err != nil {
        return nil, err
    }

    return buf, nil
}