package hub

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

	"git.yuntuo.hk/k2safe/go-omni/ws"
	"github.com/gorilla/websocket"
)

type testConnection struct {
	id   string
	done chan struct{}
	once sync.Once
	full bool
	sent atomic.Int64
}

func (c *testConnection) ID() string            { return c.id }
func (c *testConnection) Done() <-chan struct{} { return c.done }
func (c *testConnection) TrySend(ws.Message) error {
	select {
	case <-c.done:
		return ws.ErrClosed
	default:
	}
	if c.full {
		return ws.ErrQueueFull
	}
	c.sent.Add(1)
	return nil
}
func (c *testConnection) Close(int, string) error { c.once.Do(func() { close(c.done) }); return nil }
func newConnection(id string) *testConnection {
	return &testConnection{id: id, done: make(chan struct{})}
}
func testHub(t *testing.T, cfg Config) *Hub {
	t.Helper()
	h, err := New(cfg)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = h.Close() })
	return h
}

func TestHubRoutingBackpressureAndCleanup(t *testing.T) {
	h := testHub(t, Config{})
	a, b, c := newConnection("a"), newConnection("b"), newConnection("c")
	b.full = true
	for _, connection := range []*testConnection{a, b, c} {
		if err := h.Add(connection); err != nil {
			t.Fatal(err)
		}
	}
	for _, id := range []string{"a", "b"} {
		if err := h.Join(id, "group"); err != nil {
			t.Fatal(err)
		}
	}
	message := ws.Message{Type: ws.Text, Data: []byte("hello")}
	result, err := h.SendGroup("group", message)
	if err != nil || result.Attempted != 2 || result.Queued != 1 || !errors.Is(result.Errors["b"], ws.ErrQueueFull) {
		t.Fatalf("group: %#v %v", result, err)
	}
	if c.sent.Load() != 0 {
		t.Fatal("group leaked")
	}
	result, err = h.Broadcast(message)
	if err != nil || result.Attempted != 3 || result.Queued != 2 {
		t.Fatalf("broadcast: %#v %v", result, err)
	}
	if err = h.Kick("a", 1000, ""); err != nil {
		t.Fatal(err)
	}
	deadline := time.After(time.Second)
	for h.Stats().Connections != 2 {
		select {
		case <-deadline:
			t.Fatal("auto cleanup timeout")
		default:
			time.Sleep(time.Millisecond)
		}
	}
	if got := h.Stats(); got.Memberships != 1 || got.Groups != 1 {
		t.Fatalf("stats: %#v", got)
	}
	if !h.Remove("b") {
		t.Fatal("remove")
	}
	if got := h.Stats(); got.Memberships != 0 || got.Groups != 0 {
		t.Fatalf("group leak: %#v", got)
	}
	if _, err = h.Send("missing", message); !errors.Is(err, ErrNotFound) {
		t.Fatal(err)
	}
	_ = h.Close()
	if _, err = h.Broadcast(message); !errors.Is(err, ErrClosed) {
		t.Fatal(err)
	}
	select {
	case <-c.done:
		t.Fatal("hub closed borrowed connection")
	default:
	}
}

