bj_power/bj_power_mes/bj_power_mes.go

package main

import (
	"context"
	"embed"
	"errors"
	"flag"
	"fmt"
	"io/fs"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"bj_power_mes/internal/handler"
	"bj_power_mes/internal/upload"

	xhttp "bj_power_mes/common/httpx"
	xlog "bj_power_mes/common/logx"
	"bj_power_mes/constants"
	"bj_power_mes/ent/workorder"
	"bj_power_mes/internal/config"
	"bj_power_mes/internal/svc"

	"github.com/dromara/carbon/v2"
	"github.com/getlantern/systray"
	"github.com/pkg/browser"
	"github.com/zeromicro/go-zero/core/conf"
	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/rest"
	"github.com/zeromicro/go-zero/rest/httpx"
	"golang.org/x/sys/windows"
)

//go:generate go install -v github.com/josephspurrier/goversioninfo/cmd/goversioninfo
//go:generate goversioninfo -icon=app.ico -manifest=app.manifest -64 -o=app.syso

const AppName = "back_cover_pms"

var configFile = flag.String("f", "etc/bj_power_mes-api.yaml", "the config file")

// AppPath 优先用 os.Executable()，正确处理通过快捷方式 / PATH 启动的情况
// 开机自启时 os.Args[0] 可能不是绝对路径
var AppPath = func() string {
	if exe, err := os.Executable(); err == nil {
		if abs, err := filepath.Abs(exe); err == nil {
			return abs
		}
	}
	abs, _ := filepath.Abs(os.Args[0])
	return abs
}()

var baseUrl string
var svcCtx *svc.ServiceContext
var server *rest.Server

//go:embed public
var public embed.FS

//go:embed app.ico
var icon []byte

func main() {
	// 切换工作目录到 exe 所在目录
	// 开机自启时 CWD 默认是 system32，会导致 etc/bj_power_mes-api.yaml / logs 等相对路径失效
	if exe, err := os.Executable(); err == nil {
		if dir := filepath.Dir(exe); dir != "" {
			_ = os.Chdir(dir)
		}
	}

	flag.Parse()

	// carbon 默认配置
	carbon.SetDefault(carbon.Default{
		Layout:       carbon.DateTimeLayout,
		Timezone:     carbon.PRC,
		Locale:       "zh-CN",
		WeekStartsAt: carbon.Monday,
	})

	var c config.Config
	conf.MustLoad(*configFile, &c)

	baseUrl = fmt.Sprintf("http://127.0.0.1:%d", c.Port)

	logger := xlog.NewLogger(c.Logger)
	xlog.SetDefault(logger)
	logx.SetWriter(logger)

	mutex, err := createMutex(AppName)
	if err != nil {
		if errors.Is(err, ErrAlreadyExists) {
			_ = browser.OpenURL(baseUrl)
			return
		}
		MessageBox("错误", err.Error())
		return
	}
	defer procCloseHandle.Call(mutex)

	svcCtx = svc.NewServiceContext(c)

	sub, _ := fs.Sub(public, "public")
	spaHandler := xhttp.NewNotFoundHandler(http.FS(sub))

	// 静态文件：检测图片
	uploadDir := svcCtx.Config.Upload.Dir
	if uploadDir == "" {
		uploadDir = "./uploads"
	}
	uploadFS := http.Dir(uploadDir)

	notFoundHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.URL.Path, "/uploads/") {
			http.StripPrefix("/uploads/", http.FileServer(uploadFS)).ServeHTTP(w, r)
			return
		}
		spaHandler.ServeHTTP(w, r)
	})

	server = rest.MustNewServer(
		c.RestConf,
		rest.WithCors("*"),
		rest.WithNotFoundHandler(notFoundHandler),
	)
	logx.DisableStat()
	logx.SetWriter(logger) //重新设置日志输出

	defer server.Stop()

	//server.Use(middleware.CheckPermission(svcCtx))

	handler.RegisterHandlers(server, svcCtx)
	// 生产业务扩展路由（BOM/工序字典/备料/PLC/拧紧/报工/半成品/AGV/追溯 + 内部API）
	handler.RegisterProductionRoutes(server, svcCtx)

	// 注册 SSE 路由
	server.AddRoute(rest.Route{
		Method:  http.MethodGet,
		Path:    "/sse",
		Handler: svcCtx.SSEHandler.Serve,
	}, rest.WithTimeout(time.Hour*60))

	server.AddRoute(rest.Route{
		Method: http.MethodPost,
		Path:   "//document/upload",
		Handler: func(w http.ResponseWriter, r *http.Request) {
			upload.HandleInspectionUpload(svcCtx, w, r)
		},
	}, rest.WithMaxBytes(10*1024*1024))

	// 设置自定义 Validator
	httpx.SetValidator(xhttp.NewValidator())
	// 设置自定义返回处理
	httpx.SetOkHandler(xhttp.OkJsonCtx)
	// 设置自定义错误处理
	httpx.SetErrorHandlerCtx(xhttp.ErrorCtx)

	fmt.Printf("Starting server at %s:%d...\n", c.Host, c.Port)
	slog.Info(fmt.Sprintf("Starting server at %s:%d...", c.Host, c.Port))

	go upload.StartCleanup(svcCtx)
	go server.Start()

	_ = browser.OpenURL(baseUrl)

	// 启动系统托盘
	systray.Run(onReady, func() {
		slog.Info("onExit")
		if server != nil {
			server.Stop()
		}
		os.Exit(0)
	})
}

