// Package ws implements the client side of RFC 6455 WebSockets.
//
// It is deliberately small: a relay connection needs the opening handshake,
// masked client frames, fragmentation, ping/pong, close, and one bounded
// message reader. Server frames are never masked. One goroutine reads and one
// goroutine writes; WriteText serialises concurrent writers.
package ws

import (
	"bufio"
	"crypto/rand"
	"crypto/sha1"
	"crypto/tls"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"net/url"
	"strings"
	"sync"
	"time"
)

// Frame opcodes.
const (
	OpContinuation = 0x0
	OpText         = 0x1
	OpBinary       = 0x2
	OpClose        = 0x8
	OpPing         = 0x9
	OpPong         = 0xa
)

const acceptGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

var (
	// ErrProtocol is returned for a framing or handshake violation.
	ErrProtocol = errors.New("ws: protocol error")
	// ErrTooLarge is returned when a message exceeds MaxMessage.
	ErrTooLarge = errors.New("ws: message too large")
	// ErrClosed is returned after the peer closed the connection.
	ErrClosed = errors.New("ws: closed")
)

// Dialer opens WebSocket connections.
type Dialer struct {
	// Timeout bounds the TCP connect and TLS handshake, not the WebSocket
	// lifetime. Zero means no bound.
	Timeout time.Duration
	// MaxMessage bounds one (possibly fragmented) message. Zero means 8 MiB.
	MaxMessage int64
	// TLSConfig overrides the default TLS configuration for wss URLs. The
	// ServerName and an empty RootCAs are filled from the URL when unset.
	TLSConfig *tls.Config
	// Header holds extra handshake headers, for example Origin or
	// Authorization. Sec-WebSocket-* headers cannot be overridden.
	Header map[string]string
	// PingInterval sends a WebSocket ping when the connection has been idle
	// for this long. Zero means 45s; a negative value disables keepalive.
	// Relays commonly close idle sockets; a ping is enough to hold the
	// subscription open.
	PingInterval time.Duration
}

// Conn is a client-side WebSocket connection.
type Conn struct {
	conn net.Conn
	br   *bufio.Reader

	wmu  sync.Mutex
	dead bool

	maxMessage int64
	closeOnce  sync.Once
	closeErr   error

	pingStop chan struct{}
}

// Dial opens a WebSocket connection to a ws:// or wss:// URL.
func Dial(rawURL string) (*Conn, error) {
	return (&Dialer{}).Dial(rawURL)
}