func TestHubLimitsAndReusedID(t *testing.T) {
	h := testHub(t, Config{MaxConnections: 2, MaxGroups: 1, MaxMemberships: 1, MaxGroupsPerConnection: 1})
	first := newConnection("a")
	if err := h.Add(first); err != nil {
		t.Fatal(err)
	}
	if err := h.Add(first); !errors.Is(err, ErrDuplicate) {
		t.Fatal(err)
	}
	if err := h.Add(newConnection("b")); err != nil {
		t.Fatal(err)
	}
	if err := h.Add(newConnection("c")); !errors.Is(err, ErrLimit) {
		t.Fatal(err)
	}
	if err := h.Join("a", "one"); err != nil {
		t.Fatal(err)
	}
	if err := h.Join("a", "one"); err != nil {
		t.Fatal("idempotent join", err)
	}
	for _, pair := range [][2]string{{"a", "two"}, {"b", "one"}} {
		if err := h.Join(pair[0], pair[1]); !errors.Is(err, ErrLimit) {
			t.Fatal(err)
		}
	}
	h.Remove("a")
	second := newConnection("a")
	if err := h.Add(second); err != nil {
		t.Fatal(err)
	}
	_ = first.Close(1000, "")
	if err := h.Join("a", "replacement"); err != nil {
		t.Fatal(err)
	}
	if _, err := h.Send("a", ws.Message{Type: ws.Text}); err != nil {
		t.Fatal("old watcher removed replacement", err)
	}
	var nilConnection *testConnection
	if err := h.Add(nilConnection); !errors.Is(err, ErrInvalid) {
		t.Fatal(err)
	}
}

func TestHubConcurrentLifecycle(t *testing.T) {
	h := testHub(t, Config{MaxConnections: 100})
	var wg sync.WaitGroup
	for i := 0; i < 50; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			c := newConnection(fmt.Sprint(i))
			if err := h.Add(c); err != nil {
				t.Error(err)
				return
			}
			for j := 0; j < 10; j++ {
				_ = h.Join(c.ID(), "shared")
				_, _ = h.Broadcast(ws.Message{Type: ws.Text})
				h.Leave(c.ID(), "shared")
			}
			_ = c.Close(1000, "")
		}(i)
	}
	wg.Wait()
	_ = h.Close()
	if h.Stats() != (Stats{}) {
		t.Fatal("registry not released")
	}
}

func TestHubRealSessionDeliveryAndDisconnect(t *testing.T) {
	received := make(chan string, 1)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer conn.Close()
		for {
			_, data, err := conn.ReadMessage()
			if err != nil {
				return
			}
			received <- string(data)
		}
	}))
	defer server.Close()
	client, err := ws.NewClient(ws.ClientConfig{ProxyMode: "direct", Session: ws.SessionConfig{CloseTimeout: 100 * time.Millisecond}})
	if err != nil {
		t.Fatal(err)
	}
	defer func() {
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		if err := client.Shutdown(ctx); err != nil {
			t.Error(err)
		}
	}()
	s, err := client.Dial(context.Background(), "ws"+strings.TrimPrefix(server.URL, "http"), ws.Handlers{})
	if err != nil {
		t.Fatal(err)
	}
	h := testHub(t, Config{})
	if err = h.Add(s); err != nil {
		t.Fatal(err)
	}
	if err = h.Join(s.ID(), "room"); err != nil {
		t.Fatal(err)
	}
	result, err := h.SendGroup("room", ws.Message{Type: ws.Text, Data: []byte("queued")})
	if err != nil || result.Queued != 1 {
		t.Fatalf("send: %#v %v", result, err)
	}
	select {
	case got := <-received:
		if got != "queued" {
			t.Fatal(got)
		}
	case <-time.After(time.Second):
		t.Fatal("delivery timed out")
	}
	_ = s.Close(1000, "")
	deadline := time.After(time.Second)
	for h.Stats().Connections != 0 {
		select {
		case <-deadline:
			t.Fatal("disconnected Session still registered")
		default:
			time.Sleep(time.Millisecond)
		}
	}
}

func BenchmarkHubBroadcast(b *testing.B) {
	for _, count := range []int{1, 100, 1000} {
		b.Run(fmt.Sprint(count), func(b *testing.B) {
			h, _ := New(Config{})
			defer h.Close()
			for i := 0; i < count; i++ {
				if err := h.Add(newConnection(fmt.Sprint(i))); err != nil {
					b.Fatal(err)
				}
			}
			message := ws.Message{Type: ws.Text, Data: []byte("benchmark")}
			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				result, err := h.Broadcast(message)
				if err != nil || result.Queued != count {
					b.Fatal("broadcast failed")
				}
			}
		})
	}
}