visionA/local-tool/server/main.go

package main

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"
	"time"

	"visiona-local/server/internal/api"
	"visiona-local/server/internal/api/handlers"
	"visiona-local/server/internal/api/ws"
	"visiona-local/server/internal/camera"
	"visiona-local/server/internal/config"
	"visiona-local/server/internal/deps"
	"visiona-local/server/internal/device"
	"visiona-local/server/internal/flash"
	"visiona-local/server/internal/inference"
	"visiona-local/server/internal/model"
	pkglogger "visiona-local/server/pkg/logger"
	"visiona-local/server/web"
)

var (
	Version   = "dev"
	BuildTime = "unknown"
)

// baseDir returns the base directory for resolving data/ and scripts/ paths.
// In dev mode (go run), uses the working directory.
// In production (compiled binary), uses the binary's directory so the server
// works correctly regardless of the working directory.
func baseDir(devMode bool) string {
	if devMode {
		return "."
	}
	exe, err := os.Executable()
	if err != nil {
		return "."
	}
	return filepath.Dir(exe)
}

// resolveBridgeScript finds kneron_bridge.py across different packaging layouts.
//
// Possible locations (tried in order):
//  1. <base>/scripts/kneron_bridge.py              — dev mode or flat layout
//  2. <base>/../scripts/kneron_bridge.py           — Windows/Linux installer: binary in {app}/bin, scripts in {app}/scripts
//  3. <base>/../Resources/scripts/kneron_bridge.py — macOS app bundle: binary in Contents/MacOS, scripts in Contents/Resources
//  4. ./scripts/kneron_bridge.py                   — cwd fallback
func resolveBridgeScript(base string) string {
	candidates := []string{
		filepath.Join(base, "scripts", "kneron_bridge.py"),
		filepath.Join(base, "..", "scripts", "kneron_bridge.py"),
		filepath.Join(base, "..", "Resources", "scripts", "kneron_bridge.py"),
		filepath.Join(".", "scripts", "kneron_bridge.py"),
	}
	for _, c := range candidates {
		abs, err := filepath.Abs(c)
		if err != nil {
			continue
		}
		if info, err := os.Stat(abs); err == nil && !info.IsDir() {
			return abs
		}
	}
	// Nothing found — return the default so downstream logs a clear error
	return filepath.Join(base, "scripts", "kneron_bridge.py")
}

