package ws

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

func sessionPair(t testing.TB, cfg SessionConfig, h Handlers, start, compression bool) (*Session, *websocket.Conn) {
	t.Helper()
	if err := cfg.Validate(); err != nil {
		t.Fatal(err)
	}
	accepted := make(chan *Session, 1)
	up := websocket.Upgrader{EnableCompression: compression}
	host := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := up.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		s := newSession(conn, cfg, h, context.Background())
		accepted <- s
		if start {
			s.start()
		}
	}))
	dialer := websocket.Dialer{EnableCompression: compression}
	peer, _, err := dialer.Dial("ws"+strings.TrimPrefix(host.URL, "http"), nil)
	if err != nil {
		host.Close()
		t.Fatal(err)
	}
	s := <-accepted
	t.Cleanup(func() {
		s.Abort()
		s.start()
		_ = peer.Close()
		host.Close()
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		_ = s.Wait(ctx)
	})
	return s, peer
}

func sessionTestConfig() SessionConfig {
	return SessionConfig{DisableHeartbeat: true, CloseTimeout: 100 * time.Millisecond, WriteTimeout: time.Second}
}
func waitSession(t *testing.T, s *Session) error {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	err := s.Wait(ctx)
	if errors.Is(err, context.DeadlineExceeded) && ctx.Err() != nil {
		t.Fatal("session did not finish")
	}
	return err
}
func peerClose(t *testing.T, peer *websocket.Conn) {
	t.Helper()
	_ = peer.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(1000, ""), time.Now().Add(time.Second))
	_ = peer.SetReadDeadline(time.Now().Add(time.Second))
	_, _, _ = peer.ReadMessage()
}

func TestSessionEchoAndCompletion(t *testing.T) {
	var opened, closed atomic.Bool
	s, peer := sessionPair(t, sessionTestConfig(), Handlers{
		OnOpen: func(context.Context, *Session) error { opened.Store(true); return nil },
		OnMessage: func(ctx context.Context, s *Session, msg Message) error {
			if !opened.Load() {
				return errors.New("message before open")
			}
			return s.Send(ctx, msg)
		},
		OnClose: func(*Session, error) { closed.Store(true) },
	}, true, false)
	for _, kind := range []int{websocket.TextMessage, websocket.BinaryMessage} {
		if err := peer.WriteMessage(kind, []byte("hello")); err != nil {
			t.Fatal(err)
		}
		got, data, err := peer.ReadMessage()
		if err != nil || got != kind || string(data) != "hello" {
			t.Fatalf("echo: %d %q %v", got, data, err)
		}
	}
	peerClose(t, peer)
	if err := waitSession(t, s); err != nil {
		t.Fatal(err)
	}
	if !closed.Load() {
		t.Fatal("Done preceded OnClose")
	}
	stats := s.Stats()
	if stats.MessagesReceived != 2 || stats.MessagesSent != 2 || stats.QueuedMessages != 0 || stats.ReadQueuedMessages != 0 {
		t.Fatalf("stats: %+v", stats)
	}
}

func TestSessionQueueBoundsCopiesAndGracefulDrain(t *testing.T) {
	cfg := sessionTestConfig()
	cfg.WriteQueueSize = 2
	cfg.WriteQueueBytes = 4
	s, peer := sessionPair(t, cfg, Handlers{}, false, false)
	data := []byte("ab")
	if err := s.TrySend(Message{Type: Text, Data: data}); err != nil {
		t.Fatal(err)
	}
	data[0] = 'z'
	if err := s.TrySend(Message{Type: Text, Data: []byte("cd")}); err != nil {
		t.Fatal(err)
	}
	if err := s.TrySend(Message{Type: Text, Data: []byte("e")}); !errors.Is(err, ErrQueueFull) {
		t.Fatalf("overflow: %v", err)
	}
	if err := s.Close(1000, ""); err != nil {
		t.Fatal(err)
	}
	if err := s.TrySend(Message{Type: Text}); !errors.Is(err, ErrClosed) {
		t.Fatalf("closing send: %v", err)
	}
	s.start()
	for _, want := range []string{"ab", "cd"} {
		_, got, err := peer.ReadMessage()
		if err != nil || string(got) != want {
			t.Fatalf("drain: %q %v", got, err)
		}
	}
	_, _, err := peer.ReadMessage()
	var ce *websocket.CloseError
	if !errors.As(err, &ce) || ce.Code != 1000 {
		t.Fatalf("close: %v", err)
	}
	if err := waitSession(t, s); err != nil {
		t.Fatal(err)
	}
}

