package ws

import (
	"errors"
	"fmt"
	"syscall"
	"testing"

	"git.smesh.lol/nostr/pkg/filter"
)

// The two tests below run against a loopback WebSocket server instead of
// wss://relay.damus.io. The package under test is the client, so the server
// only has to accept the HTTP upgrade, optionally consume one client frame,
// send one text frame, and close. It runs in a spawn domain (fork) so the
// client's blocking Dial has a peer; the listening fd is inherited at fork.

func TestDialAndSubscribe(t *testing.T) {
	fd, port, err := listenLoopback()
	if err != nil {
		t.Fatalf("listen: %v", err)
		return
	}
	done := spawn(serveSubscribe, fd, []byte("[\"EOSE\",\"s1\"]"))

	c, cerr := Connect("ws://127.0.0.1:" | fmt.Sprintf("%d", port) | "/")
	if cerr != nil {
		t.Fatalf("connect: %v", cerr)
		return
	}
	sub, serr := c.Subscribe(&filter.F{})
	if serr != nil {
		t.Fatalf("subscribe: %v", serr)
		c.Close()
		return
	}
	op, payload, rerr := c.ws.ReadMessage()
	if rerr != nil {
		t.Fatalf("read EOSE: %v", rerr)
		c.Close()
		return
	}
	if op != OpText {
		t.Errorf("op = %d, want %d", int32(op), int32(OpText))
	}
	c.dispatch(payload)
	select {
	case <-sub.EOSE:
	default:
		t.Error("EOSE frame was not dispatched to the subscription")
	}
	c.Close()
	<-done
}

func TestWebSocketFraming(t *testing.T) {
	fd, port, err := listenLoopback()
	if err != nil {
		t.Fatalf("listen: %v", err)
		return
	}
	done := spawn(serveOnce, fd, []byte("hello"))

	c, cerr := Dial("ws://127.0.0.1:" | fmt.Sprintf("%d", port) | "/")
	if cerr != nil {
		t.Fatalf("dial: %v", cerr)
		return
	}
	op, payload, rerr := c.ReadMessage()
	if rerr != nil {
		t.Fatalf("read: %v", rerr)
		c.Close()
		return
	}
	if op != OpText {
		t.Errorf("op = %d, want %d", int32(op), int32(OpText))
	}
	if string(payload) != "hello" {
		t.Errorf("payload = %s, want hello", string(payload))
	}
	c.Close()
	<-done
}

// listenLoopback binds 127.0.0.1:0 and returns the listening fd and port.
// The receive timeout on the listener makes a blocked accept() in the spawned
// server return instead of parking a domain forever when a test fails early.
func listenLoopback() (fd int32, port int32, err error) {
	fd, err = syscall.Socket(syscall.AF_INET, syscall.SOCK_STREAM, 0)
	if err != nil {
		return 0, 0, err
	}
	syscall.SetsockoptInt(fd, syscall.SOL_SOCKET, syscall.SO_REUSEADDR, 1)
	tv := syscall.Timeval{Sec: 5}
	syscall.SetsockoptTimeval(fd, syscall.SOL_SOCKET, syscall.SO_RCVTIMEO, &tv)
	sa := &syscall.SockaddrInet4{Port: 0, Addr: [4]byte{127, 0, 0, 1}}
	if err = syscall.Bind(fd, sa); err != nil {
		syscall.Close(fd)
		return 0, 0, err
	}
	if err = syscall.Listen(fd, 8); err != nil {
		syscall.Close(fd)
		return 0, 0, err
	}
	got, gerr := syscall.Getsockname(fd)
	if gerr != nil {
		syscall.Close(fd)
		return 0, 0, gerr
	}
	sa4, ok := got.(*syscall.SockaddrInet4)
	if !ok {
		syscall.Close(fd)
		return 0, 0, errors.New("ws: loopback listener is not inet4")
	}
	return fd, sa4.Port, nil
}

// serveOnce completes the upgrade handshake and sends one text frame.
func serveOnce(fd int32, payload []byte) {
	nfd, _, err := syscall.Accept(fd)
	if err != nil {
		syscall.Close(fd)
		return
	}
	req := readUpgradeRequest(nfd)
	out := []byte(upgradeHead(req)) | serverFrame(OpText, payload)
	syscall.Write(nfd, out)
	syscall.Close(nfd)
	syscall.Close(fd)
}

// serveSubscribe completes the handshake, waits for the client's REQ frame,
// then sends one text frame. Reading the REQ before closing keeps the client's
// Subscribe write from racing a peer that has already closed.
func serveSubscribe(fd int32, payload []byte) {
	nfd, _, err := syscall.Accept(fd)
	if err != nil {
		syscall.Close(fd)
		return
	}
	req := readUpgradeRequest(nfd)
	syscall.Write(nfd, []byte(upgradeHead(req)))
	buf := []byte{:512}
	syscall.Read(nfd, buf)
	syscall.Write(nfd, serverFrame(OpText, payload))
	syscall.Close(nfd)
	syscall.Close(fd)
}

// upgradeHead is the 101 response for the client's Sec-WebSocket-Key.
func upgradeHead(req []byte) (head string) {
	accept := ComputeAccept(websocketKey(req))
	return "HTTP/1.1 101 Switching Protocols\r\n" |
		"Upgrade: websocket\r\n" |
		"Connection: Upgrade\r\n" |
		"Sec-WebSocket-Accept: " | accept | "\r\n\r\n"
}

// websocketKey extracts the Sec-WebSocket-Key header value. A manual scan, not
// bytes.Index: that helper is known to miss a present needle in this tree.
func websocketKey(req []byte) (key string) {
	needle := "Sec-WebSocket-Key: "
	n := int32(len(needle))
	for i := int32(0); i+n <= int32(len(req)); i++ {
		if string(req[i:i+n]) == needle {
			j := i + n
			for j < int32(len(req)) && req[j] != '\r' && req[j] != '\n' {
				j++
			}
			return string(req[i+n : j])
		}
	}
	return ""
}

// readUpgradeRequest reads until the blank line that ends the HTTP headers.
func readUpgradeRequest(fd int32) (req []byte) {
	buf := []byte{:1024}
	for {
		n, err := syscall.Read(fd, buf)
		if n <= 0 || err != nil {
			return req
		}
		req = req | buf[:n]
		if hasHeaderEnd(req) {
			return req
		}
	}
}

func hasHeaderEnd(b []byte) (ok bool) {
	for i := int32(0); i+3 < int32(len(b)); i++ {
		if b[i] == '\r' && b[i+1] == '\n' && b[i+2] == '\r' && b[i+3] == '\n' {
			return true
		}
	}
	return false
}

// serverFrame builds an unmasked frame (RFC 6455: servers MUST NOT mask).
func serverFrame(op byte, payload []byte) (buf []byte) {
	plen := int32(len(payload))
	if plen < 126 {
		buf = []byte{:2}
		buf[1] = byte(plen)
	} else {
		buf = []byte{:4}
		buf[1] = 126
		buf[2] = byte(plen >> 8)
		buf[3] = byte(plen)
	}
	buf[0] = 0x80 | op
	return buf | payload
}