edubox/agent/ui.go

package main

import (
	_ "embed"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/exec"

	"github.com/gorilla/websocket"
)

//go:embed ui/index.html
var uiHTML string

var upgrader = websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}

func startUI(dataDir, nodeID, serverAddr string) {
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/html")
		fmt.Fprint(w, uiHTML)
	})

	http.HandleFunc("/api/config", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		switch r.Method {
		case http.MethodGet:
			cfg, _, err := loadOrCreateConfig(dataDir)
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			// Do not expose the auth key in plain GET unless requested; for local UI it is fine.
			json.NewEncoder(w).Encode(cfg)
		case http.MethodPost:
			var cfg AgentConfig
			if err := json.NewDecoder(r.Body).Decode(&cfg); err != nil {
				http.Error(w, err.Error(), http.StatusBadRequest)
				return
			}
			if cfg.DataDir == "" {
				cfg.DataDir = dataDir
			}
			if err := saveConfig(dataDir, &cfg); err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			w.WriteHeader(http.StatusNoContent)
		default:
			http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		}
	})

	http.HandleFunc("/api/restart", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
			return
		}
		w.WriteHeader(http.StatusNoContent)
		go func() {
			cmd := exec.Command(os.Args[0], os.Args[1:]...)
			cmd.Stdout = os.Stdout
			cmd.Stderr = os.Stderr
			cmd.Stdin = os.Stdin
			if err := cmd.Start(); err != nil {
				log.Printf("Restart failed: %v", err)
				return
			}
			os.Exit(0)
		}()
	})

	http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			log.Printf("UI WS upgrade error: %v", err)
			return
		}
		defer conn.Close()
		log.Printf("UI client connected from %s", r.RemoteAddr)

		// Register notifier to forward activation results from main WS to this UI connection
		notifierID := registerUINotifier(func(msg map[string]interface{}) {
			log.Printf("UI notifier forwarding to browser: %+v", msg)
			if err := conn.WriteJSON(msg); err != nil {
				log.Printf("UI notify error: %v", err)
			} else {
				log.Printf("UI notifier sent successfully")
			}
		})
		defer unregisterUINotifier(notifierID)

		for {
			var msg map[string]interface{}
			if err := conn.ReadJSON(&msg); err != nil {
				log.Printf("UI client disconnected: %v", err)
				break
			}
			action, _ := msg["action"].(string)
			log.Printf("UI received action: %s", action)
			switch action {
			case "check":
				act, err := loadActivation(dataDir)
				if err == nil && act.Activated {
					conn.WriteJSON(map[string]interface{}{"action": "activated", "studentName": act.StudentName})
				} else {
					conn.WriteJSON(map[string]interface{}{"action": "not_activated"})
				}
			case "activate":
				code, _ := msg["code"].(string)
				log.Printf("UI handling activate with code: %s", code)
				if err := sendMessage(WSMessage{Action: "activate", NodeID: nodeID, Code: code}); err != nil {
					log.Printf("UI sendMessage failed: %v", err)
					conn.WriteJSON(map[string]interface{}{"action": "activation_failed", "error": err.Error()})
				} else {
					log.Printf("UI sendMessage succeeded, waiting for server response...")
				}
			case "instances":
				listInstances(dataDir, conn)
			}
		}
	})

	port := "7070"
	log.Printf("%s UI starting on http://localhost:%s", APP_NAME, port)
	if err := http.ListenAndServe("127.0.0.1:"+port, nil); err != nil {
		log.Fatalf("%s UI server error: %v", APP_NAME, err)
	}
}

func listInstances(dataDir string, conn *websocket.Conn) {
	instances, err := loadInstances(dataDir)
	if err != nil {
		log.Printf("loadInstances error: %v", err)
		conn.WriteJSON(map[string]interface{}{"action": "instances_list", "instances": []interface{}{}})
		return
	}

	var list []map[string]interface{}
	for _, inst := range instances {
		status := getInstanceStatus(dataDir, inst.ID)
		if status != inst.Status {
			inst.Status = status
			_ = upsertInstance(dataDir, inst)
		}
		list = append(list, map[string]interface{}{
			"id":           inst.ID,
			"templateName": inst.TemplateName,
			"port":         inst.Port,
			"status":       inst.Status,
			"url":          instanceURL(inst),
		})
	}

	conn.WriteJSON(map[string]interface{}{"action": "instances_list", "instances": list})
}