package mock

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	log "github.com/sirupsen/logrus"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true // Allow connections from any origin
	},
}

// MockWebSocketServer manages WebSocket connections and mock broadcasts
type MockWebSocketServer struct {
	discovery    *MockNodeDiscovery
	clients      map[*websocket.Conn]bool
	mutex        sync.RWMutex
	writeMutex   sync.Mutex
	shutdownChan chan struct{}
	shutdownOnce sync.Once
	logger       *log.Logger
}

// NewMockWebSocketServer creates a new mock WebSocket server
func NewMockWebSocketServer(discovery *MockNodeDiscovery) *MockWebSocketServer {
	mws := &MockWebSocketServer{
		discovery:    discovery,
		clients:      make(map[*websocket.Conn]bool),
		shutdownChan: make(chan struct{}),
		logger:       log.New(),
	}

	// Register callback for node updates
	discovery.AddCallback(mws.handleNodeUpdate)

	// Start periodic broadcasts
	go mws.startPeriodicBroadcasts()

	return mws
}

// HandleWebSocket handles WebSocket upgrade and connection
func (mws *MockWebSocketServer) HandleWebSocket(w http.ResponseWriter, r *http.Request) error {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		mws.logger.WithError(err).Error("Failed to upgrade WebSocket connection")
		return err
	}

	mws.mutex.Lock()
	mws.clients[conn] = true
	mws.mutex.Unlock()

	mws.logger.Debug("Mock WebSocket client connected")

	// Send current cluster state to newly connected client
	go mws.sendCurrentClusterState(conn)

	// Handle client messages and disconnection
	go mws.handleClient(conn)

	return nil
}

// handleClient handles messages from a WebSocket client
func (mws *MockWebSocketServer) handleClient(conn *websocket.Conn) {
	defer func() {
		mws.mutex.Lock()
		delete(mws.clients, conn)
		mws.mutex.Unlock()
		conn.Close()
		mws.logger.Debug("Mock WebSocket client disconnected")
	}()

	// Set read deadline and pong handler
	conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	conn.SetPongHandler(func(string) error {
		conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		return nil
	})

	// Start ping routine
	go func() {
		ticker := time.NewTicker(54 * time.Second)
		defer ticker.Stop()

		for {
			select {
			case <-mws.shutdownChan:
				return
			case <-ticker.C:
				conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
				if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
					return
				}
			}
		}
	}()

	// Read messages
	for {
		_, _, err := conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				mws.logger.WithError(err).Error("WebSocket error")
			}
			break
		}
	}
}

// sendCurrentClusterState sends the current cluster state to a newly connected client
func (mws *MockWebSocketServer) sendCurrentClusterState(conn *websocket.Conn) {
	nodes := mws.discovery.GetNodes()
	if len(nodes) == 0 {
		return
	}

	// Convert to mock format
	mockNodes := make(map[string]*NodeInfo)
	for ip, node := range nodes {
		mockNodes[ip] = &NodeInfo{
			IP:       node.IP,
			Hostname: node.Hostname,
			Status:   string(node.Status),
			Latency:  node.Latency,
			LastSeen: node.LastSeen,
			Labels:   node.Labels,
		}
	}

	members := GenerateMockClusterMembers(mockNodes)

	message := struct {
		Type        string      `json:"type"`
		Members     interface{} `json:"members"`
		PrimaryNode string      `json:"primaryNode"`
		TotalNodes  int         `json:"totalNodes"`
		Timestamp   string      `json:"timestamp"`
	}{
		Type:        "cluster_update",
		Members:     members,
		PrimaryNode: mws.discovery.GetPrimaryNode(),
		TotalNodes:  len(nodes),
		Timestamp:   time.Now().Format(time.RFC3339),
	}

	data, err := json.Marshal(message)
	if err != nil {
		mws.logger.WithError(err).Error("Failed to marshal cluster data")
		return
	}

	conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
	if err := conn.WriteMessage(websocket.TextMessage, data); err != nil {
		mws.logger.WithError(err).Error("Failed to send initial cluster state")
	}
}

// handleNodeUpdate is called when node information changes
func (mws *MockWebSocketServer) handleNodeUpdate(nodeIP, action string) {
	mws.logger.WithFields(log.Fields{
		"node_ip": nodeIP,
		"action":  action,
	}).Debug("Mock node update received, broadcasting to WebSocket clients")

	// Broadcast cluster update
	mws.BroadcastClusterUpdate()

	// Also broadcast node discovery event
	mws.broadcastNodeDiscovery(nodeIP, action)
}