func main() {
	cfg := config.Load()

	logger := pkglogger.New(cfg.LogLevel)

	logger.Info("Starting visionA-local Server %s (built: %s)", Version, BuildTime)
	logger.Info("Mock mode: %v, Mock camera: %v, Dev mode: %v, Python mode: %s",
		cfg.MockMode, cfg.MockCamera, cfg.DevMode, cfg.PythonMode)

	// 把 VISIONA_BUNDLE_BIN_DIR 加到 PATH，讓 exec.Command("yt-dlp") / exec.Command("ffmpeg")
	// 能透過 LookPath 找到 bundle 內的 binary（Go 1.19+ Windows 不再搜 cwd）。
	if bundleBin := os.Getenv("VISIONA_BUNDLE_BIN_DIR"); bundleBin != "" {
		sep := string(os.PathListSeparator)
		os.Setenv("PATH", bundleBin+sep+os.Getenv("PATH"))
		logger.Info("Added VISIONA_BUNDLE_BIN_DIR to PATH: %s", bundleBin)
	}

	// Check external dependencies
	deps.PrintStartupReport(logger)

	// Resolve base directory and data directory
	base := baseDir(cfg.DevMode)
	dataDir := cfg.DataDir
	if dataDir == "" {
		dataDir = filepath.Join(base, "data")
	}

	// Initialize model repository (built-in models from JSON)
	modelRepo := model.NewRepository(filepath.Join(dataDir, "models.json"))
	logger.Info("Loaded %d built-in models", modelRepo.Count())

	// Initialize model store (custom uploaded models)
	customModelDir := cfg.ModelDir
	if customModelDir == "" {
		customModelDir = filepath.Join(dataDir, "custom-models")
	}
	modelStore := model.NewModelStore(customModelDir)
	customModels, err := modelStore.LoadCustomModels()
	if err != nil {
		logger.Warn("Failed to load custom models: %v", err)
	}
	for _, m := range customModels {
		modelRepo.Add(m)
	}
	if len(customModels) > 0 {
		logger.Info("Loaded %d custom models", len(customModels))
	}

	// Initialize WebSocket hub (before device manager so log broadcaster is ready)
	wsHub := ws.NewHub()
	go wsHub.Run()

	// Initialize log broadcaster for real-time log streaming
	logBroadcaster := pkglogger.NewBroadcaster(500, func(entry pkglogger.LogEntry) {
		wsHub.BroadcastToRoom("server-logs", entry)
	})
	logger.SetBroadcaster(logBroadcaster)

	// Initialize device manager
	registry := device.NewRegistry()
	bridgeScript := resolveBridgeScript(base)
	logger.Info("Kneron bridge script: %s", bridgeScript)
	deviceMgr := device.NewManager(registry, cfg.MockMode, cfg.MockDeviceCount, bridgeScript)
	deviceMgr.SetLogBroadcaster(logBroadcaster)
	deviceMgr.Start()

	// Initialize camera manager
	cameraMgr := camera.NewManager(cfg.MockCamera)

	// Initialize services
	flashSvc := flash.NewService(deviceMgr, modelRepo, dataDir)
	inferenceSvc := inference.NewService(deviceMgr)

	// Determine static file system for embedded frontend
	var staticFS http.FileSystem
	if !cfg.DevMode {
		staticFS = web.StaticFS()
		logger.Info("Serving embedded frontend static files")
	} else {
		logger.Info("Dev mode: frontend static serving disabled (use Wails dev server)")
	}

	// Build HTTP server (needed for graceful shutdown and restart)
	var httpServer *http.Server
	restartRequested := make(chan struct{}, 1)

	shutdownFn := func() {
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		inferenceSvc.StopAll()
		cameraMgr.Close()
		if httpServer != nil {
			_ = httpServer.Shutdown(ctx)
		}
	}

	restartFn := func() {
		// Signal the main goroutine to perform exec after server shutdown
		select {
		case restartRequested <- struct{}{}:
		default:
		}
		shutdownFn()
	}

	// Resolve python bin (used by InstallDriver handler on Windows).
	// Priority: VISIONA_PYTHON env var (set by Wails shell) → cfg.PythonBin (--python flag)
	pythonBinForSystem := os.Getenv("VISIONA_PYTHON")
	if pythonBinForSystem == "" {
		pythonBinForSystem = cfg.PythonBin
	}

	// Create system handler with injected version and restart function
	systemHandler := handlers.NewSystemHandler(Version, BuildTime, pythonBinForSystem, restartFn)

	// Create router
	r := api.NewRouter(modelRepo, modelStore, deviceMgr, cameraMgr, flashSvc, inferenceSvc, wsHub, staticFS, logBroadcaster, systemHandler)

	// Configure HTTP server (bind to localhost only)
	addr := cfg.Addr()
	httpServer = &http.Server{
		Addr:    addr,
		Handler: r,
	}

	// Handle OS signals for graceful shutdown
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		sig := <-quit
		logger.Info("Received signal %v, shutting down gracefully...", sig)
		shutdownFn()
		os.Exit(0)
	}()

	// Kill existing process on the port if occupied
	killExistingProcess(addr, logger)

	// Start server
	logger.Info("Server listening on %s", addr)
	if err := httpServer.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
		log.Fatalf("Failed to start server: %v", err)
	}

	// If restart was requested, exec the same binary to replace this process
	select {
	case <-restartRequested:
		logger.Info("Performing self-restart via exec...")
		exe, err := os.Executable()
		if err != nil {
			log.Fatalf("Failed to get executable path: %v", err)
		}
		exe, err = filepath.EvalSymlinks(exe)
		if err != nil {
			log.Fatalf("Failed to resolve executable symlinks: %v", err)
		}
		_ = syscall.Exec(exe, os.Args, os.Environ())
		log.Fatalf("syscall.Exec failed")
	default:
		// Normal shutdown, just exit
	}
}

// killExistingProcess checks if the port is already in use and kills the
// occupying process so the server can start cleanly.
func killExistingProcess(addr string, logger *pkglogger.Logger) {
	// Extract port from addr (e.g. "127.0.0.1:3721" → "3721")
	_, port, err := net.SplitHostPort(addr)
	if err != nil {
		return
	}

	// Quick check: try to listen — if it works, port is free
	ln, err := net.Listen("tcp", addr)
	if err == nil {
		ln.Close()
		return
	}

	// Port is occupied, find and kill the process
	logger.Info("Port %s is in use, killing existing process...", port)

	var cmd *exec.Cmd
	if runtime.GOOS == "windows" {
		cmd = exec.Command("cmd", "/C", fmt.Sprintf("for /f \"tokens=5\" %%a in ('netstat -ano ^| findstr :%s') do taskkill /F /PID %%a", port))
	} else {
		cmd = exec.Command("sh", "-c", fmt.Sprintf("lsof -ti:%s | xargs kill -9 2>/dev/null", port))
	}

	output, err := cmd.CombinedOutput()
	if err != nil {
		logger.Warn("Failed to kill process on port %s: %v (%s)", port, err, strings.TrimSpace(string(output)))
		return
	}

	// Wait briefly for port to be released
	time.Sleep(500 * time.Millisecond)
	logger.Info("Previous process killed, port %s is now free", port)
}