// Package pushy is a miniature of Netflix's Pushy: a WebSocket server
// (RFC 6455 handshake + framing, stdlib-only), a push registry mapping
// device -> node, async delivery via a message processor, and graceful
// drain with backoff+jitter. It exists to make the persistent-connection
// proxy model legible and testable.
package pushy

import (
	"bufio"
	"crypto/sha1"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
)

// wsGUID is the RFC 6455 magic constant used to derive the accept key.
const wsGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

// AcceptKey computes Sec-WebSocket-Accept from the client's
// Sec-WebSocket-Key: base64(sha1(key + GUID)). (RFC 6455 §1.3.)
func AcceptKey(key string) string {
	h := sha1.New()
	io.WriteString(h, key+wsGUID)
	return base64.StdEncoding.EncodeToString(h.Sum(nil))
}

// WebSocket opcodes (RFC 6455 §5.2).
const (
	OpContinuation byte = 0x0
	OpText         byte = 0x1
	OpBinary       byte = 0x2
	OpClose        byte = 0x8
	OpPing         byte = 0x9
	OpPong         byte = 0xA
)

// WriteFrame writes a single (final) WebSocket frame. Server->client
// frames are unmasked; client->server frames MUST be masked (RFC 6455
// §5.3) — `mask` controls that, and maskKey supplies the 4-byte key.
func WriteFrame(w io.Writer, opcode byte, payload []byte, mask bool, maskKey [4]byte) error {
	var hdr []byte
	hdr = append(hdr, 0x80|opcode) // FIN=1
	n := len(payload)
	maskBit := byte(0)
	if mask {
		maskBit = 0x80
	}
	switch {
	case n < 126:
		hdr = append(hdr, maskBit|byte(n))
	case n < 65536:
		hdr = append(hdr, maskBit|126)
		var ext [2]byte
		binary.BigEndian.PutUint16(ext[:], uint16(n))
		hdr = append(hdr, ext[:]...)
	default:
		hdr = append(hdr, maskBit|127)
		var ext [8]byte
		binary.BigEndian.PutUint64(ext[:], uint64(n))
		hdr = append(hdr, ext[:]...)
	}
	if mask {
		hdr = append(hdr, maskKey[:]...)
	}
	if _, err := w.Write(hdr); err != nil {
		return err
	}
	if !mask {
		_, err := w.Write(payload)
		return err
	}
	masked := make([]byte, n)
	for i := 0; i < n; i++ {
		masked[i] = payload[i] ^ maskKey[i%4]
	}
	_, err := w.Write(masked)
	return err
}

// ReadFrame reads a single WebSocket frame, returning its opcode, payload
// (unmasked), and the FIN bit.
func ReadFrame(r io.Reader) (opcode byte, payload []byte, fin bool, err error) {
	var h2 [2]byte
	if _, err = io.ReadFull(r, h2[:]); err != nil {
		return
	}
	fin = h2[0]&0x80 != 0
	opcode = h2[0] & 0x0f
	masked := h2[1]&0x80 != 0
	n := int(h2[1] & 0x7f)
	switch n {
	case 126:
		var ext [2]byte
		if _, err = io.ReadFull(r, ext[:]); err != nil {
			return
		}
		n = int(binary.BigEndian.Uint16(ext[:]))
	case 127:
		var ext [8]byte
		if _, err = io.ReadFull(r, ext[:]); err != nil {
			return
		}
		n = int(binary.BigEndian.Uint64(ext[:]))
	}
	var key [4]byte
	if masked {
		if _, err = io.ReadFull(r, key[:]); err != nil {
			return
		}
	}
	payload = make([]byte, n)
	if _, err = io.ReadFull(r, payload); err != nil {
		return
	}
	if masked {
		for i := 0; i < n; i++ {
			payload[i] ^= key[i%4]
		}
	}
	return
}

// Conn is an established WebSocket connection over a net.Conn.
type Conn struct {
	raw      net.Conn
	br       *bufio.Reader
	isServer bool
}

