package notification

import (
	"bytes"
	"sync"
	"testing"
	"time"
)

// MockConnection implements the Connection interface for testing.
type MockConnection struct {
	mu       sync.Mutex
	Messages [][]byte
	Closed   bool
}

func (m *MockConnection) WriteMessage(message []byte) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.Messages = append(m.Messages, message)
	return nil
}

func (m *MockConnection) Close() error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.Closed = true
	return nil
}

func (m *MockConnection) GetMessages() [][]byte {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.Messages
}

func TestHubBroadcasting(t *testing.T) {
	hub := NewHub()

	// Start Hub in a background goroutine
	go hub.Run()

	// Create two mock connections and clients
	conn1 := &MockConnection{}
	client1 := &Client{
		Hub:  hub,
		Conn: conn1,
		Send: make(chan []byte, 10),
	}

	conn2 := &MockConnection{}
	client2 := &Client{
		Hub:  hub,
		Conn: conn2,
		Send: make(chan []byte, 10),
	}

	// Start write pumps
	go client1.WritePump()
	go client2.WritePump()

	// Register clients
	hub.Register <- client1
	hub.Register <- client2

	// Give a tiny moment for registration to process in the goroutine
	time.Sleep(10 * time.Millisecond)

	// Broadcast first event
	event1 := []byte("Welcome to Level 4!")
	hub.Broadcast <- event1

	// Give a moment for write pumps to consume channels
	time.Sleep(10 * time.Millisecond)

	// Verify both clients got event 1
	if len(conn1.GetMessages()) != 1 || !bytes.Equal(conn1.GetMessages()[0], event1) {
		t.Errorf("Client 1 did not receive correct message, got: %q", conn1.GetMessages())
	}
	if len(conn2.GetMessages()) != 1 || !bytes.Equal(conn2.GetMessages()[0], event1) {
		t.Errorf("Client 2 did not receive correct message, got: %q", conn2.GetMessages())
	}

	// Unregister client 1
	hub.Unregister <- client1
	time.Sleep(10 * time.Millisecond)

	// Broadcast second event
	event2 := []byte("Only Client 2 should see this")
	hub.Broadcast <- event2
	time.Sleep(10 * time.Millisecond)

	// Verify client 1 did not get event 2
	if len(conn1.GetMessages()) != 1 {
		t.Errorf("Client 1 received messages after being unregistered: %q", conn1.GetMessages())
	}

	// Verify client 2 got event 2 (so total messages is 2)
	if len(conn2.GetMessages()) != 2 || !bytes.Equal(conn2.GetMessages()[1], event2) {
		t.Errorf("Client 2 did not receive second message, got: %q", conn2.GetMessages())
	}
}