package ws

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

// listenTest picks an explicit loopback port. TinyGo's listener reports :0
// rather than the assigned port, so a random explicit port with retries is the
// portable way to get one for a test.
func listenTest(t *testing.T) net.Listener {
	t.Helper()
	for i := 0; i < 40; i++ {
		addr := fmt.Sprintf("127.0.0.1:%d", 20000+mrand.Intn(2000))
		ln, err := net.Listen("tcp", addr)
		if err == nil {
			return ln
		}
	}
	t.Fatal("could not bind a test port")
	return nil
}

func newFramePair(t *testing.T, max int64) (*Conn, net.Conn) {
	t.Helper()
	ln := listenTest(t)
	defer ln.Close()
	type accepted struct {
		c net.Conn
	}
	ch := make(chan net.Conn, 1)
	go func() {
		c, err := ln.Accept()
		if err == nil {
			ch <- c
		}
	}()
	client, err := net.Dial("tcp", ln.Addr().String())
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	server := <-ch
	brw := bufio.NewReadWriter(bufio.NewReader(server), bufio.NewWriter(server))
	c := &Conn{conn: server, br: brw.Reader, bw: brw.Writer, MaxMessage: max}
	t.Cleanup(func() {
		client.Close()
		server.Close()
	})
	return c, client
}

func clientFrame(masked, fin bool, op byte, payload []byte) []byte {
	b0 := op
	if fin {
		b0 |= 0x80
	}
	out := []byte{b0}
	n := len(payload)
	maskBit := byte(0)
	if masked {
		maskBit = 0x80
	}
	switch {
	case n < 126:
		out = append(out, maskBit|byte(n))
	case n <= 0xffff:
		out = append(out, maskBit|126, byte(n>>8), byte(n))
	default:
		out = append(out, maskBit|127)
		var ext [8]byte
		binary.BigEndian.PutUint64(ext[:], uint64(n))
		out = append(out, ext[:]...)
	}
	if !masked {
		return append(out, payload...)
	}
	var key [4]byte
	rand.Read(key[:])
	out = append(out, key[:]...)
	for i, b := range payload {
		out = append(out, b^key[i&3])
	}
	return out
}

func writeClient(t *testing.T, c net.Conn, masked, fin bool, op byte, payload []byte) {
	t.Helper()
	if _, err := c.Write(clientFrame(masked, fin, op, payload)); err != nil {
		t.Fatalf("client write: %v", err)
	}
}

func readServerFrame(t *testing.T, c net.Conn) (op byte, payload []byte) {
	t.Helper()
	c.SetReadDeadline(time.Now().Add(3 * time.Second))
	hdr := make([]byte, 2)
	if _, err := io.ReadFull(c, hdr); err != nil {
		t.Fatalf("server frame header: %v", err)
	}
	if hdr[1]&0x80 != 0 {
		t.Fatal("server frame is masked")
	}
	n := int64(hdr[1] & 0x7f)
	switch n {
	case 126:
		ext := make([]byte, 2)
		io.ReadFull(c, ext)
		n = int64(binary.BigEndian.Uint16(ext))
	case 127:
		ext := make([]byte, 8)
		io.ReadFull(c, ext)
		n = int64(binary.BigEndian.Uint64(ext))
	}
	payload = make([]byte, n)
	if _, err := io.ReadFull(c, payload); err != nil {
		t.Fatalf("server frame payload: %v", err)
	}
	return hdr[0] & 0x0f, payload
}

