package relay

import (
	"bufio"
	"crypto/rand"
	"encoding/base64"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"strings"
	"testing"
	"time"
)

// testWS is a minimal RFC 6455 client used by the integration tests.
type testWS struct {
	conn net.Conn
	br   *bufio.Reader
}

func dialWS(t *testing.T, addr string) *testWS {
	t.Helper()
	conn, err := net.DialTimeout("tcp", addr, 3*time.Second)
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	keyBytes := make([]byte, 16)
	rand.Read(keyBytes)
	key := base64.StdEncoding.EncodeToString(keyBytes)
	req := "GET / HTTP/1.1\r\nHost: " + addr +
		"\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: " + key +
		"\r\nSec-WebSocket-Version: 13\r\n\r\n"
	if _, err := conn.Write([]byte(req)); err != nil {
		t.Fatalf("handshake write: %v", err)
	}
	br := bufio.NewReader(conn)
	line, err := br.ReadString('\n')
	if err != nil {
		t.Fatalf("handshake read: %v", err)
	}
	if !strings.Contains(line, "101") {
		t.Fatalf("handshake status: %q", line)
	}
	for {
		h, err := br.ReadString('\n')
		if err != nil {
			t.Fatalf("handshake headers: %v", err)
		}
		if h == "\r\n" || h == "\n" {
			break
		}
	}
	return &testWS{conn: conn, br: br}
}

func (c *testWS) writeText(payload []byte) error {
	mask := make([]byte, 4)
	rand.Read(mask)
	hdr := make([]byte, 0, 14)
	hdr = append(hdr, 0x81) // FIN + text
	switch {
	case len(payload) < 126:
		hdr = append(hdr, 0x80|byte(len(payload)))
	case len(payload) <= 0xffff:
		hdr = append(hdr, 0x80|126, byte(len(payload)>>8), byte(len(payload)))
	default:
		hdr = append(hdr, 0x80|127)
		var ext [8]byte
		binary.BigEndian.PutUint64(ext[:], uint64(len(payload)))
		hdr = append(hdr, ext[:]...)
	}
	hdr = append(hdr, mask...)
	body := make([]byte, len(payload))
	for i, b := range payload {
		body[i] = b ^ mask[i&3]
	}
	if _, err := c.conn.Write(append(hdr, body...)); err != nil {
		return err
	}
	return nil
}

func (c *testWS) read() ([]byte, error) {
	c.conn.SetReadDeadline(time.Now().Add(5 * time.Second))
	for {
		hdr := make([]byte, 2)
		if _, err := io.ReadFull(c.br, hdr); err != nil {
			return nil, err
		}
		op := hdr[0] & 0x0f
		length := int64(hdr[1] & 0x7f)
		switch length {
		case 126:
			ext := make([]byte, 2)
			if _, err := io.ReadFull(c.br, ext); err != nil {
				return nil, err
			}
			length = int64(binary.BigEndian.Uint16(ext))
		case 127:
			ext := make([]byte, 8)
			if _, err := io.ReadFull(c.br, ext); err != nil {
				return nil, err
			}
			length = int64(binary.BigEndian.Uint64(ext))
		}
		payload := make([]byte, length)
		if _, err := io.ReadFull(c.br, payload); err != nil {
			return nil, err
		}
		switch op {
		case 0x9: // ping
			if err := c.writeFrame(0xa, payload); err != nil {
				return nil, err
			}
			continue
		case 0x8: // close
			return nil, io.EOF
		case 0x1, 0x2:
			return payload, nil
		default:
			continue
		}
	}
}

func (c *testWS) writeFrame(op byte, payload []byte) error {
	mask := make([]byte, 4)
	rand.Read(mask)
	hdr := []byte{0x80 | op, 0x80 | byte(len(payload))}
	hdr = append(hdr, mask...)
	body := make([]byte, len(payload))
	for i, b := range payload {
		body[i] = b ^ mask[i&3]
	}
	_, err := c.conn.Write(append(hdr, body...))
	return err
}

// expect reads messages until one starts with prefix and returns it.
func (c *testWS) expect(prefix string) string {
	for i := 0; i < 20; i++ {
		msg, err := c.read()
		if err != nil {
			panic(fmt.Sprintf("read while waiting for %q: %v", prefix, err))
		}
		if strings.HasPrefix(string(msg), prefix) {
			return string(msg)
		}
	}
	panic("expected message not seen: " + prefix)
}

func (c *testWS) close() { c.conn.Close() }