package ws

import (
	"encoding/binary"
	"io"
	"unicode/utf8"
)

// Reader parses client frames without writing anything. Control frames are
// returned to the caller so that the socket's writer domain remains the only
// writer on the connection.
type Reader struct {
	br         *bufioReader
	MaxMessage int64
	fragOp     byte
	frag       []byte
	inFrag     bool
}

// bufioReader is the small interface Reader needs; it avoids importing bufio
// in this file's public signature while still accepting *bufio.Reader.
type bufioReader struct {
	ReadFull func([]byte) (int, error)
}

// NewReader wraps a buffered reader.
func NewReader(br interface{ Read([]byte) (int, error) }, max int64) *Reader {
	if max <= 0 {
		max = 1 << 20
	}
	return &Reader{
		br:         &bufioReader{ReadFull: func(p []byte) (int, error) { return io.ReadFull(br, p) }},
		MaxMessage: max,
	}
}

// Read returns the next frame. Data messages are reassembled across
// continuations; ping, pong and close frames are returned immediately.
func (r *Reader) Read() (op byte, payload []byte, err error) {
	for {
		hdr := make([]byte, 2)
		if _, err := r.br.ReadFull(hdr); err != nil {
			return 0, nil, err
		}
		fin := hdr[0]&0x80 != 0
		op := hdr[0] & 0x0f
		if hdr[0]&0x70 != 0 {
			return 0, nil, ErrProtocol
		}
		if hdr[1]&0x80 == 0 {
			return 0, nil, ErrProtocol
		}
		length := int64(hdr[1] & 0x7f)
		switch length {
		case 126:
			ext := make([]byte, 2)
			if _, err := r.br.ReadFull(ext); err != nil {
				return 0, nil, err
			}
			length = int64(binary.BigEndian.Uint16(ext))
		case 127:
			ext := make([]byte, 8)
			if _, err := r.br.ReadFull(ext); err != nil {
				return 0, nil, err
			}
			v := binary.BigEndian.Uint64(ext)
			if v > 1<<62 {
				return 0, nil, ErrTooLarge
			}
			length = int64(v)
		}
		control := op&0x08 != 0
		if control && (!fin || length > 125) {
			return 0, nil, ErrProtocol
		}
		if length > r.MaxMessage || int64(len(r.frag))+length > r.MaxMessage {
			return 0, nil, ErrTooLarge
		}
		mask := make([]byte, 4)
		if _, err := r.br.ReadFull(mask); err != nil {
			return 0, nil, err
		}
		data := make([]byte, length)
		if _, err := r.br.ReadFull(data); err != nil {
			return 0, nil, err
		}
		for i := range data {
			data[i] ^= mask[i&3]
		}

		switch op {
		case OpPing, OpPong, OpClose:
			return op, data, nil
		case OpText, OpBinary:
			if r.inFrag {
				return 0, nil, ErrProtocol
			}
			if fin {
				if op == OpText && !utf8.Valid(data) {
					return 0, nil, ErrProtocol
				}
				return op, data, nil
			}
			r.fragOp = op
			r.frag = append(r.frag[:0], data...)
			r.inFrag = true
		case OpContinuation:
			if !r.inFrag {
				return 0, nil, ErrProtocol
			}
			r.frag = append(r.frag, data...)
			if fin {
				if r.fragOp == OpText && !utf8.Valid(r.frag) {
					return 0, nil, ErrProtocol
				}
				out := append([]byte(nil), r.frag...)
				r.inFrag = false
				r.frag = r.frag[:0]
				return r.fragOp, out, nil
			}
		default:
			return 0, nil, ErrProtocol
		}
	}
}

// AppendFrame appends one unmasked server frame to dst.
func AppendFrame(dst []byte, op byte, payload []byte) []byte {
	dst = append(dst, 0x80|op)
	switch {
	case len(payload) < 126:
		dst = append(dst, byte(len(payload)))
	case len(payload) <= 0xffff:
		dst = append(dst, 126, byte(len(payload)>>8), byte(len(payload)))
	default:
		dst = append(dst, 127)
		var ext [8]byte
		binary.BigEndian.PutUint64(ext[:], uint64(len(payload)))
		dst = append(dst, ext[:]...)
	}
	return append(dst, payload...)
}

// AppendClose appends a close frame with a status code and reason.
func AppendClose(dst []byte, code uint16, reason string) []byte {
	payload := make([]byte, 2+len(reason))
	binary.BigEndian.PutUint16(payload[:2], code)
	copy(payload[2:], reason)
	return AppendFrame(dst, OpClose, payload)
}