// startPeriodicBroadcasts sends periodic updates to keep clients informed
func (mws *MockWebSocketServer) startPeriodicBroadcasts() {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-mws.shutdownChan:
			return
		case <-ticker.C:
			mws.BroadcastClusterUpdate()
		}
	}
}

// BroadcastClusterUpdate sends cluster updates to all connected clients
func (mws *MockWebSocketServer) BroadcastClusterUpdate() {
	mws.mutex.RLock()
	clients := make([]*websocket.Conn, 0, len(mws.clients))
	for client := range mws.clients {
		clients = append(clients, client)
	}
	mws.mutex.RUnlock()

	if len(clients) == 0 {
		return
	}

	nodes := mws.discovery.GetNodes()
	
	// Convert to mock format
	mockNodes := make(map[string]*NodeInfo)
	for ip, node := range nodes {
		mockNodes[ip] = &NodeInfo{
			IP:       node.IP,
			Hostname: node.Hostname,
			Status:   string(node.Status),
			Latency:  node.Latency,
			LastSeen: node.LastSeen,
			Labels:   node.Labels,
		}
	}

	members := GenerateMockClusterMembers(mockNodes)

	message := struct {
		Type        string      `json:"type"`
		Members     interface{} `json:"members"`
		PrimaryNode string      `json:"primaryNode"`
		TotalNodes  int         `json:"totalNodes"`
		Timestamp   string      `json:"timestamp"`
	}{
		Type:        "cluster_update",
		Members:     members,
		PrimaryNode: mws.discovery.GetPrimaryNode(),
		TotalNodes:  len(nodes),
		Timestamp:   time.Now().Format(time.RFC3339),
	}

	data, err := json.Marshal(message)
	if err != nil {
		mws.logger.WithError(err).Error("Failed to marshal cluster update")
		return
	}

	mws.logger.WithField("clients", len(clients)).Debug("Broadcasting mock cluster update")

	// Send to all clients with write synchronization
	mws.writeMutex.Lock()
	defer mws.writeMutex.Unlock()

	for _, client := range clients {
		client.SetWriteDeadline(time.Now().Add(5 * time.Second))
		if err := client.WriteMessage(websocket.TextMessage, data); err != nil {
			mws.logger.WithError(err).Error("Failed to send cluster update to client")
		}
	}
}

// broadcastNodeDiscovery sends node discovery events to all clients
func (mws *MockWebSocketServer) broadcastNodeDiscovery(nodeIP, action string) {
	mws.mutex.RLock()
	clients := make([]*websocket.Conn, 0, len(mws.clients))
	for client := range mws.clients {
		clients = append(clients, client)
	}
	mws.mutex.RUnlock()

	if len(clients) == 0 {
		return
	}

	message := struct {
		Type      string `json:"type"`
		Action    string `json:"action"`
		NodeIP    string `json:"nodeIp"`
		Timestamp string `json:"timestamp"`
	}{
		Type:      "node_discovery",
		Action:    action,
		NodeIP:    nodeIP,
		Timestamp: time.Now().Format(time.RFC3339),
	}

	data, err := json.Marshal(message)
	if err != nil {
		mws.logger.WithError(err).Error("Failed to marshal node discovery event")
		return
	}

	// Send to all clients with write synchronization
	mws.writeMutex.Lock()
	defer mws.writeMutex.Unlock()

	for _, client := range clients {
		client.SetWriteDeadline(time.Now().Add(5 * time.Second))
		if err := client.WriteMessage(websocket.TextMessage, data); err != nil {
			mws.logger.WithError(err).Error("Failed to send node discovery event to client")
		}
	}
}