func checkCanShutdown() bool {
	if svcCtx == nil {
		return true
	}
	count, err := svcCtx.EntClient.WorkOrder.Query().
		Where(workorder.StatusIn(constants.WorkOrderStatus_InProgress, constants.WorkOrderStatus_Pausing)).
		Count(context.Background())
	if err != nil {
		slog.Warn("check shutdown: query failed", "error", err)
		return true
	}
	if count > 0 {
		MessageBox("警告", fmt.Sprintf("当前有 %d 个工单正在执行（IN_PROGRESS/PAUSING），请先暂停或完成工单后再退出程序。", count))
		return false
	}
	return true
}

func onReady() {
	systray.SetIcon(icon)
	systray.SetTemplateIcon(icon, icon)
	systray.SetTitle("单元管控系统")
	systray.SetTooltip("单元管控系统")
	mOpen := systray.AddMenuItem("打开管理页面", "")
	systray.AddSeparator()

	mAutoStart := systray.AddMenuItem("开机启动", "")
	if isAutoStartEnabled(AppName) {
		mAutoStart.Check()
	} else {
		mAutoStart.Uncheck()
	}
	mQuit := systray.AddMenuItem("退出", "退出服务")

	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		for {
			select {
			case s := <-signalChan:
				slog.Info("receive signal: " + s.String())
				if checkCanShutdown() {
					systray.Quit()
				}
				return
			case <-mQuit.ClickedCh:
				if checkCanShutdown() {
					systray.Quit()
					return
				}
				// 不能退出时继续监听托盘菜单，否则再次右键点退出没反应
			case <-mOpen.ClickedCh:
				err := browser.OpenURL(baseUrl)
				if err != nil {
					logx.Error(err.Error())
				}
			case <-mAutoStart.ClickedCh:
				if mAutoStart.Checked() {
					err := removeAutoStart(AppName)
					if err != nil {
						if errors.Is(err, windows.ERROR_ACCESS_DENIED) {
							MessageBox("错误", "请用管理员权限运行此程序。")
						} else {
							MessageBox("错误", err.Error())
						}
						slog.Error("设置失败： " + err.Error())
					} else {
						mAutoStart.Uncheck()
						slog.Info("已经成功取消自动运行。")
					}
				} else {
					err := setAutoStart(AppName, AppPath)
					if err != nil {
						if errors.Is(err, windows.ERROR_ACCESS_DENIED) {
							MessageBox("错误", "请用管理员权限运行此程序。")
						} else {
							MessageBox("错误", err.Error())
							slog.Error("设置失败： " + err.Error())
						}
					} else {
						mAutoStart.Check()
						slog.Info("已经成功设置为自动运行。")
					}
				}
			}
		}
	}()
}