// Dial opens a WebSocket connection with d's settings.
func (d *Dialer) Dial(rawURL string) (*Conn, error) {
	u, err := url.Parse(rawURL)
	if err != nil {
		return nil, fmt.Errorf("ws: parse url: %w", err)
	}
	var tlsInUse bool
	switch strings.ToLower(u.Scheme) {
	case "ws":
	case "wss":
		tlsInUse = true
	default:
		return nil, fmt.Errorf("ws: unsupported scheme %q", u.Scheme)
	}
	host := u.Host
	if host == "" {
		return nil, errors.New("ws: url has no host")
	}
	hostname := u.Hostname()
	port := u.Port()
	if port == "" {
		if tlsInUse {
			port = "443"
		} else {
			port = "80"
		}
	}
	addr := net.JoinHostPort(hostname, port)

	var conn net.Conn
	if d.Timeout > 0 {
		conn, err = net.DialTimeout("tcp", addr, d.Timeout)
	} else {
		conn, err = net.Dial("tcp", addr)
	}
	if err != nil {
		return nil, fmt.Errorf("ws: dial %s: %w", addr, err)
	}
	deadline := time.Time{}
	if d.Timeout > 0 {
		deadline = time.Now().Add(d.Timeout)
		conn.SetDeadline(deadline)
	}
	if tlsInUse {
		cfg := &tls.Config{}
		if d.TLSConfig != nil {
			cfg = d.TLSConfig.Clone()
		}
		if cfg.ServerName == "" {
			cfg.ServerName = hostname
		}
		tconn := tls.Client(conn, cfg)
		if err := tconn.Handshake(); err != nil {
			conn.Close()
			return nil, fmt.Errorf("ws: tls handshake: %w", err)
		}
		conn = tconn
	}

	path := u.RequestURI()
	if path == "" {
		path = "/"
	}
	keyBytes := make([]byte, 16)
	if _, err := rand.Read(keyBytes); err != nil {
		conn.Close()
		return nil, fmt.Errorf("ws: random key: %w", err)
	}
	key := base64.StdEncoding.EncodeToString(keyBytes)

	var sb strings.Builder
	sb.WriteString("GET " + path + " HTTP/1.1\r\n")
	sb.WriteString("Host: " + host + "\r\n")
	sb.WriteString("Upgrade: websocket\r\n")
	sb.WriteString("Connection: Upgrade\r\n")
	sb.WriteString("Sec-WebSocket-Key: " + key + "\r\n")
	sb.WriteString("Sec-WebSocket-Version: 13\r\n")
	for name, value := range d.Header {
		lower := strings.ToLower(name)
		if strings.HasPrefix(lower, "sec-websocket-") ||
			lower == "upgrade" || lower == "connection" || lower == "host" {
			continue
		}
		sb.WriteString(name + ": " + value + "\r\n")
	}
	sb.WriteString("\r\n")
	if _, err := conn.Write([]byte(sb.String())); err != nil {
		conn.Close()
		return nil, fmt.Errorf("ws: handshake write: %w", err)
	}

	br := bufio.NewReaderSize(conn, 16<<10)
	status, err := br.ReadString('\n')
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("ws: handshake status: %w", err)
	}
	if !strings.Contains(status, " 101 ") {
		conn.Close()
		return nil, fmt.Errorf("%w: handshake status %q", ErrProtocol, strings.TrimSpace(status))
	}
	headers := make(map[string]string, 16)
	for {
		line, err := br.ReadString('\n')
		if err != nil {
			conn.Close()
			return nil, fmt.Errorf("ws: handshake headers: %w", err)
		}
		line = strings.TrimRight(line, "\r\n")
		if line == "" {
			break
		}
		name, value, ok := strings.Cut(line, ":")
		if !ok {
			conn.Close()
			return nil, fmt.Errorf("%w: malformed header %q", ErrProtocol, line)
		}
		headers[strings.ToLower(strings.TrimSpace(name))] = strings.TrimSpace(value)
	}
	if !headerContains(headers["upgrade"], "websocket") ||
		!headerContains(headers["connection"], "upgrade") {
		conn.Close()
		return nil, fmt.Errorf("%w: missing upgrade headers", ErrProtocol)
	}
	expected := acceptKey(key)
	if headers["sec-websocket-accept"] != expected {
		conn.Close()
		return nil, fmt.Errorf("%w: bad accept key", ErrProtocol)
	}
	if d.Timeout > 0 {
		conn.SetDeadline(time.Time{})
	}
	c := &Conn{conn: conn, br: br, maxMessage: d.MaxMessage}
	interval := d.PingInterval
	if interval == 0 {
		interval = 45 * time.Second
	}
	if interval > 0 {
		c.pingStop = make(chan struct{})
		go c.keepalive(interval)
	}
	return c, nil
}

func (c *Conn) keepalive(interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for {
		select {
		case <-c.pingStop:
			return
		case <-ticker.C:
			if err := c.writeFrame(OpPing, []byte("nostrino")); err != nil {
				c.closeConn()
				return
			}
		}
	}
}

func headerContains(value, token string) bool {
	for _, part := range strings.Split(value, ",") {
		if strings.EqualFold(strings.TrimSpace(part), token) {
			return true
		}
	}
	return false
}

func acceptKey(key string) string {
	h := sha1.Sum([]byte(key + acceptGUID))
	return base64.StdEncoding.EncodeToString(h[:])
}

// WriteText writes one unfragmented text message.
func (c *Conn) WriteText(payload []byte) error {
	return c.writeFrame(OpText, payload)
}

// WriteBinary writes one unfragmented binary message.
func (c *Conn) WriteBinary(payload []byte) error {
	return c.writeFrame(OpBinary, payload)
}

func (c *Conn) writeFrame(opcode byte, payload []byte) error {
	c.wmu.Lock()
	defer c.wmu.Unlock()
	if c.dead {
		return ErrClosed
	}
	var hdr [14]byte
	hdr[0] = 0x80 | opcode
	n := len(payload)
	var hlen int
	switch {
	case n < 126:
		hdr[1] = 0x80 | byte(n)
		hlen = 2
	case n <= 0xffff:
		hdr[1] = 0x80 | 126
		binary.BigEndian.PutUint16(hdr[2:4], uint16(n))
		hlen = 4
	default:
		hdr[1] = 0x80 | 127
		binary.BigEndian.PutUint64(hdr[2:10], uint64(n))
		hlen = 10
	}
	var mask [4]byte
	if _, err := rand.Read(mask[:]); err != nil {
		return fmt.Errorf("ws: mask: %w", err)
	}
	copy(hdr[hlen:hlen+4], mask[:])
	hlen += 4
	buf := make([]byte, hlen+n)
	copy(buf, hdr[:hlen])
	for i := 0; i < n; i++ {
		buf[hlen+i] = payload[i] ^ mask[i&3]
	}
	if _, err := c.conn.Write(buf); err != nil {
		return err
	}
	return nil
}