func TestSessionCanceledQueueDoesNotCloseConnection(t *testing.T) {
	s, peer := sessionPair(t, sessionTestConfig(), Handlers{}, false, false)
	ctx, cancel := context.WithCancel(context.Background())
	result := make(chan error, 1)
	go func() { result <- s.Send(ctx, Message{Type: Text, Data: []byte("skip")}) }()
	deadline := time.Now().Add(time.Second)
	for s.Stats().QueuedMessages == 0 {
		if time.Now().After(deadline) {
			t.Fatal("send never admitted")
		}
		time.Sleep(time.Millisecond)
	}
	cancel()
	if err := <-result; !errors.Is(err, context.Canceled) {
		t.Fatal(err)
	}
	if err := s.TrySend(Message{Type: Text, Data: []byte("next")}); err != nil {
		t.Fatal(err)
	}
	s.start()
	_, data, err := peer.ReadMessage()
	if err != nil || string(data) != "next" {
		t.Fatalf("got %q %v", data, err)
	}
	peerClose(t, peer)
	if err := waitSession(t, s); err != nil {
		t.Fatal(err)
	}
}

func TestSessionConcurrentSendAndJSON(t *testing.T) {
	s, peer := sessionPair(t, sessionTestConfig(), Handlers{}, true, false)
	const count = 24
	results := make(chan error, count)
	for i := 0; i < count; i++ {
		go func(i int) { results <- s.SendJSON(context.Background(), i) }(i)
	}
	seen := make(map[string]bool)
	for i := 0; i < count; i++ {
		_, data, err := peer.ReadMessage()
		if err != nil {
			t.Fatal(err)
		}
		if seen[string(data)] {
			t.Fatal("duplicate write")
		}
		seen[string(data)] = true
	}
	for i := 0; i < count; i++ {
		if err := <-results; err != nil {
			t.Fatal(err)
		}
	}
	peerClose(t, peer)
	if err := waitSession(t, s); err != nil {
		t.Fatal(err)
	}
}

func TestSessionBoundedDispatcherDoesNotBlockPing(t *testing.T) {
	entered := make(chan struct{})
	release := make(chan struct{})
	var once sync.Once
	s, peer := sessionPair(t, sessionTestConfig(), Handlers{OnMessage: func(ctx context.Context, _ *Session, _ Message) error {
		once.Do(func() { close(entered) })
		select {
		case <-release:
		case <-ctx.Done():
		}
		return nil
	}}, true, false)
	t.Cleanup(func() { close(release) })
	if err := peer.WriteMessage(websocket.TextMessage, []byte("block callback")); err != nil {
		t.Fatal(err)
	}
	<-entered
	pong := make(chan struct{}, 1)
	peer.SetPongHandler(func(data string) error {
		if data == "probe" {
			pong <- struct{}{}
		}
		return nil
	})
	readDone := make(chan struct{})
	go func() { defer close(readDone); _, _, _ = peer.ReadMessage() }()
	if err := peer.WriteControl(websocket.PingMessage, []byte("probe"), time.Now().Add(time.Second)); err != nil {
		t.Fatal(err)
	}
	select {
	case <-pong:
	case <-time.After(time.Second):
		t.Fatal("blocked callback starved control frames")
	}
	_ = s.Close(1000, "")
	<-readDone
	if err := waitSession(t, s); err != nil {
		t.Fatal(err)
	}
}

func TestSessionInboundQueueFullCancelsCallback(t *testing.T) {
	cfg := sessionTestConfig()
	cfg.ReadQueueSize = 1
	cfg.ReadQueueBytes = 8
	entered := make(chan struct{})
	s, peer := sessionPair(t, cfg, Handlers{OnMessage: func(ctx context.Context, _ *Session, _ Message) error { close(entered); <-ctx.Done(); return nil }}, true, false)
	_ = peer.WriteMessage(websocket.TextMessage, []byte("one"))
	<-entered
	_ = peer.WriteMessage(websocket.TextMessage, []byte("two"))
	_, _, err := peer.ReadMessage()
	var ce *websocket.CloseError
	if !errors.As(err, &ce) || ce.Code != 1008 {
		t.Fatalf("close: %v", err)
	}
	if err := waitSession(t, s); !errors.Is(err, ErrQueueFull) {
		t.Fatalf("session: %v", err)
	}
}

