package web_console

import (
	"bytes"
	"context"
	"net"
	"net/http"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

func newWebSocketPair(t *testing.T) (clientConn *websocket.Conn, serverConn *websocket.Conn, cleanup func()) {
	t.Helper()

	serverConnCh := make(chan *websocket.Conn, 1)

	mux := http.NewServeMux()
	mux.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			t.Errorf("upgrader.Upgrade: %v", err)
			return
		}
		serverConnCh <- conn
	})

	ln := newPipeListener()
	srv := &http.Server{Handler: mux}
	go func() {
		_ = srv.Serve(ln)
	}()

	dialer := websocket.Dialer{
		NetDialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			serverNetConn, clientNetConn := net.Pipe()
			select {
			case ln.conns <- serverNetConn:
				return clientNetConn, nil
			case <-ln.closed:
				_ = serverNetConn.Close()
				_ = clientNetConn.Close()
				return nil, net.ErrClosed
			case <-ctx.Done():
				_ = serverNetConn.Close()
				_ = clientNetConn.Close()
				return nil, ctx.Err()
			}
		},
	}

	clientConn, _, err := dialer.Dial("ws://pipe/ws", nil)
	if err != nil {
		_ = srv.Close()
		t.Fatalf("websocket.Dial: %v", err)
	}

	select {
	case serverConn = <-serverConnCh:
	case <-time.After(2 * time.Second):
		_ = clientConn.Close()
		_ = srv.Close()
		t.Fatalf("timed out waiting for server websocket conn")
	}

	cleanup = func() {
		_ = clientConn.Close()
		_ = serverConn.Close()
		_ = srv.Close()
	}
	return clientConn, serverConn, cleanup
}

func TestWSWriteStderr_ForwardsBinaryMessage(t *testing.T) {
	clientConn, serverConn, cleanup := newWebSocketPair(t)
	defer cleanup()

	sshCli := &sshClient{
		conn:     serverConn,
		sessErr:  strings.NewReader("err\n"),
		closeSig: make(chan struct{}, 1),
	}

	done := make(chan error, 1)
	go func() {
		done <- sshCli.wsWriteStderr()
	}()

	_ = clientConn.SetReadDeadline(time.Now().Add(2 * time.Second))
	msgType, payload, err := clientConn.ReadMessage()
	if err != nil {
		t.Fatalf("clientConn.ReadMessage: %v", err)
	}
	if msgType != websocket.BinaryMessage {
		t.Fatalf("expected BinaryMessage, got %d", msgType)
	}
	if string(payload) != "err\n" {
		t.Fatalf("unexpected payload: %q", string(payload))
	}

	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatalf("wsWriteStderr did not exit")
	}
}

func TestWSWrite_SendsLargeChunkInOneFrame(t *testing.T) {
	clientConn, serverConn, cleanup := newWebSocketPair(t)
	defer cleanup()

	payload := bytes.Repeat([]byte("a"), 8*1024)
	sshCli := &sshClient{
		conn:     serverConn,
		sessOut:  bytes.NewReader(payload),
		closeSig: make(chan struct{}, 1),
	}

	done := make(chan error, 1)
	go func() {
		done <- sshCli.wsWrite()
	}()

	_ = clientConn.SetReadDeadline(time.Now().Add(2 * time.Second))
	msgType, got, err := clientConn.ReadMessage()
	if err != nil {
		t.Fatalf("clientConn.ReadMessage: %v", err)
	}
	if msgType != websocket.BinaryMessage {
		t.Fatalf("expected BinaryMessage, got %d", msgType)
	}
	if !bytes.Equal(got, payload) {
		t.Fatalf("unexpected payload length=%d want=%d", len(got), len(payload))
	}

	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatalf("wsWrite did not exit")
	}
}

func TestWriteBinaryMessage_SerializesConcurrentWrites(t *testing.T) {
	clientConn, serverConn, cleanup := newWebSocketPair(t)
	defer cleanup()

	sshCli := &sshClient{
		conn:     serverConn,
		closeSig: make(chan struct{}, 1),
	}

	const n = 32
	var wg sync.WaitGroup
	errCh := make(chan error, n)
	for i := 0; i < n; i++ {
		wg.Add(1)
		payload := []byte{byte(i)}
		go func() {
			defer wg.Done()
			errCh <- sshCli.writeBinaryMessage(payload)
		}()
	}

	got := make(map[byte]struct{}, n)
	_ = clientConn.SetReadDeadline(time.Now().Add(5 * time.Second))
	for len(got) < n {
		msgType, payload, err := clientConn.ReadMessage()
		if err != nil {
			t.Fatalf("clientConn.ReadMessage: %v", err)
		}
		if msgType != websocket.BinaryMessage {
			t.Fatalf("expected BinaryMessage, got %d", msgType)
		}
		if len(payload) != 1 {
			t.Fatalf("expected 1-byte payload, got %d bytes", len(payload))
		}
		got[payload[0]] = struct{}{}
	}

	wg.Wait()
	close(errCh)
	for err := range errCh {
		if err != nil {
			t.Fatalf("writeBinaryMessage: %v", err)
		}
	}

	if len(got) != n {
		t.Fatalf("unexpected message count: got=%d want=%d", len(got), n)
	}
}