// ReadMessage reads the next complete text or binary message. Control frames
// are handled internally: ping is answered with pong and close ends the
// connection. A close frame is reported as io.EOF after ErrClosed is stored.
func (c *Conn) ReadMessage() ([]byte, error) {
	var message []byte
	var messageOp byte
	inMessage := false
	for {
		op, fin, payload, err := c.readFrame()
		if err != nil {
			return nil, err
		}
		switch op {
		case OpPing:
			if err := c.writeFrame(OpPong, payload); err != nil {
				return nil, err
			}
			continue
		case OpPong:
			continue
		case OpClose:
			c.writeFrame(OpClose, payload)
			c.closeConn()
			return nil, io.EOF
		case OpText, OpBinary:
			if inMessage {
				return nil, fmt.Errorf("%w: new data frame inside a message", ErrProtocol)
			}
			messageOp = op
			message = append(message[:0], payload...)
			inMessage = true
		case OpContinuation:
			if !inMessage {
				return nil, fmt.Errorf("%w: continuation without a message", ErrProtocol)
			}
			message = append(message, payload...)
		default:
			return nil, fmt.Errorf("%w: unknown opcode %#x", ErrProtocol, op)
		}
		if int64(len(message)) > c.max() {
			return nil, ErrTooLarge
		}
		if fin {
			if messageOp == 0 {
				return nil, fmt.Errorf("%w: empty message", ErrProtocol)
			}
			return message, nil
		}
	}
}

func (c *Conn) max() int64 {
	if c.maxMessage > 0 {
		return c.maxMessage
	}
	return 8 << 20
}

func (c *Conn) readFrame() (op byte, fin bool, payload []byte, err error) {
	var hdr [2]byte
	if _, err = io.ReadFull(c.br, hdr[:]); err != nil {
		return 0, false, nil, err
	}
	fin = hdr[0]&0x80 != 0
	if hdr[0]&0x70 != 0 {
		return 0, false, nil, fmt.Errorf("%w: reserved bits set", ErrProtocol)
	}
	op = hdr[0] & 0x0f
	masked := hdr[1]&0x80 != 0
	length := uint64(hdr[1] & 0x7f)
	switch length {
	case 126:
		var ext [2]byte
		if _, err = io.ReadFull(c.br, ext[:]); err != nil {
			return 0, false, nil, err
		}
		length = uint64(binary.BigEndian.Uint16(ext[:]))
	case 127:
		var ext [8]byte
		if _, err = io.ReadFull(c.br, ext[:]); err != nil {
			return 0, false, nil, err
		}
		length = binary.BigEndian.Uint64(ext[:])
		if length>>63 != 0 {
			return 0, false, nil, fmt.Errorf("%w: negative length", ErrProtocol)
		}
	}
	if length > uint64(c.max()) {
		return 0, false, nil, ErrTooLarge
	}
	if op >= OpClose && (!fin || length > 125) {
		return 0, false, nil, fmt.Errorf("%w: bad control frame", ErrProtocol)
	}
	if masked {
		// A server must not mask frames. Tolerate and unmask rather than fail,
		// because some reverse proxies do it; the RFC violation is harmless here.
		var mask [4]byte
		if _, err = io.ReadFull(c.br, mask[:]); err != nil {
			return 0, false, nil, err
		}
		payload = make([]byte, length)
		if _, err = io.ReadFull(c.br, payload); err != nil {
			return 0, false, nil, err
		}
		for i := range payload {
			payload[i] ^= mask[i&3]
		}
		return op, fin, payload, nil
	}
	payload = make([]byte, length)
	if _, err = io.ReadFull(c.br, payload); err != nil {
		return 0, false, nil, err
	}
	return op, fin, payload, nil
}

// Close sends a close frame and closes the underlying socket.
func (c *Conn) Close() error {
	c.closeOnce.Do(func() {
		if c.pingStop != nil {
			close(c.pingStop)
		}
		c.writeFrame(OpClose, nil)
		c.closeErr = c.closeConn()
	})
	return c.closeErr
}

func (c *Conn) closeConn() error {
	c.wmu.Lock()
	c.dead = true
	c.wmu.Unlock()
	return c.conn.Close()
}

// SetReadDeadline sets the deadline for future reads.
func (c *Conn) SetReadDeadline(t time.Time) error { return c.conn.SetReadDeadline(t) }

// SetWriteDeadline sets the deadline for future writes.
func (c *Conn) SetWriteDeadline(t time.Time) error { return c.conn.SetWriteDeadline(t) }

// SetDeadline sets both deadlines.
func (c *Conn) SetDeadline(t time.Time) error { return c.conn.SetDeadline(t) }