package websocket

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

// MockWebSocketHandler is a mock implementation for testing
type MockWebSocketHandler struct {
	mu         sync.Mutex
	messages   []map[string]interface{}
	broadcasts []string
}

func (m *MockWebSocketHandler) BroadcastMessage(messageType string, data interface{}) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.broadcasts = append(m.broadcasts, messageType)
	m.messages = append(m.messages, map[string]interface{}{
		"type": messageType,
		"data": data,
	})
}

func (m *MockWebSocketHandler) GetMessages() []map[string]interface{} {
	m.mu.Lock()
	defer m.mu.Unlock()
	result := make([]map[string]interface{}, len(m.messages))
	copy(result, m.messages)
	return result
}

func (m *MockWebSocketHandler) GetBroadcasts() []string {
	m.mu.Lock()
	defer m.mu.Unlock()
	result := make([]string, len(m.broadcasts))
	copy(result, m.broadcasts)
	return result
}

func (m *MockWebSocketHandler) Clear() {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.messages = make([]map[string]interface{}, 0)
	m.broadcasts = make([]string, 0)
}

func NewMockWebSocketHandler() *MockWebSocketHandler {
	return &MockWebSocketHandler{
		messages:   make([]map[string]interface{}, 0),
		broadcasts: make([]string, 0),
	}
}

func TestBroadcaster_StartHeartbeat(t *testing.T) {
	mockHandler := NewMockWebSocketHandler()
	broadcaster := NewBroadcaster(mockHandler)

	// Start heartbeat with short interval for testing
	broadcaster.StartHeartbeat(100 * time.Millisecond)

	// Wait for a few heartbeats
	time.Sleep(350 * time.Millisecond)

	// Stop the broadcaster
	broadcaster.Stop()

	// Check if heartbeats were sent
	messages := mockHandler.GetMessages()
	if len(messages) == 0 {
		t.Error("Expected heartbeat messages, but got none")
	}

	// Check that all messages are heartbeat type
	broadcasts := mockHandler.GetBroadcasts()
	for _, msgType := range broadcasts {
		if msgType != "heartbeat" {
			t.Errorf("Expected heartbeat message type, got %s", msgType)
		}
	}

	t.Logf("Received %d heartbeat messages", len(messages))
}

func TestBroadcaster_BroadcastNotification(t *testing.T) {
	mockHandler := NewMockWebSocketHandler()
	broadcaster := NewBroadcaster(mockHandler)

	// Send a notification
	broadcaster.BroadcastNotification("Test Title", "Test Message", "info")

	// Check if notification was sent
	messages := mockHandler.GetMessages()
	if len(messages) != 1 {
		t.Errorf("Expected 1 message, got %d", len(messages))
		return
	}

	msg := messages[0]
	if msg["type"] != "notification" {
		t.Errorf("Expected message type 'notification', got %s", msg["type"])
	}

	data := msg["data"].(map[string]interface{})
	if data["title"] != "Test Title" {
		t.Errorf("Expected title 'Test Title', got %s", data["title"])
	}
	if data["message"] != "Test Message" {
		t.Errorf("Expected message 'Test Message', got %s", data["message"])
	}
	if data["level"] != "info" {
		t.Errorf("Expected level 'info', got %s", data["level"])
	}

	t.Logf("Notification sent successfully: %+v", data)
}

func TestBroadcaster_SimulateDataStream(t *testing.T) {
	mockHandler := NewMockWebSocketHandler()
	broadcaster := NewBroadcaster(mockHandler)

	// Start data stream with short interval for testing
	broadcaster.SimulateDataStream()

	// Wait for a few data points
	time.Sleep(550 * time.Millisecond)

	// Stop the broadcaster
	broadcaster.Stop()

	// Check if data stream messages were sent
	messages := mockHandler.GetMessages()
	if len(messages) == 0 {
		t.Error("Expected data stream messages, but got none")
	}

	// Check that all messages are data_stream type
	broadcasts := mockHandler.GetBroadcasts()
	for _, msgType := range broadcasts {
		if msgType != "data_stream" {
			t.Errorf("Expected data_stream message type, got %s", msgType)
		}
	}

	// Check data structure
	for i, msg := range messages {
		data := msg["data"].(map[string]interface{})
		if data["type"] != "simulated_data" {
			t.Errorf("Expected data type 'simulated_data', got %s", data["type"])
		}
		if id, ok := data["id"].(int); ok {
			if id != i+1 {
				t.Errorf("Expected id %d, got %d", i+1, id)
			}
		}
		if value, ok := data["value"].(int); ok {
			expectedValue := (i + 1) * 10
			if value != expectedValue {
				t.Errorf("Expected value %d, got %d", expectedValue, value)
			}
		}
	}

	t.Logf("Received %d data stream messages", len(messages))
}

func TestBroadcaster_Stop(t *testing.T) {
	mockHandler := NewMockWebSocketHandler()
	broadcaster := NewBroadcaster(mockHandler)

	// Start heartbeat
	broadcaster.StartHeartbeat(50 * time.Millisecond)

	// Wait a bit
	time.Sleep(100 * time.Millisecond)

	// Stop the broadcaster
	broadcaster.Stop()

	// Clear previous messages
	mockHandler.Clear()

	// Wait a bit more to ensure no new messages are sent
	time.Sleep(200 * time.Millisecond)

	// Check that no new messages were sent after stopping
	messages := mockHandler.GetMessages()
	if len(messages) > 0 {
		t.Errorf("Expected no messages after stopping, but got %d", len(messages))
	}

	// Clear quit channel to allow reuse in tests
	broadcaster.quit = make(chan struct{})

	t.Log("Broadcaster stopped successfully")
}

func TestBroadcaster_MultipleOperations(t *testing.T) {
	mockHandler := NewMockWebSocketHandler()
	broadcaster := NewBroadcaster(mockHandler)

	// Start heartbeat
	broadcaster.StartHeartbeat(100 * time.Millisecond)

	// Send notification
	broadcaster.BroadcastNotification("Test", "Message", "warning")

	// Start data stream
	broadcaster.SimulateDataStream()

	// Wait for some activity
	time.Sleep(350 * time.Millisecond)

	// Stop everything
	broadcaster.Stop()

	// Check results
	messages := mockHandler.GetMessages()
	if len(messages) == 0 {
		t.Error("Expected messages from multiple operations, but got none")
	}

	broadcasts := mockHandler.GetBroadcasts()
	hasHeartbeat := false
	hasNotification := false
	hasDataStream := false

	for _, msgType := range broadcasts {
		switch msgType {
		case "heartbeat":
			hasHeartbeat = true
		case "notification":
			hasNotification = true
		case "data_stream":
			hasDataStream = true
		}
	}

	if !hasHeartbeat {
		t.Error("Expected heartbeat messages")
	}
	if !hasNotification {
		t.Error("Expected notification message")
	}
	if !hasDataStream {
		t.Error("Expected data stream messages")
	}

	t.Logf("Multiple operations test passed: %d total messages", len(messages))
}