edubox/agent/main.go

package main

import (
	"flag"
	"io"
	"log"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"sync"
	"syscall"
	"time"
)

// version is injected at build time via -ldflags "-X main.version=X.Y.Z"
var version = "dev"

const APP_NAME = "studioE5"

var (
	dataDir   = flag.String("data-dir", "./studioE5-data", "Répertoire de données")
	uiEnabled = flag.Bool("ui", true, "Activer l'interface locale HTMX")
	noTray    = flag.Bool("no-tray", false, "Désactiver l'icône dans la barre système (mode console)")
)

func main() {
	flag.Parse()

	dd := *dataDir
	if !filepath.IsAbs(dd) {
		ex, err := os.Executable()
		if err == nil {
			dd = filepath.Join(filepath.Dir(ex), dd)
		} else {
			wd, _ := os.Getwd()
			dd = filepath.Join(wd, dd)
		}
	}
	*dataDir = dd

	if err := os.MkdirAll(*dataDir, 0755); err != nil {
		log.Fatalf("Cannot create data-dir: %v", err)
	}

	// Redirect agent logs to a file so the console can be hidden on Windows.
	agentLogPath := filepath.Join(*dataDir, "agent.log")
	if agentLogFile, err := os.OpenFile(agentLogPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644); err == nil {
		log.SetOutput(io.MultiWriter(agentLogFile, uiLogWriter{}))
	} else {
		log.Printf("Cannot open agent log file %s: %v", agentLogPath, err)
	}

	cfg, _, err := loadOrCreateConfig(*dataDir)
	if err != nil {
		log.Fatalf("Cannot load config: %v", err)
	}

	if err := saveConfig(*dataDir, cfg); err != nil {
		log.Fatalf("Cannot save config: %v", err)
	}

	log.Printf("[%s Agent] version=%s node=%s data-dir=%s server=%s", APP_NAME, version, cfg.NodeID, *dataDir, cfg.Server)

	if *uiEnabled {
		go startUI(*dataDir, cfg.NodeID, cfg.Server)
	}

	go startWebSocket(cfg.Server, cfg.NodeID, *dataDir, cfg.HeadscaleURL, cfg.HeadscaleAuthKey)

	shutdownCh := make(chan struct{})

	// Capture Ctrl+C / SIGTERM so a console window close or service stop
	// triggers the same cleanup path as the tray "Quit" menu.
	go func() {
		sigCh := make(chan os.Signal, 1)
		signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
		<-sigCh
		log.Println("Shutdown signal received")
		close(shutdownCh)
	}()

	var cleanupWg sync.WaitGroup
	cleanupWg.Add(1)
	go func() {
		defer cleanupWg.Done()
		<-shutdownCh
		log.Println("Cleaning up before exit...")

		// Stop Tailscale so the next agent start does not conflict on the
		// same socket/state.
		stopTailscale()

		// Stop any running instances so containers are not left behind, but keep
		// their volumes intact so data survives the next agent start.
		if inst, err := loadInstances(*dataDir); err == nil {
			for id, info := range inst {
				if info.Status == "running" {
					log.Printf("Stopping instance %s", id)
					_ = dockerComposeStop(*dataDir, id)
					info.Status = "stopped"
					_ = sendMessage(WSMessage{Action: "instance_stopped", InstanceID: id})
				}
			}
			_ = saveInstances(*dataDir, inst)
		}
		log.Println("Cleanup complete")
	}()

	if *noTray {
		log.Printf("[%s Agent] running in console mode (no tray)", APP_NAME)
		<-shutdownCh
		cleanupWg.Wait()
		return
	}

	// Run tray on its own locked OS thread; keep main blocked so the process
	// does not exit when systray is not available (e.g. headless Linux).
	go func() {
		runtime.LockOSThread()
		defer runtime.UnlockOSThread()
		runTray(APP_NAME, shutdownCh)
	}()

	<-shutdownCh
	cleanupWg.Wait()
}

func startTailscaleAndReport(dataDir, nodeID, headscaleURL, authKey string) {
	ip, err := startTailscale(dataDir, nodeID, headscaleURL, authKey)
	if err != nil {
		log.Printf("Tailscale error: %v", err)
		return
	}
	log.Printf("Tailscale IP obtained: %s", ip)

	for {
		if err := sendMessage(WSMessage{Action: "tailscale_ip", NodeID: nodeID, TailscaleIP: ip}); err != nil {
			log.Printf("Waiting for WebSocket to send tailscale_ip...")
			time.Sleep(1 * time.Second)
			continue
		}
		log.Printf("Sent tailscale_ip to server: %s", ip)
		break
	}

	// Reconfigure tailscale serve for any instances that were left running
	// (e.g. after an agent restart while containers kept running).
	if inst, err := loadInstances(dataDir); err == nil {
		for id, info := range inst {
			if info.Status == "running" {
				log.Printf("Reconfiguring tailscale serve for running instance %s on port %d", id, info.Port)
				if err := setupTailscaleServe(info.Port); err != nil {
					log.Printf("setupTailscaleServe error for %s: %v", id, err)
				}
			}
		}
	}

	// Notify the local UI that the service status has changed.
	broadcastUI(map[string]interface{}{
		"action": "status",
		"status": buildUIStatus(dataDir),
	})
}