// ServerHandshake performs the RFC 6455 server upgrade on conn: it reads
// the client's HTTP GET, validates the upgrade, and writes the 101
// response. Returns the established Conn.
func ServerHandshake(conn net.Conn) (*Conn, error) {
	br := bufio.NewReader(conn)
	key := ""
	upgrade := false
	for {
		line, err := br.ReadString('\n')
		if err != nil {
			return nil, err
		}
		line = strings.TrimRight(line, "\r\n")
		if line == "" {
			break // end of headers
		}
		l := strings.ToLower(line)
		switch {
		case strings.HasPrefix(l, "sec-websocket-key:"):
			key = strings.TrimSpace(line[len("sec-websocket-key:"):])
		case strings.HasPrefix(l, "upgrade:") && strings.Contains(l, "websocket"):
			upgrade = true
		}
	}
	if !upgrade || key == "" {
		return nil, errors.New("pushy: not a websocket upgrade")
	}
	resp := "HTTP/1.1 101 Switching Protocols\r\n" +
		"Upgrade: websocket\r\nConnection: Upgrade\r\n" +
		"Sec-WebSocket-Accept: " + AcceptKey(key) + "\r\n\r\n"
	if _, err := conn.Write([]byte(resp)); err != nil {
		return nil, err
	}
	return &Conn{raw: conn, br: br, isServer: true}, nil
}

// ClientHandshake performs the client side (used in tests).
func ClientHandshake(conn net.Conn, host, path string) (*Conn, error) {
	// A fixed key keeps tests deterministic; production uses a random nonce.
	key := base64.StdEncoding.EncodeToString([]byte("0123456789abcdef"))
	req := fmt.Sprintf("GET %s HTTP/1.1\r\nHost: %s\r\nUpgrade: websocket\r\n"+
		"Connection: Upgrade\r\nSec-WebSocket-Key: %s\r\nSec-WebSocket-Version: 13\r\n\r\n",
		path, host, key)
	if _, err := conn.Write([]byte(req)); err != nil {
		return nil, err
	}
	br := bufio.NewReader(conn)
	accept := ""
	for {
		line, err := br.ReadString('\n')
		if err != nil {
			return nil, err
		}
		line = strings.TrimRight(line, "\r\n")
		if line == "" {
			break
		}
		if strings.HasPrefix(strings.ToLower(line), "sec-websocket-accept:") {
			accept = strings.TrimSpace(line[len("sec-websocket-accept:"):])
		}
	}
	if accept != AcceptKey(key) {
		return nil, errors.New("pushy: bad accept key")
	}
	return &Conn{raw: conn, br: br, isServer: false}, nil
}

// WriteMessage writes one message. Server side is unmasked; client side
// masks (with a fixed key here for test determinism).
func (c *Conn) WriteMessage(opcode byte, data []byte) error {
	if c.isServer {
		return WriteFrame(c.raw, opcode, data, false, [4]byte{})
	}
	return WriteFrame(c.raw, opcode, data, true, [4]byte{0xde, 0xad, 0xbe, 0xef})
}

// ReadMessage returns the next data message, transparently answering
// pings and surfacing a close as io.EOF.
func (c *Conn) ReadMessage() (opcode byte, data []byte, err error) {
	for {
		op, payload, fin, err := ReadFrame(c.br)
		if err != nil {
			return 0, nil, err
		}
		switch op {
		case OpPing:
			_ = c.WriteMessage(OpPong, payload) // liveness reply
			continue
		case OpPong:
			continue
		case OpClose:
			return OpClose, payload, io.EOF
		}
		if !fin {
			// Minimal continuation handling: read until FIN.
			full := append([]byte{}, payload...)
			for !fin {
				_, more, f, e := ReadFrame(c.br)
				if e != nil {
					return 0, nil, e
				}
				full = append(full, more...)
				fin = f
			}
			return op, full, nil
		}
		return op, payload, nil
	}
}

// Close sends a close frame and closes the underlying connection.
func (c *Conn) Close() error {
	_ = c.WriteMessage(OpClose, nil)
	return c.raw.Close()
}