func TestSessionCompressedReadLimitAndInvalidUTF8(t *testing.T) {
	for _, tc := range []struct {
		name  string
		data  []byte
		code  int
		cause error
	}{
		{"compressed", []byte(strings.Repeat("x", 4096)), 1009, ErrMessageTooLarge},
		{"utf8", []byte{0xff}, 1007, ErrInvalidMessage},
	} {
		t.Run(tc.name, func(t *testing.T) {
			cfg := sessionTestConfig()
			cfg.MaxMessageBytes = 128
			s, peer := sessionPair(t, cfg, Handlers{}, true, true)
			if err := peer.WriteMessage(websocket.TextMessage, tc.data); err != nil {
				t.Fatal(err)
			}
			_, _, err := peer.ReadMessage()
			var ce *websocket.CloseError
			if !errors.As(err, &ce) || ce.Code != tc.code {
				t.Fatalf("close: %v", err)
			}
			if err := waitSession(t, s); !errors.Is(err, tc.cause) {
				t.Fatalf("limit: %v", err)
			}
		})
	}
}

func TestSessionCallbackPanicAndCloseBarrier(t *testing.T) {
	closeEntered := make(chan struct{})
	release := make(chan struct{})
	s, peer := sessionPair(t, sessionTestConfig(), Handlers{
		OnMessage: func(context.Context, *Session, Message) error { panic("secret peer data") },
		OnClose:   func(*Session, error) { close(closeEntered); <-release },
	}, true, false)
	_ = peer.WriteMessage(websocket.TextMessage, []byte("panic"))
	_, _, _ = peer.ReadMessage()
	<-closeEntered
	select {
	case <-s.Done():
		t.Fatal("Done ignored running OnClose")
	default:
	}
	close(release)
	if err := waitSession(t, s); err == nil || strings.Contains(err.Error(), "secret") {
		t.Fatalf("panic error: %v", err)
	}
}

func TestSessionHeartbeatDetectsUnresponsivePeer(t *testing.T) {
	cfg := SessionConfig{PingInterval: 10 * time.Millisecond, PongTimeout: 30 * time.Millisecond, CloseTimeout: 30 * time.Millisecond}
	s, _ := sessionPair(t, cfg, Handlers{}, true, false)
	if err := waitSession(t, s); err == nil {
		t.Fatal("peer without reader remained healthy")
	}
}

func TestSessionWriteCancellationStopsConnection(t *testing.T) {
	cfg := sessionTestConfig()
	cfg.MaxMessageBytes = 16 << 20
	cfg.WriteQueueBytes = 16 << 20
	s, _ := sessionPair(t, cfg, Handlers{}, true, false)
	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	err := s.Send(ctx, Message{Type: Binary, Data: make([]byte, 16<<20)})
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("send: %v", err)
	}
	if err := waitSession(t, s); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("write cancellation: %v", err)
	}
}

func TestSessionValidationAndRedactedErrors(t *testing.T) {
	for _, cfg := range []SessionConfig{{MaxMessageBytes: -1}, {ReadQueueSize: -1}, {WriteQueueBytes: -1}, {WriteTimeout: -1}, {PingInterval: time.Second, PongTimeout: time.Second}} {
		if cfg.Validate() == nil {
			t.Fatalf("invalid config accepted: %+v", cfg)
		}
	}
	s, _ := sessionPair(t, sessionTestConfig(), Handlers{}, false, false)
	for _, msg := range []Message{{Type: 9}, {Type: Text, Data: []byte{0xff}}} {
		if !errors.Is(s.TrySend(msg), ErrInvalidMessage) {
			t.Fatal("invalid message accepted")
		}
	}
	if !errors.Is(s.TrySend(Message{Type: Binary, Data: make([]byte, (1<<20)+1)}), ErrMessageTooLarge) {
		t.Fatal("large outbound accepted")
	}
	for _, code := range []int{1005, 1006, 1015, 5000} {
		if s.Close(code, "") == nil {
			t.Fatalf("invalid close %d", code)
		}
	}
	secret := fmt.Errorf("https://user:password@example.test?token=secret")
	for _, err := range []error{&OpError{Op: "dial", Cause: secret}, &HandshakeError{StatusCode: 401, Cause: secret}, &CloseError{Code: 1008, Cause: secret}} {
		if strings.Contains(err.Error(), "secret") || !errors.Is(err, secret) {
			t.Fatalf("error wrapping: %v", err)
		}
	}
}

func BenchmarkSessionSend(b *testing.B) {
	s, peer := sessionPair(b, sessionTestConfig(), Handlers{}, true, false)
	readerDone := make(chan struct{})
	go func() {
		defer close(readerDone)
		for {
			if _, _, err := peer.ReadMessage(); err != nil {
				return
			}
		}
	}()
	msg := Message{Type: Binary, Data: make([]byte, 256)}
	b.ReportAllocs()
	b.SetBytes(256)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if err := s.Send(context.Background(), msg); err != nil {
			b.Fatal(err)
		}
	}
	b.StopTimer()
	_ = s.Close(1000, "")
	<-readerDone
}