func TestUpgradeHandshake(t *testing.T) {
	const key = "dGhlIHNhbXBsZSBub25jZQ=="
	headers := map[string]string{
		"upgrade":               "websocket",
		"connection":            "keep-alive, Upgrade",
		"sec-websocket-version": "13",
		"sec-websocket-key":     key,
	}
	server, client := net.Pipe()
	defer server.Close()
	defer client.Close()
	go func() {
		brw := bufio.NewReadWriter(bufio.NewReader(server), bufio.NewWriter(server))
		Upgrade(server, brw, headers)
	}()

	br := bufio.NewReader(client)
	line, err := br.ReadString('\n')
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(line, "101") {
		t.Fatalf("status: %q", line)
	}
	got := map[string]string{}
	for {
		h, err := br.ReadString('\n')
		if err != nil {
			t.Fatal(err)
		}
		if h == "\r\n" {
			break
		}
		parts := strings.SplitN(strings.TrimSpace(h), ":", 2)
		got[strings.ToLower(parts[0])] = strings.TrimSpace(parts[1])
	}
	// RFC 6455 example accept value for this key.
	if got["sec-websocket-accept"] != "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=" {
		t.Fatalf("accept: %q", got["sec-websocket-accept"])
	}
	if got["upgrade"] != "websocket" || got["connection"] != "Upgrade" {
		t.Fatalf("upgrade headers: %+v", got)
	}
}

func TestUpgradeRejectsBadHeaders(t *testing.T) {
	base := map[string]string{
		"upgrade":               "websocket",
		"connection":            "Upgrade",
		"sec-websocket-version": "13",
		"sec-websocket-key":     "dGhlIHNhbXBsZSBub25jZQ==",
	}
	clone := func() map[string]string {
		m := make(map[string]string, len(base))
		for k, v := range base {
			m[k] = v
		}
		return m
	}
	cases := []struct {
		name string
		mod  func(map[string]string)
	}{
		{"no upgrade", func(m map[string]string) { delete(m, "upgrade") }},
		{"wrong upgrade", func(m map[string]string) { m["upgrade"] = "h2c" }},
		{"no connection", func(m map[string]string) { delete(m, "connection") }},
		{"old version", func(m map[string]string) { m["sec-websocket-version"] = "8" }},
		{"no key", func(m map[string]string) { delete(m, "sec-websocket-key") }},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			h := clone()
			tc.mod(h)
			server, client := net.Pipe()
			defer server.Close()
			defer client.Close()
			brw := bufio.NewReadWriter(bufio.NewReader(server), bufio.NewWriter(server))
			if c, err := Upgrade(server, brw, h); err == nil || c != nil {
				t.Fatalf("expected rejection, got conn=%v err=%v", c, err)
			}
		})
	}
}

func TestReadText(t *testing.T) {
	c, client := newFramePair(t, 1<<20)
	writeClient(t, client, true, true, OpText, []byte("hello"))
	op, data, err := c.ReadMessage()
	if err != nil || op != OpText || string(data) != "hello" {
		t.Fatalf("op=%d data=%q err=%v", op, data, err)
	}
}

func TestReadFragmentedText(t *testing.T) {
	c, client := newFramePair(t, 1<<20)
	writeClient(t, client, true, false, OpText, []byte("hel"))
	writeClient(t, client, true, true, OpContinuation, []byte("lo"))
	op, data, err := c.ReadMessage()
	if err != nil || op != OpText || string(data) != "hello" {
		t.Fatalf("op=%d data=%q err=%v", op, data, err)
	}
}

func TestPingPongThenText(t *testing.T) {
	c, client := newFramePair(t, 1<<20)
	writeClient(t, client, true, true, OpPing, []byte("pingpayload"))
	writeClient(t, client, true, true, OpText, []byte("after"))
	// ReadMessage answers the ping before returning the text.
	op, data, err := c.ReadMessage()
	if err != nil || op != OpText || string(data) != "after" {
		t.Fatalf("op=%d data=%q err=%v", op, data, err)
	}
	br := bufio.NewReader(client)
	client.SetReadDeadline(time.Now().Add(2 * time.Second))
	hdr := make([]byte, 2)
	if _, err := io.ReadFull(br, hdr); err != nil {
		t.Fatalf("pong header: %v", err)
	}
	if hdr[0]&0x0f != OpPong {
		t.Fatalf("expected pong, got opcode %d", hdr[0]&0x0f)
	}
	payload := make([]byte, hdr[1]&0x7f)
	io.ReadFull(br, payload)
	if string(payload) != "pingpayload" {
		t.Fatalf("pong payload %q", payload)
	}
}

