package ws

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

	"github.com/gorilla/websocket"
)

func managedServer(t testing.TB, cfg ServerConfig, h Handlers) (*Server, *httptest.Server, string) {
	t.Helper()
	s, err := NewServer(cfg)
	if err != nil {
		t.Fatal(err)
	}
	httpServer := httptest.NewServer(s.Handler(h))
	t.Cleanup(func() {
		httpServer.Close()
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		if err := s.Shutdown(ctx); err != nil {
			t.Error(err)
		}
	})
	return s, httpServer, "ws" + strings.TrimPrefix(httpServer.URL, "http")
}

func TestServerOriginAndProtocolPolicy(t *testing.T) {
	_, host, endpoint := managedServer(t, ServerConfig{Subprotocols: []string{"test.v1"}, RequireSubprotocol: true}, Handlers{})
	for _, tt := range []struct {
		name, origin string
		code         int
	}{
		{"service", "", 101}, {"same", host.URL, 101}, {"foreign", "https://foreign.invalid", 403}, {"null", "null", 403}, {"path", host.URL + "/", 403}, {"comma", host.URL + ", https://evil.invalid", 403},
	} {
		t.Run(tt.name, func(t *testing.T) {
			d := websocket.Dialer{Subprotocols: []string{"test.v1"}}
			h := http.Header{}
			if tt.origin != "" {
				h.Set("Origin", tt.origin)
			}
			c, r, err := d.Dial(endpoint, h)
			if c != nil {
				defer c.Close()
			}
			if r == nil || r.StatusCode != tt.code {
				t.Fatalf("status=%v err=%v", r, err)
			}
			if tt.code == 101 && (err != nil || c.Subprotocol() != "test.v1") {
				t.Fatal("negotiation failed", err)
			}
		})
	}
	_, r, err := websocket.DefaultDialer.Dial(endpoint, nil)
	if err == nil || r.StatusCode != 400 {
		t.Fatal("missing subprotocol accepted")
	}
	d := websocket.Dialer{Subprotocols: []string{"test.v1"}}
	_, r, err = d.Dial(endpoint, http.Header{"Origin": []string{host.URL, host.URL}})
	if err == nil || r.StatusCode != 403 {
		t.Fatal("duplicate origin accepted")
	}
}

func TestServerAllowedOriginsAndAuthentication(t *testing.T) {
	s, _, endpoint := managedServer(t, ServerConfig{AllowedOrigins: []string{"https://allowed.invalid"}, RequireOrigin: true, Prepare: func(ctx context.Context, r *http.Request) (Handlers, error) {
		if r.Header.Get("Authorization") != "Bearer local-test" {
			return Handlers{}, &RejectError{StatusCode: 403}
		}
		return Handlers{OnMessage: func(ctx context.Context, s *Session, m Message) error { return s.Send(ctx, m) }}, nil
	}}, Handlers{})
	for _, h := range []http.Header{nil, {"Origin": []string{"https://allowed.invalid"}}, {"Authorization": []string{"Bearer local-test"}}, {"Origin": []string{"https://wrong.invalid"}, "Authorization": []string{"Bearer local-test"}}} {
		c, r, e := websocket.DefaultDialer.Dial(endpoint, h)
		if c != nil {
			c.Close()
		}
		if e == nil || r.StatusCode != 403 {
			t.Fatal("unauthorized handshake accepted")
		}
	}
	c, _, err := websocket.DefaultDialer.Dial(endpoint, http.Header{"Origin": []string{"https://allowed.invalid"}, "Authorization": []string{"Bearer local-test"}})
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	if err = c.WriteMessage(websocket.TextMessage, []byte("echo")); err != nil {
		t.Fatal(err)
	}
	c.SetReadDeadline(time.Now().Add(time.Second))
	_, body, err := c.ReadMessage()
	if err != nil || string(body) != "echo" {
		t.Fatal(string(body), err)
	}
	if s.Stats().Accepted != 1 {
		t.Fatal(s.Stats())
	}
}

