package web_console

import (
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"errors"
	"net"
	"net/http"
	"net/url"
	"sync"
	"testing"
	"time"

	"github.com/golang-jwt/jwt/v4"
	"github.com/gorilla/websocket"
)

type pipeAddr string

func (a pipeAddr) Network() string { return "pipe" }
func (a pipeAddr) String() string  { return string(a) }

type pipeListener struct {
	conns  chan net.Conn
	closed chan struct{}
	once   sync.Once
}

func newPipeListener() *pipeListener {
	return &pipeListener{
		conns:  make(chan net.Conn, 1),
		closed: make(chan struct{}),
	}
}

func (l *pipeListener) Accept() (net.Conn, error) {
	select {
	case c := <-l.conns:
		return c, nil
	case <-l.closed:
		return nil, net.ErrClosed
	}
}

func (l *pipeListener) Close() error {
	l.once.Do(func() { close(l.closed) })
	return nil
}

func (l *pipeListener) Addr() net.Addr { return pipeAddr("pipe") }

func TestWebSocketReadLimit_RejectsOversizedFirstMessage(t *testing.T) {
	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("ed25519.GenerateKey: %v", err)
	}

	token := jwt.NewWithClaims(jwt.SigningMethodEdDSA, jwt.MapClaims{
		"region": "test",
		"m_ip":   "127.0.0.1",
		"m_user": "root",
		"exp":    time.Now().Add(time.Minute).Unix(),
	})
	signed, err := token.SignedString(priv)
	if err != nil {
		t.Fatalf("token.SignedString: %v", err)
	}

	handler := &sshHandler{
		jwtPublicKey: pub,
		sshIdentityFiles: map[string]string{
			"test": "/nonexistent",
		},
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/ssh-server/web", handler.webSocket)

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

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

	wsURL := "ws://pipe/ssh-server/web?auth=" + url.QueryEscape(signed)
	conn, _, err := dialer.Dial(wsURL, nil)
	if err != nil {
		t.Fatalf("websocket.Dial: %v", err)
	}
	defer conn.Close()

	oversized := make([]byte, maxWSReadLimit+1)
	// Keep the JSON valid and pad with whitespace so the server must read through the whole message
	// (and thus hit the read limit) instead of failing JSON parsing early.
	header := []byte("1:{\"high\":1,\"width\":1}")
	copy(oversized, header)
	for i := len(header); i < len(oversized); i++ {
		oversized[i] = ' '
	}

	_ = conn.WriteMessage(websocket.TextMessage, oversized)

	_ = conn.SetReadDeadline(time.Now().Add(2 * time.Second))
	_, _, err = conn.ReadMessage()
	if err == nil {
		t.Fatalf("expected close error")
	}
	if errors.Is(err, net.ErrClosed) {
		t.Fatalf("expected close frame, got net.ErrClosed: %v", err)
	}
	if !websocket.IsCloseError(err, websocket.CloseMessageTooBig) {
		t.Fatalf("expected CloseMessageTooBig, got %v", err)
	}
}