package main

import (
	"log"
	"time"

	"github.com/gorilla/websocket"
)

type WSMessage struct {
	Action        string `json:"action"`
	NodeID        string `json:"nodeId,omitempty"`
	Code          string `json:"code,omitempty"`
	InstanceID    string `json:"instanceId,omitempty"`
	Type          string `json:"type,omitempty"`
	Port          int    `json:"port,omitempty"`
	ComposeConfig string `json:"composeConfig,omitempty"`
	StudentName   string `json:"studentName,omitempty"`
	Error         string `json:"error,omitempty"`
}

func startWebSocket(serverAddr, nodeID, dataDir string) {
	for {
		conn, _, err := websocket.DefaultDialer.Dial(serverAddr, nil)
		if err != nil {
			log.Printf("WS connect error: %v, retrying in 5s...", err)
			time.Sleep(5 * time.Second)
			continue
		}

		log.Printf("WS connected to %s", serverAddr)

		// Register
		if err := conn.WriteJSON(WSMessage{Action: "register", NodeID: nodeID}); err != nil {
			log.Printf("WS register error: %v", err)
			conn.Close()
			continue
		}

		// Activation flow
		act, err := loadActivation(dataDir)
		if err != nil || !act.Activated {
			log.Println("Waiting for activation...")
		} else {
			log.Printf("Already activated as %s", act.StudentName)
		}

		// Heartbeat goroutine
		done := make(chan struct{})
		go func() {
			ticker := time.NewTicker(10 * time.Second)
			defer ticker.Stop()
			for {
				select {
				case <-ticker.C:
					if err := conn.WriteJSON(WSMessage{Action: "heartbeat", NodeID: nodeID}); err != nil {
						return
					}
				case <-done:
					return
				}
			}
		}()

		// Read loop
		for {
			var msg WSMessage
			if err := conn.ReadJSON(&msg); err != nil {
				log.Printf("WS read error: %v", err)
				break
			}
			handleMessage(conn, msg, dataDir, nodeID)
		}

		close(done)
		conn.Close()
		log.Println("WS disconnected, reconnecting in 5s...")
		time.Sleep(5 * time.Second)
	}
}

func handleMessage(conn *websocket.Conn, msg WSMessage, dataDir, nodeID string) {
	switch msg.Action {
	case "activate":
		// handled by UI, but server can also push activation response
		if msg.StudentName != "" {
			act := &Activation{Activated: true, StudentName: msg.StudentName, Code: msg.Code}
			saveActivation(dataDir, act)
			log.Printf("Activated as %s", msg.StudentName)
		}
	case "start":
		log.Printf("Start instance %s on port %d", msg.InstanceID, msg.Port)
		if err := writeCompose(dataDir, msg.InstanceID, msg.ComposeConfig); err != nil {
			log.Printf("writeCompose error: %v", err)
			conn.WriteJSON(WSMessage{Action: "instance_error", InstanceID: msg.InstanceID, Error: err.Error()})
			return
		}
		if err := dockerComposeUp(dataDir, msg.InstanceID); err != nil {
			log.Printf("dockerComposeUp error: %v", err)
			conn.WriteJSON(WSMessage{Action: "instance_error", InstanceID: msg.InstanceID, Error: err.Error()})
			return
		}
		conn.WriteJSON(WSMessage{Action: "instance_started", InstanceID: msg.InstanceID, Port: msg.Port})
	case "stop":
		log.Printf("Stop instance %s", msg.InstanceID)
		if err := dockerComposeDown(dataDir, msg.InstanceID); err != nil {
			log.Printf("dockerComposeDown error: %v", err)
		}
	case "reset":
		log.Printf("Reset instance %s", msg.InstanceID)
		dockerComposeRm(dataDir, msg.InstanceID)
		if err := writeCompose(dataDir, msg.InstanceID, msg.ComposeConfig); err != nil {
			log.Printf("writeCompose error: %v", err)
			return
		}
		if err := dockerComposeUp(dataDir, msg.InstanceID); err != nil {
			log.Printf("dockerComposeUp error: %v", err)
			conn.WriteJSON(WSMessage{Action: "instance_error", InstanceID: msg.InstanceID, Error: err.Error()})
			return
		}
		conn.WriteJSON(WSMessage{Action: "instance_started", InstanceID: msg.InstanceID, Port: msg.Port})
	default:
		log.Printf("Unknown action: %s", msg.Action)
	}
}