YunShu/main.go

package main

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"syscall"

	"hub.gaomia.site/titor/YunShu/pkg/style"
	"hub.gaomia.site/titor/YunShu/pkg/termui"
	"gopkg.in/yaml.v3"
)

const version = "2.3.0"

func init() {
	kernel32 := syscall.NewLazyDLL("kernel32.dll")
	setConsoleCP := kernel32.NewProc("SetConsoleOutputCP")
	setConsoleCP.Call(65001)
}

func printHelp() {
	fmt.Println()
	fmt.Println(style.Cyan.Render("☁ 云枢·Agent"), style.Dim.Render("v"+version))
	fmt.Println()
	fmt.Println(style.Bold.Render("用法:"))
	fmt.Println("  yunshu [命令] [查询内容]")
	fmt.Println()
	fmt.Println(style.Bold.Render("命令:"))
	fmt.Println("  onboard       交互式初始化配置")
	fmt.Println("  log           查看日志 (--top, --level, --clear, --watch)")
	fmt.Println("  help, -h      显示帮助信息")
	fmt.Println("  version, -v   显示版本号")
	fmt.Println()
	fmt.Println(style.Bold.Render("示例:"))
	fmt.Println("  yunshu \"北京今天天气\"   ", style.Dim.Render("单次天气查询"))
	fmt.Println("  yunshu                 ", style.Dim.Render("启动交互模式"))
	fmt.Println("  yunshu log             ", style.Dim.Render("查看日志"))
	fmt.Println("  yunshu log --watch     ", style.Dim.Render("实时监听日志"))
	fmt.Println("  yunshu onboard         ", style.Dim.Render("重新初始化配置"))
	fmt.Println()
	fmt.Println(style.Bold.Render("配置文件:"), "~/.config/yunshu/config.yml")
	fmt.Println()
}

func printVersion() {
	fmt.Println("yunshu", version)
}

func migrateMemoryJSON() {
	memoryPath := filepath.Join(ConfigDir(), "memory.json")
	data, err := os.ReadFile(memoryPath)
	if err != nil {
		return
	}

	var store map[string]any
	if err := json.Unmarshal(data, &store); err != nil {
		return
	}

	// personality → config/soul.md
	if v, ok := store["personality"]; ok {
		dir := filepath.Join(ConfigDir(), "config")
		os.MkdirAll(dir, 0755)
		os.WriteFile(filepath.Join(dir, "soul.md"), []byte("# AI 灵魂\n\n"+fmt.Sprint(v)+"\n"), 0644)
	}

	// dialog_context → session/dialog.yml
	if v, ok := store["dialog_context"]; ok {
		dir := filepath.Join(ConfigDir(), "session")
		os.MkdirAll(dir, 0755)
		if out, err := yaml.Marshal(v); err == nil {
			os.WriteFile(filepath.Join(dir, "dialog.yml"), out, 0644)
		}
	}

	// agent_errors → log.yml
	if v, ok := store["agent_errors"]; ok {
		if out, err := yaml.Marshal(v); err == nil {
			os.WriteFile(filepath.Join(ConfigDir(), "log.yml"), out, 0644)
		}
	}

	os.Remove(memoryPath)

	// 确保 user.md 模板存在
	ensureUserConfig()
}

func migrateFilePaths() {
	dir := ConfigDir()
	ensureSessionDir := func() string {
		sd := filepath.Join(dir, "session")
		os.MkdirAll(sd, 0755)
		return sd
	}

	readFile := func(p string) []byte {
		d, _ := os.ReadFile(p)
		return d
	}
	writeFile := func(p string, d []byte, perm os.FileMode) {
		if len(d) > 0 {
			os.WriteFile(p, d, perm)
		}
	}

	// config.yaml → config.yml
	oldYaml := filepath.Join(dir, "config.yaml")
	newYml := filepath.Join(dir, "config.yml")
	if _, err := os.Stat(oldYaml); err == nil {
		if _, err := os.Stat(newYml); os.IsNotExist(err) {
			writeFile(newYml, readFile(oldYaml), 0600)
		}
		os.Remove(oldYaml)
	}

	// session.json → session/session.json
	oldSess := filepath.Join(dir, "session.json")
	newSess := filepath.Join(ensureSessionDir(), "session.json")
	if _, err := os.Stat(oldSess); err == nil {
		if _, err := os.Stat(newSess); os.IsNotExist(err) {
			writeFile(newSess, readFile(oldSess), 0644)
		}
		os.Remove(oldSess)
	}

	// context/dialog.yaml → session/dialog.yml
	oldDlg := filepath.Join(dir, "context", "dialog.yaml")
	newDlg := filepath.Join(ensureSessionDir(), "dialog.yml")
	if _, err := os.Stat(oldDlg); err == nil {
		if _, err := os.Stat(newDlg); os.IsNotExist(err) {
			writeFile(newDlg, readFile(oldDlg), 0644)
		}
		os.Remove(oldDlg)
		os.Remove(filepath.Join(dir, "context"))
	}

	// log.yaml → log.yml
	oldLog := filepath.Join(dir, "log.yaml")
	newLog := filepath.Join(dir, "log.yml")
	if _, err := os.Stat(oldLog); err == nil {
		if _, err := os.Stat(newLog); os.IsNotExist(err) {
			writeFile(newLog, readFile(oldLog), 0644)
		}
		os.Remove(oldLog)
	}
}

