ai-status-light/cmd/monitor/main.go

package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"sort"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"AI-Status-Light/internal/api"
	"AI-Status-Light/internal/database"
	"AI-Status-Light/internal/discovery"
	"AI-Status-Light/internal/event"
	"AI-Status-Light/internal/monitor"
	mqttcli "AI-Status-Light/internal/mqtt"
)

const defaultDBPath = "./data/config.db"

var Version = "dev"

func main() {
	if len(os.Args) < 2 {
		printUsage()
		return
	}

	switch os.Args[1] {
	case "monitor":
		runMonitor(os.Args[2:])
	case "config":
		runConfig(os.Args[2:])
	case "serve":
		runServe(os.Args[2:])
	case "version":
		fmt.Printf("opencode-monitor %s\n", Version)
	default:
		printUsage()
	}
}

func printUsage() {
	fmt.Printf("opencode-monitor %s\n\n", Version)
	fmt.Println("用法: opencode-monitor <命令> [选项]")
	fmt.Println("")
	fmt.Println("命令:")
	fmt.Println("  monitor    启动监控")
	fmt.Println("  config     管理 MQTT 配置")
	fmt.Println("  serve      启动 API 服务")
	fmt.Println("  version    显示版本信息")
	fmt.Println("")
	fmt.Println("运行 'opencode-monitor <命令> -h' 查看命令帮助")
}

func runServe(args []string) {
	fs := flag.NewFlagSet("serve", flag.ExitOnError)
	addr := fs.String("addr", ":8080", "监听地址")
	dbPath := fs.String("db", defaultDBPath, "数据库路径")
	fs.Parse(args)

	db, err := database.New(*dbPath)
	if err != nil {
		fmt.Printf("打开数据库失败: %v\n", err)
		return
	}
	defer db.Close()

	server := api.New(db, *addr)
	fmt.Printf("API 服务已启动: %s\n", *addr)
	fmt.Println("接口文档:")
	fmt.Println("  GET    /api/health       - 健康检查")
	fmt.Println("  GET    /api/mqtt         - 获取所有配置")
	fmt.Println("  POST   /api/mqtt         - 创建配置")
	fmt.Println("  GET    /api/mqtt/:id     - 获取单个配置")
	fmt.Println("  PUT    /api/mqtt/:id     - 更新配置")
	fmt.Println("  DELETE /api/mqtt/:id     - 删除配置")

	if err := server.Start(); err != nil {
		fmt.Printf("服务启动失败: %v\n", err)
	}
}