func TestServerStopDuringHandshake(t *testing.T) {
	entered, release := make(chan struct{}), make(chan struct{})
	s, _, endpoint := managedServer(t, ServerConfig{Prepare: func(context.Context, *http.Request) (Handlers, error) {
		close(entered)
		<-release
		return Handlers{}, nil
	}}, Handlers{})
	dialDone := make(chan struct{})
	go func() {
		defer close(dialDone)
		c, _, _ := websocket.DefaultDialer.Dial(endpoint, nil)
		if c != nil {
			defer c.Close()
			c.ReadMessage()
		}
	}()
	<-entered
	_ = s.Stop(context.Background())
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
	defer cancel()
	if err := s.Drain(ctx); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatal("handshake was not tracked", err)
	}
	close(release)
	select {
	case <-dialDone:
	case <-time.After(time.Second):
		t.Fatal("accepted connection escaped stop")
	}
	ctx2, cancel2 := context.WithTimeout(context.Background(), time.Second)
	defer cancel2()
	if err := s.Drain(ctx2); err != nil {
		t.Fatal(err)
	}
}

func TestServerLimitsAndShutdownCallbacks(t *testing.T) {
	entered, release := make(chan struct{}), make(chan struct{})
	s, _, endpoint := managedServer(t, ServerConfig{MaxConnections: 1, Session: SessionConfig{CloseTimeout: 20 * time.Millisecond}}, Handlers{OnMessage: func(context.Context, *Session, Message) error { close(entered); <-release; return nil }})
	c, _, err := websocket.DefaultDialer.Dial(endpoint, nil)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	_, r, err := websocket.DefaultDialer.Dial(endpoint, nil)
	if err == nil || r.StatusCode != 503 {
		t.Fatal("connection cap failed")
	}
	if err = c.WriteMessage(websocket.TextMessage, []byte("hold")); err != nil {
		t.Fatal(err)
	}
	<-entered
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
	defer cancel()
	if err = s.Shutdown(ctx); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatal("blocked callback incorrectly drained", err)
	}
	close(release)
	ctx2, cancel2 := context.WithTimeout(context.Background(), time.Second)
	defer cancel2()
	if err = s.Drain(ctx2); err != nil {
		t.Fatal(err)
	}
	_, r, err = websocket.DefaultDialer.Dial(endpoint, nil)
	if err == nil || r.StatusCode != 503 {
		t.Fatal("server admitted after shutdown")
	}
}

func TestServerHandshakeRateAndPanic(t *testing.T) {
	_, _, endpoint := managedServer(t, ServerConfig{HandshakeRate: 0.001, HandshakeBurst: 1, Prepare: func(context.Context, *http.Request) (Handlers, error) { panic("SECRET") }}, Handlers{})
	_, r, e := websocket.DefaultDialer.Dial(endpoint, nil)
	if e == nil || r.StatusCode != 401 {
		t.Fatal("prepare panic not contained", e)
	}
	_, r, e = websocket.DefaultDialer.Dial(endpoint, nil)
	if e == nil || r.StatusCode != 429 {
		t.Fatal("rate cap not enforced", e)
	}
}

func TestServerUnsupportedWriter(t *testing.T) {
	s, err := NewServer(ServerConfig{})
	if err != nil {
		t.Fatal(err)
	}
	r := httptest.NewRequest("GET", "http://localhost/ws", nil)
	r.Header.Set("Upgrade", "websocket")
	r.Header.Set("Connection", "Upgrade")
	w := httptest.NewRecorder()
	_, err = s.Accept(w, r, Handlers{})
	if !errors.Is(err, ErrUnsupportedTransport) || w.Code != 501 {
		t.Fatal(w.Code, err)
	}
	if err = s.Shutdown(context.Background()); err != nil {
		t.Fatal(err)
	}
}