// BroadcastFirmwareUploadStatus sends firmware upload status updates to all clients
func (mws *MockWebSocketServer) BroadcastFirmwareUploadStatus(nodeIP, status, filename string, fileSize int) {
	mws.mutex.RLock()
	clients := make([]*websocket.Conn, 0, len(mws.clients))
	for client := range mws.clients {
		clients = append(clients, client)
	}
	mws.mutex.RUnlock()

	if len(clients) == 0 {
		return
	}

	message := struct {
		Type      string `json:"type"`
		NodeIP    string `json:"nodeIp"`
		Status    string `json:"status"`
		Filename  string `json:"filename"`
		FileSize  int    `json:"fileSize"`
		Timestamp string `json:"timestamp"`
	}{
		Type:      "firmware_upload_status",
		NodeIP:    nodeIP,
		Status:    status,
		Filename:  filename,
		FileSize:  fileSize,
		Timestamp: time.Now().Format(time.RFC3339),
	}

	data, err := json.Marshal(message)
	if err != nil {
		mws.logger.WithError(err).Error("Failed to marshal firmware upload status")
		return
	}

	mws.logger.WithFields(log.Fields{
		"node_ip":  nodeIP,
		"status":   status,
		"clients":  len(clients),
	}).Debug("Broadcasting mock firmware upload status")

	// Send to all clients with write synchronization
	mws.writeMutex.Lock()
	defer mws.writeMutex.Unlock()

	for _, client := range clients {
		client.SetWriteDeadline(time.Now().Add(5 * time.Second))
		if err := client.WriteMessage(websocket.TextMessage, data); err != nil {
			mws.logger.WithError(err).Error("Failed to send firmware upload status to client")
		}
	}
}

// BroadcastRolloutProgress sends rollout progress updates to all clients
func (mws *MockWebSocketServer) BroadcastRolloutProgress(rolloutID, nodeIP, status string, current, total int) {
	mws.mutex.RLock()
	clients := make([]*websocket.Conn, 0, len(mws.clients))
	for client := range mws.clients {
		clients = append(clients, client)
	}
	mws.mutex.RUnlock()

	if len(clients) == 0 {
		return
	}

	message := struct {
		Type      string `json:"type"`
		RolloutID string `json:"rolloutId"`
		NodeIP    string `json:"nodeIp"`
		Status    string `json:"status"`
		Current   int    `json:"current"`
		Total     int    `json:"total"`
		Progress  int    `json:"progress"`
		Timestamp string `json:"timestamp"`
	}{
		Type:      "rollout_progress",
		RolloutID: rolloutID,
		NodeIP:    nodeIP,
		Status:    status,
		Current:   current,
		Total:     total,
		Progress:  calculateProgress(current, total, status),
		Timestamp: time.Now().Format(time.RFC3339),
	}

	data, err := json.Marshal(message)
	if err != nil {
		mws.logger.WithError(err).Error("Failed to marshal rollout progress")
		return
	}

	mws.logger.WithFields(log.Fields{
		"rollout_id": rolloutID,
		"node_ip":    nodeIP,
		"status":     status,
		"progress":   fmt.Sprintf("%d/%d", current, total),
	}).Debug("Broadcasting mock rollout progress")

	// Send to all clients with write synchronization
	mws.writeMutex.Lock()
	defer mws.writeMutex.Unlock()

	for _, client := range clients {
		client.SetWriteDeadline(time.Now().Add(5 * time.Second))
		if err := client.WriteMessage(websocket.TextMessage, data); err != nil {
			mws.logger.WithError(err).Error("Failed to send rollout progress to client")
		}
	}
}

// calculateProgress calculates the correct progress percentage based on current status
func calculateProgress(current, total int, status string) int {
	if total == 0 {
		return 0
	}

	// Base progress is based on completed nodes
	completedNodes := current - 1
	if status == "completed" {
		completedNodes = current
	}

	// Calculate base progress (completed nodes / total nodes)
	baseProgress := float64(completedNodes) / float64(total) * 100

	// If currently updating labels or uploading, add partial progress for the current node
	if status == "updating_labels" {
		nodeProgress := 100.0 / float64(total) * 0.25
		baseProgress += nodeProgress
	} else if status == "uploading" {
		nodeProgress := 100.0 / float64(total) * 0.5
		baseProgress += nodeProgress
	}

	// Ensure we don't exceed 100%
	if baseProgress > 100 {
		baseProgress = 100
	}

	return int(baseProgress)
}

// GetClientCount returns the number of connected WebSocket clients
func (mws *MockWebSocketServer) GetClientCount() int {
	mws.mutex.RLock()
	defer mws.mutex.RUnlock()
	return len(mws.clients)
}

// Shutdown gracefully shuts down the WebSocket server
func (mws *MockWebSocketServer) Shutdown(ctx context.Context) error {
	mws.shutdownOnce.Do(func() {
		mws.logger.Info("Shutting down mock WebSocket server")
		close(mws.shutdownChan)

		mws.mutex.Lock()
		clients := make([]*websocket.Conn, 0, len(mws.clients))
		for client := range mws.clients {
			clients = append(clients, client)
		}
		mws.mutex.Unlock()

		// Close all client connections
		for _, client := range clients {
			client.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, "Server shutting down"))
			client.Close()
		}
	})

	return nil
}