func TestUnmaskedFrameRejected(t *testing.T) {
	c, client := newFramePair(t, 1<<20)
	writeClient(t, client, false, true, OpText, []byte("unmasked"))
	if _, _, err := c.ReadMessage(); err == nil {
		t.Fatal("unmasked frame accepted")
	}
	op, payload := readServerFrame(t, client)
	if op != OpClose || len(payload) < 2 || binary.BigEndian.Uint16(payload) != 1002 {
		t.Fatalf("expected close 1002, op=%d payload=%x", op, payload)
	}
}

func TestOversizedMessageRejected(t *testing.T) {
	c, client := newFramePair(t, 8)
	writeClient(t, client, true, true, OpText, []byte("this is longer than eight"))
	if _, _, err := c.ReadMessage(); err != ErrTooLarge {
		t.Fatalf("err=%v, want ErrTooLarge", err)
	}
	op, payload := readServerFrame(t, client)
	if op != OpClose || binary.BigEndian.Uint16(payload) != 1009 {
		t.Fatalf("expected close 1009, op=%d payload=%x", op, payload)
	}
}

func TestInvalidUTF8Rejected(t *testing.T) {
	c, client := newFramePair(t, 1<<20)
	writeClient(t, client, true, true, OpText, []byte{0xff, 0xfe})
	if _, _, err := c.ReadMessage(); err == nil {
		t.Fatal("invalid UTF-8 accepted")
	}
	op, payload := readServerFrame(t, client)
	if op != OpClose || binary.BigEndian.Uint16(payload) != 1007 {
		t.Fatalf("expected close 1007, op=%d payload=%x", op, payload)
	}
}

func TestClientClose(t *testing.T) {
	c, client := newFramePair(t, 1<<20)
	payload := []byte{0x03, 0xe8} // 1000, no reason
	writeClient(t, client, true, true, OpClose, payload)
	if _, _, err := c.ReadMessage(); err != ErrClosed {
		t.Fatalf("err=%v, want ErrClosed", err)
	}
	op, reply := readServerFrame(t, client)
	if op != OpClose || len(reply) < 2 || binary.BigEndian.Uint16(reply) != 1000 {
		t.Fatalf("close reply: op=%d payload=%x", op, reply)
	}
}

func TestWriteFrameLengths(t *testing.T) {
	sizes := []int{0, 5, 125, 126, 4096, 70000}
	for _, n := range sizes {
		c, client := newFramePair(t, 1<<20)
		payload := make([]byte, n)
		for i := range payload {
			payload[i] = byte(i * 7)
		}
		if err := c.WriteText(payload); err != nil {
			t.Fatalf("size %d: %v", n, err)
		}
		op, got := readServerFrame(t, client)
		if op != OpText || len(got) != n {
			t.Fatalf("size %d: op=%d len=%d", n, op, len(got))
		}
		for i := range got {
			if got[i] != payload[i] {
				t.Fatalf("size %d: payload mismatch", n)
			}
		}
		client.Close()
	}
}

func TestCloseIsIdempotent(t *testing.T) {
	c, _ := newFramePair(t, 1<<20)
	if err := c.Close(1000, "bye"); err != nil {
		t.Fatal(err)
	}
	if err := c.Close(1000, "again"); err != nil {
		t.Fatalf("second close: %v", err)
	}
	if err := c.WriteText([]byte("x")); err != ErrClosed {
		t.Fatalf("write after close: %v", err)
	}
}

func TestRemoteAddr(t *testing.T) {
	c, _ := newFramePair(t, 1<<20)
	if c.RemoteAddr() == "" {
		t.Fatal("empty remote address")
	}
	c.SetReadDeadline(time.Now().Add(time.Second))
	c.SetWriteDeadline(time.Now().Add(time.Second))
}