func TestServerConcurrentShutdown(t *testing.T) {
	s, _, endpoint := managedServer(t, ServerConfig{Session: SessionConfig{CloseTimeout: 10 * time.Millisecond}}, Handlers{})
	var wg sync.WaitGroup
	for i := 0; i < 40; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			c, _, _ := websocket.DefaultDialer.Dial(endpoint, nil)
			if c != nil {
				defer c.Close()
				c.SetReadDeadline(time.Now().Add(time.Second))
				c.ReadMessage()
			}
		}()
	}
	_ = s.Stop(context.Background())
	wg.Wait()
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	if err := s.Drain(ctx); err != nil {
		t.Fatal(err)
	}
	if stats := s.Stats(); stats.ActiveConnections != 0 || stats.Handshakes != 0 {
		t.Fatal(stats)
	}
}

func TestServerConfigValidation(t *testing.T) {
	for _, cfg := range []ServerConfig{{AllowedOrigins: []string{"*"}}, {AllowedOrigins: []string{"https://user:secret@host"}}, {AllowedOrigins: []string{"https://host/path"}}, {HandshakeRate: -1}, {HandshakeBurst: 1}, {RequireSubprotocol: true}, {Subprotocols: []string{"bad\r\n"}}, {MaxConnections: -1}} {
		if err := cfg.Validate(); err == nil {
			t.Fatal("bad config accepted")
		}
	}
}

func TestServerParallelConnections(t *testing.T) {
	s, _, endpoint := managedServer(t, ServerConfig{MaxConnections: 256}, Handlers{OnMessage: func(ctx context.Context, s *Session, m Message) error { return s.Send(ctx, m) }})
	var workers sync.WaitGroup
	for i := 0; i < 128; i++ {
		workers.Add(1)
		go func(id int) {
			defer workers.Done()
			c, _, err := websocket.DefaultDialer.Dial(endpoint, nil)
			if err != nil {
				t.Error(err)
				return
			}
			defer c.Close()
			c.SetReadDeadline(time.Now().Add(15 * time.Second))
			c.SetWriteDeadline(time.Now().Add(15 * time.Second))
			for n := 0; n < 50; n++ {
				value := fmt.Sprintf("%d:%d", id, n)
				if err := c.WriteMessage(websocket.TextMessage, []byte(value)); err != nil {
					t.Error(err)
					return
				}
				_, got, err := c.ReadMessage()
				if err != nil || string(got) != value {
					t.Errorf("session isolation: %s %v", got, err)
					return
				}
			}
		}(i)
	}
	workers.Wait()
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	if err := s.Shutdown(ctx); err != nil {
		t.Fatal(err)
	}
	if s.Stats().Accepted != 128 {
		t.Fatal(s.Stats())
	}
}

func TestServerTypedNilRejection(t *testing.T) {
	_, _, endpoint := managedServer(t, ServerConfig{Prepare: func(context.Context, *http.Request) (Handlers, error) { var err *RejectError; return Handlers{}, err }}, Handlers{})
	_, r, err := websocket.DefaultDialer.Dial(endpoint, nil)
	if err == nil || r == nil || r.StatusCode != 401 {
		t.Fatal(r, err)
	}
}

func FuzzOrigin(f *testing.F) {
	for _, v := range []string{"https://example.com", "http://localhost:8080", "null", "https://[::1]:443", "https://host:invalid", "https://evil@safe"} {
		f.Add(v)
	}
	f.Fuzz(func(t *testing.T, v string) {
		a, ok := canonicalOrigin(v)
		if ok {
			b, ok2 := canonicalOrigin(a)
			if !ok2 || a != b {
				t.Fatal("not canonical")
			}
		}
	})
}

func BenchmarkServerRoundTrip(b *testing.B) {
	for _, size := range []int{64, 1024, 64 << 10} {
		b.Run(fmt.Sprint(size), func(b *testing.B) {
			_, _, endpoint := managedServer(b, ServerConfig{}, Handlers{OnMessage: func(ctx context.Context, s *Session, m Message) error { return s.Send(ctx, m) }})
			c, _, err := websocket.DefaultDialer.Dial(endpoint, nil)
			if err != nil {
				b.Fatal(err)
			}
			defer c.Close()
			payload := make([]byte, size)
			b.SetBytes(int64(size) * 2)
			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				if err = c.WriteMessage(websocket.BinaryMessage, payload); err != nil {
					b.Fatal(err)
				}
				if _, _, err = c.ReadMessage(); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}