func runMonitor(args []string) {
	fs := flag.NewFlagSet("monitor", flag.ExitOnError)
	host := fs.String("host", "127.0.0.1", "主机地址")
	portsFlag := fs.String("ports", "", "端口列表，逗号分隔 (如: 4096,4097,4098)")
	scanFlag := fs.String("scan", "", "扫描端口范围 (如: 4096-4100)")
	pidFlag := fs.Bool("pid", false, "通过进程号查找端口")
	intervalFlag := fs.Int("interval", 5, "动态扫描间隔(秒), 默认5")
	dbPath := fs.String("db", defaultDBPath, "数据库路径")
	apiAddr := fs.String("api-addr", "", "API 服务地址 (如: :8080)")
	fs.Parse(args)

	var scanRange *[2]int
	if *scanFlag != "" {
		parts := strings.Split(*scanFlag, "-")
		if len(parts) == 2 {
			start, err1 := strconv.Atoi(parts[0])
			end, err2 := strconv.Atoi(parts[1])
			if err1 == nil && err2 == nil {
				scanRange = &[2]int{start, end}
			}
		}
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

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

	db, err := database.New(*dbPath)
	if err != nil {
		fmt.Printf("打开数据库失败: %v\n", err)
		return
	}
	defer db.Close()

	if *apiAddr != "" {
		apiServer := api.New(db, *apiAddr)
		go func() {
			fmt.Printf("API 服务已启动: %s\n", *apiAddr)
			if err := apiServer.Start(); err != nil {
				fmt.Printf("API 服务失败: %v\n", err)
			}
		}()
	}

	var mqttClient *mqttcli.Client
	cfg, err := db.GetMQTTConfig()
	if err != nil {
		fmt.Printf("读取 MQTT 配置失败: %v\n", err)
	} else if cfg != nil {
		mqttClient = mqttcli.NewFromConfig(cfg)
		if err := mqttClient.Connect(); err != nil {
			fmt.Printf("MQTT 连接失败: %v\n", err)
			mqttClient = nil
		} else {
			defer mqttClient.Disconnect()
			fmt.Printf("MQTT 已连接: %s (主题: %s)\n", cfg.Broker, cfg.Topic)
		}
	}

	callback := createCallback(mqttClient)

	if *portsFlag != "" {
		runFixedMode(ctx, *host, *portsFlag, callback, sigChan)
		cancel()
		return
	}

	if scanRange != nil {
		fmt.Printf("扫描端口范围 %d-%d...\n", scanRange[0], scanRange[1])
	} else if *pidFlag {
		fmt.Println("通过进程号查找 OpenCode 端口...")
	} else {
		fmt.Println("查找 OpenCode 实例...")
	}

	runDynamicMode(ctx, *host, scanRange, *intervalFlag, callback, sigChan)
	cancel()
}

func runConfig(args []string) {
	if len(args) < 1 {
		printConfigUsage()
		return
	}

	dbPath := defaultDBPath
	for i, arg := range args {
		if arg == "--db" && i+1 < len(args) {
			dbPath = args[i+1]
			break
		}
	}

	db, err := database.New(dbPath)
	if err != nil {
		fmt.Printf("打开数据库失败: %v\n", err)
		return
	}
	defer db.Close()

	switch args[0] {
	case "list":
		configs, err := db.ListMQTTConfigs()
		if err != nil {
			fmt.Printf("查询失败: %v\n", err)
			return
		}
		if len(configs) == 0 {
			fmt.Println("未配置 MQTT")
			return
		}
		for _, cfg := range configs {
			status := "禁用"
			if cfg.Enabled {
				status = "启用"
			}
			auth := ""
			if cfg.Username != "" {
				auth = fmt.Sprintf(" [认证: %s]", cfg.Username)
			}
			fmt.Printf("[%d] %s | %s | %s%s | %s\n", cfg.ID, cfg.Broker, cfg.ClientID, cfg.Topic, auth, status)
		}

	case "set":
		fs := flag.NewFlagSet("config set", flag.ExitOnError)
		broker := fs.String("broker", "", "MQTT Broker 地址 (如: tcp://127.0.0.1:1883)")
		clientID := fs.String("client-id", "opencode-monitor", "MQTT 客户端 ID")
		username := fs.String("username", "", "MQTT 用户名")
		password := fs.String("password", "", "MQTT 密码")
		topic := fs.String("topic", "opencode/status", "MQTT 主题")
		enabled := fs.Bool("enabled", true, "是否启用")
		fs.Parse(args[1:])

		if *broker == "" {
			fmt.Println("必须指定 --broker")
			return
		}

		cfg := &database.MQTTConfig{
			Broker:   *broker,
			ClientID: *clientID,
			Username: *username,
			Password: *password,
			Topic:    *topic,
			Enabled:  *enabled,
		}

		if err := db.SaveMQTTConfig(cfg); err != nil {
			fmt.Printf("保存失败: %v\n", err)
			return
		}
		fmt.Println("配置已保存")

	case "delete":
		if len(args) < 2 {
			fmt.Println("必须指定配置 ID")
			return
		}
		id, err := strconv.Atoi(args[1])
		if err != nil {
			fmt.Println("无效的 ID")
			return
		}
		if err := db.DeleteMQTTConfig(id); err != nil {
			fmt.Printf("删除失败: %v\n", err)
			return
		}
		fmt.Println("配置已删除")

	default:
		printConfigUsage()
	}
}

func printConfigUsage() {
	fmt.Println("用法: opencode-monitor config <子命令> [选项]")
	fmt.Println("")
	fmt.Println("子命令:")
	fmt.Println("  list              列出所有配置")
	fmt.Println("  set               设置 MQTT 配置")
	fmt.Println("  delete <id>       删除配置")
	fmt.Println("")
	fmt.Println("选项:")
	fmt.Println("  --broker          MQTT Broker 地址")
	fmt.Println("  --client-id       MQTT 客户端 ID")
	fmt.Println("  --username        MQTT 用户名")
	fmt.Println("  --password        MQTT 密码")
	fmt.Println("  --topic           MQTT 主题")
	fmt.Println("  --enabled         是否启用 (true/false)")
	fmt.Println("  --db              数据库路径")
}

func createCallback(mqttClient *mqttcli.Client) monitor.EventCallback {
	return func(port int, evt *event.SSEEvent) {
		msg := event.FormatEvent(port, evt)
		if msg != "" {
			fmt.Println(msg)
		}

		if mqttClient == nil {
			return
		}

		var status string

		switch evt.Type {
		case "session.status":
			if s, ok := evt.Properties["status"].(map[string]interface{}); ok {
				status = event.ParseStatus(s)
			}
		case "session.idle":
			status = "空闲"
		case "message.part.updated":
			if part, ok := evt.Properties["part"].(map[string]interface{}); ok {
				switch part["type"].(string) {
				case "tool":
					if st, ok := part["state"].(map[string]interface{}); ok {
						status = event.ParseToolState(st)
					}
				case "reasoning":
					status = "思考中"
				}
			}
		case "session.error":
			status = "错误"
		}

		if status != "" && msg != "" {
			if err := mqttClient.PublishRaw(mqttClient.GetTopic(), msg); err != nil {
				fmt.Printf("MQTT 发送失败: %v\n", err)
			}
		}
	}
}

func runFixedMode(ctx context.Context, host string, portsFlag string, callback monitor.EventCallback, sigChan chan os.Signal) {
	var ports []int
	for _, p := range strings.Split(portsFlag, ",") {
		p = strings.TrimSpace(p)
		if port, err := strconv.Atoi(p); err == nil {
			ports = append(ports, port)
		}
	}

	if len(ports) == 0 {
		fmt.Println("未指定端口")
		return
	}

	fmt.Printf("监控端口: %v\n", ports)
	fmt.Println("Ctrl+C 停止")
	fmt.Println(strings.Repeat("-", 40))

	var wg sync.WaitGroup
	for _, port := range ports {
		wg.Add(1)
		go func(p int) {
			defer wg.Done()
			m := monitor.New(host, p, callback)
			m.Run(ctx)
		}(port)
	}

	<-sigChan
	fmt.Println("\n已停止")
	wg.Wait()
}

func runDynamicMode(ctx context.Context, host string, scanRange *[2]int, interval int, callback monitor.EventCallback, sigChan chan os.Signal) {
	scanner := discovery.NewScanner(host, scanRange)

	monitoredPorts := make(map[int]bool)
	runningMonitors := make(map[int]context.CancelFunc)
	var mu sync.Mutex

	startMonitor := func(port int) {
		mu.Lock()
		if _, exists := runningMonitors[port]; exists {
			mu.Unlock()
			return
		}
		monitorCtx, monitorCancel := context.WithCancel(ctx)
		runningMonitors[port] = monitorCancel
		mu.Unlock()

		go func() {
			fmt.Printf("开始监控端口: %d\n", port)
			m := monitor.New(host, port, callback)
			m.Run(monitorCtx)
			// 连接断开后清理记录，允许重新连接
			mu.Lock()
			delete(runningMonitors, port)
			delete(monitoredPorts, port)
			mu.Unlock()
			fmt.Printf("端口 %d 监控已停止，等待重新连接\n", port)
		}()
	}

	initial := scanner.Discover()
	if len(initial) == 0 {
		fmt.Println("未找到运行中的 OpenCode 实例")
		fmt.Println("请先执行: opencode serve --port 4096")
		fmt.Println("启动后会自动检测，等待中...")
	}

	for _, port := range initial {
		monitoredPorts[port] = true
		startMonitor(port)
	}

	if len(monitoredPorts) > 0 {
		ports := make([]int, 0, len(monitoredPorts))
		for p := range monitoredPorts {
			ports = append(ports, p)
		}
		sort.Ints(ports)
		fmt.Printf("找到 %d 个实例: %v\n", len(monitoredPorts), ports)
	}

	fmt.Printf("每 %d 秒扫描新实例，Ctrl+C 停止\n", interval)
	fmt.Println(strings.Repeat("-", 40))

	scanTicker := time.NewTicker(time.Duration(interval) * time.Second)
	defer scanTicker.Stop()

	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			case <-scanTicker.C:
				newPorts := scanner.Discover()
				for _, port := range newPorts {
					mu.Lock()
					if !monitoredPorts[port] {
						monitoredPorts[port] = true
						startMonitor(port)
					}
					mu.Unlock()
				}
			}
		}
	}()

	<-sigChan
	fmt.Println("\n已停止")
}