func getMainAgent() *AgentDef {
	r := ScanAgents()
	def := r.GetMain("dialog")
	if def == nil {
		if m := r.ListMains(); len(m) > 0 {
			def = m[0]
		}
	}
	return def
}

func main() {
	args := os.Args[1:]

	if len(args) > 0 {
		switch args[0] {
		case "onboard":
			runOnboard()
			return
		case "log":
			runLogCmd(args[1:])
			return
		case "help", "--help", "-h":
			printHelp()
			return
		case "version", "--version", "-v":
			printVersion()
			return
		default:
			if strings.HasPrefix(args[0], "-") {
				fmt.Fprintln(os.Stderr, style.Red.Render("未知选项: "+args[0]))
				fmt.Fprintln(os.Stderr, "可用命令: onboard, log, help, version")
				os.Exit(1)
			}
		}
	}

	// 迁移：旧目录、文件路径、旧格式 — 在 LoadConfig 前执行
	migrateOldConfig()
	migrateFilePaths()
	migrateMemoryJSON()

	cfg, err := LoadConfig()
	if err != nil {
		// 如果 config.yml 也不存在，才是真没配置
		fmt.Fprintln(os.Stderr, style.Red.Render("未找到配置文件。请先运行:"))
		fmt.Fprintln(os.Stderr, "  yunshu onboard")
		os.Exit(1)
	}
	_ = cfg

	GenerateToolsYAML()

	def := getMainAgent()
	if def == nil {
		fmt.Fprintln(os.Stderr, style.Red.Render("未找到主持者 Agent (type: main)"))
		fmt.Fprintln(os.Stderr, "请检查 agents/ 目录下是否有 type: main 的 .md 文件")
		os.Exit(1)
	}
	originalSystemPrompt := def.SystemPrompt

	if len(args) > 0 {
		logToStderr = true
		subs := ScanAgents().ListSubs()
		def.SystemPrompt = originalSystemPrompt + BuildSubAgentPrompt(subs)
		ClearSession()
		query := strings.Join(args, " ")
		if err := RunAgent(def, query); err != nil {
			fmt.Fprintln(os.Stderr, style.Red.Render("错误: "+err.Error()))
			os.Exit(1)
		}
		return
	}

	logToStderr = false
	fmt.Println()
	fmt.Println(style.Cyan.Render("☁ 云枢·Agent"), style.Dim.Render("· 天气情报官"))
	fmt.Println(style.Dim.Render("  /exit 退出，// 开头的行不发给 LLM"))
	fmt.Println()
	ClearSession()

	for {
		// 热加载：每轮重新扫描 agent 文件
		r := ScanAgents()
		if d := r.GetMain("dialog"); d != nil {
			def = d
		} else if mains := r.ListMains(); len(mains) > 0 {
			def = mains[0]
		}
		def.SystemPrompt = originalSystemPrompt + BuildSubAgentPrompt(r.ListSubs())

		fmt.Print(style.Cyan.Render("❯ "))
		input := termui.ReadLine()
		input = strings.TrimSpace(input)
		if input == "" {
			continue
		}

		if strings.HasPrefix(input, "//") {
			continue
		}

		if strings.HasPrefix(input, "/log") {
			arg := strings.TrimSpace(strings.TrimPrefix(input, "/log"))
			switch arg {
			case "on":
				logToStderr = true
				fmt.Println(style.Green.Render("日志显示已开启"))
			case "off":
				logToStderr = false
				fmt.Println(style.Yellow.Render("日志显示已关闭"))
			default:
				fmt.Println("用法:")
				fmt.Println("  /log on   开启日志显示")
				fmt.Println("  /log off  关闭日志显示")
			}
			fmt.Println()
			continue
		}

		switch input {
		case "/exit", "exit", "quit":
			fmt.Println("再见！")
			fmt.Println()
			return
		case "/clear":
			ClearSession()
			fmt.Print("\033[2J\033[H")
			fmt.Println(style.Dim.Render("会话已清空"))
			fmt.Println()
			continue
		case "/help":
			fmt.Println("可用命令:")
			fmt.Println("  /exit      退出")
			fmt.Println("  /clear     清空会话")
			fmt.Println("  /log on|off 控制日志显示")
			fmt.Println("  /help      显示帮助")
			fmt.Println("  //         不发给 LLM 的注释行")
			fmt.Println()
			continue
		}

		if err := RunAgent(def, input); err != nil {
			fmt.Fprintln(os.Stderr, style.Red.Render("错误: "+err.Error()))
		}
		fmt.Println()
	}
}