cmd/sync/main.go

// sync 是本地 ↔ Hugging Face 中转之间的凭证/状态同步工具。
//
// 用法（典型）：
//
//	wb2api-sync status          # 只看差异，不写任何东西（有差异时退出码 2）
//	wb2api-sync pull            # 远端 → 本地（本地启动器在拉起网关前先跑这个）
//	wb2api-sync push            # 本地 → 远端
//	wb2api-sync sync            # 双向对账（缺的补上，冲突按判据取新的一方）
//	wb2api-sync watch -interval 5m   # 常驻，周期性 sync
//
// 配置来源：默认读 ./config.json 的 auth_dir / state_file 与 sync 段；
// 命令行标志可逐项覆盖（未给 config.json 时也能用标志 + 环境变量跑起来）。
//
// 为什么读取走桶、写入可走网关：见 internal/syncer 包注释（Xet 上传协议）。
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/linguo2625469/workbuddy2api-panel/internal/syncer"
)

// 退出码：0 一切就绪/已对齐；1 出错；2 有差异但未应用（status 模式，供脚本判分支）。
const (
	exitOK      = 0
	exitError   = 1
	exitPending = 2
)

// fileConfig 已移到 internal/syncer.LoadConfigFile：启动器（cmd/launcher）与
// 本工具必须用完全同一套"从 config.json 取同步配置"的规则，否则两边对
// auth_dir 的理解一旦漂移，就会出现"同步写 A 目录、网关读 B 目录"的诡异故障。

func main() {
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintf(os.Stderr, "[wb2api-sync] 失败: %v\n", err)
		os.Exit(exitError)
	}
}

// options 是解析后的运行参数。
type options struct {
	mode       syncer.Mode
	config     syncer.Config
	dryRun     bool
	reportOnly bool // status/diff/plan：按双向算一遍但不落地
	jsonOut    bool
	quiet      bool
	watch      bool
	interval   time.Duration
	loaded     string // 实际读到的 config.json 路径（空 = 没读）
}

func run(args []string) error {
	// 子命令可以是第一个位置参数；也允许省略（默认 sync）。
	cmd := "sync"
	if len(args) > 0 && !strings.HasPrefix(args[0], "-") {
		cmd = args[0]
		args = args[1:]
	}

	fs := flag.NewFlagSet("wb2api-sync", flag.ContinueOnError)
	fs.SetOutput(os.Stderr)

	var (
		configPath = fs.String("config", "config.json", "配置文件路径（读 auth_dir/state_file 与 sync 段）")
		noConfig   = fs.Bool("no-config", false, "不读配置文件，只用标志/环境变量")

		bucket    = fs.String("bucket", "", "HF 桶 ID，如 a3216/wb2api-storage")
		endpoint  = fs.String("endpoint", "", "HF 端点（默认 https://huggingface.co）")
		token     = fs.String("token", "", "HF token（建议改用 -token-file 或 HF_TOKEN 环境变量）")
		tokenFile = fs.String("token-file", "", "存放 HF token 的文件")

		authDir = fs.String("auth-dir", "", "本地凭证目录（默认取 config.json 的 auth_dir）")
		dataDir = fs.String("data-dir", "", "本地数据目录（默认取 state_file 所在目录）")
		include = fs.String("include-data", "", "额外同步的 data/ 文件名，逗号分隔（默认只同步凭证）")

		writer   = fs.String("writer", "", "写入通道：auto（默认）/ hf / relay")
		hfBin    = fs.String("hf-bin", "", "hf 可执行文件，默认 hf")
		relayURL = fs.String("relay-url", "", "中转网关地址（relay 通道用）")
		relayKey = fs.String("relay-key", "", "中转网关 api_key（留空时自动从桶的 config.json 读）")
		restart  = fs.String("restart-space", "", "推了新凭证后重启该 Space（如 a3216/gcli2api），让它加载新账号")
		caBundle = fs.String("ca-bundle", "", "额外 CA 证书包（PEM）；本机跑抓包代理时给 hf CLI 用")

		timeout  = fs.String("timeout", "", "单次网络超时，默认 60s")
		insecure = fs.Bool("insecure-skip-verify", false, "跳过 TLS 校验（仅供本地抓包代理排障）")

		dryRun   = fs.Bool("dry-run", false, "只展示将要做什么，不写任何东西")
		jsonOut  = fs.Bool("json", false, "以 JSON 输出结果")
		quiet    = fs.Bool("quiet", false, "只输出汇总行")
		watch    = fs.Bool("watch", false, "常驻循环同步")
		interval = fs.Duration("interval", 5*time.Minute, "watch 模式的同步间隔")
	)

	if cmd == "help" || cmd == "-h" || cmd == "--help" {
		printHelp(fs)
		return nil
	}
	if err := fs.Parse(args); err != nil {
		return err
	}

	mode, reportOnly, err := syncer.ParseCommand(cmd)
	if err != nil {
		return err
	}

	opt := options{
		mode: mode, dryRun: *dryRun || reportOnly, jsonOut: *jsonOut, quiet: *quiet,
		watch: *watch || cmd == "watch", interval: *interval, reportOnly: reportOnly,
	}
	if opt.watch {
		// watch 必须是双向的：单方向常驻没有意义。
		opt.mode = syncer.ModeSync
		opt.reportOnly = false
	}

	// 1) 读配置文件打底
	if !*noConfig && *configPath != "" {
		fc, err := syncer.LoadConfigFile(*configPath)
		switch {
		case err == nil:
			opt.loaded = *configPath
			opt.config = fc
		case os.IsNotExist(err):
			// 没有 config.json 不是错误：可以用标志 + 环境变量把同步跑起来。
		default:
			return fmt.Errorf("读取 %s 失败: %w", *configPath, err)
		}
	}

	// 2) 标志覆盖
	if *bucket != "" {
		opt.config.Bucket = *bucket
	}
	if *endpoint != "" {
		opt.config.Endpoint = *endpoint
	}
	if *token != "" {
		opt.config.Token = *token
	}
	if *tokenFile != "" {
		opt.config.TokenFile = *tokenFile
	}
	if *authDir != "" {
		opt.config.AuthDir = *authDir
	}
	if *dataDir != "" {
		opt.config.DataDir = *dataDir
	}
	if *include != "" {
		opt.config.IncludeData = splitCSV(*include)
	}
	if *writer != "" {
		opt.config.Writer = *writer
	}
	if *hfBin != "" {
		opt.config.HFBin = *hfBin
	}
	if *relayURL != "" {
		opt.config.RelayURL = *relayURL
	}
	if *relayKey != "" {
		opt.config.RelayKey = *relayKey
	}
	if *restart != "" {
		opt.config.RestartSpace = *restart
	}
	if *caBundle != "" {
		opt.config.CABundle = *caBundle
	}
	if *timeout != "" {
		opt.config.Timeout = *timeout
	}
	if *insecure {
		opt.config.InsecureSkipVerify = true
	}
	// 环境变量兜底（方便 CI / 临时用，不必把 token 落盘）
	if opt.config.Token == "" && opt.config.TokenFile == "" {
		if t := strings.TrimSpace(os.Getenv("HF_TOKEN")); t != "" {
			opt.config.Token = t
		}
	}
	opt.config.Normalize()

	if opt.config.AuthDir == "" {
		opt.config.AuthDir = "auths"
	}
	if opt.config.DataDir == "" {
		opt.config.DataDir = "data"
	}
	if opt.config.Bucket == "" {
		return fmt.Errorf("未配置桶：请在 config.json 的 sync.bucket 填桶 ID（形如 a3216/wb2api-storage），或用 -bucket")
	}

	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer cancel()

	eng, err := syncer.New(opt.config)
	if err != nil {
		return err
	}
	opt.config = eng.Config() // 拿回归一化后的值（含 token_file 解出来的 token）

	// relay 通道缺 api_key 时，从桶里的 config.json 自动取——省掉手工抄一遍密钥。
	if needsRelayKey(eng) {
		if key, kerr := eng.FetchRemoteAPIKey(ctx); kerr == nil {
			opt.config.RelayKey = key
			if eng, err = syncer.New(opt.config); err != nil {
				return err
			}
		}
	}

	if opt.watch {
		return runWatch(ctx, eng, opt)
	}
	return runOnce(ctx, eng, opt)
}

// needsRelayKey 报告当前是否"选了 relay 通道但还没 key"。
func needsRelayKey(eng *syncer.Engine) bool {
	return eng.WriterName() == "relay" && eng.Config().RelayKey == ""
}

// runOnce 跑一轮并输出结果。
func runOnce(ctx context.Context, eng *syncer.Engine, opt options) error {
	plan, res, err := eng.Sync(ctx, opt.mode, opt.dryRun)
	if err != nil {
		return err
	}
	if opt.jsonOut {
		printJSON(plan, res, opt)
	} else {
		printHuman(plan, res, opt, eng)
	}
	if res.Failed > 0 {
		return fmt.Errorf("%d 个文件同步失败", res.Failed)
	}
	// 推了新账号 → 让 Space 重启一次，否则它只在启动时扫凭证目录，加载不到新号。
	if !opt.dryRun && opt.config.RestartSpace != "" {
		if n := uploadedAuths(res); n > 0 {
			if err := syncer.RestartSpace(ctx, opt.config.Endpoint, opt.config.RestartSpace,
				opt.config.Token, 90*time.Second, opt.config.InsecureSkipVerify); err != nil {
				return fmt.Errorf("凭证已上传，但重启 Space 失败（新账号要等它下次重启才生效）: %w", err)
			}
			if !opt.jsonOut {
				fmt.Printf("\n已请求重启 Space %s（推送了 %d 个凭证）——冷启动约 30~60s 后新账号生效。\n",
					opt.config.RestartSpace, n)
			}
		}
	}
	// status 模式：有差异时用退出码 2 告诉脚本"尚未对齐"（与 git diff --exit-code 同风格）。
	if opt.reportOnly && plan.Changes() > 0 {
		os.Exit(exitPending)
	}
	return nil
}

// uploadedAuths 统计本轮真正写成功的 auths/ 文件数——只有这些才需要触发 Space 重启。
func uploadedAuths(res *syncer.Result) int {
	n := 0
	for _, a := range res.Actions {
		if a.Op == syncer.OpUpload && a.Applied && a.Err == "" && strings.HasPrefix(a.Path, "auths/") {
			n++
		}
	}
	return n
}

// runWatch 常驻循环。只有发生变化时才打印，避免刷屏。
func runWatch(ctx context.Context, eng *syncer.Engine, opt options) error {
	if !opt.jsonOut {
		fmt.Printf("[wb2api-sync] watch 启动：每 %s 同步一次（Ctrl-C 退出）\n", opt.interval)
	}
	t := time.NewTicker(opt.interval)
	defer t.Stop()
	for {
		plan, res, err := eng.Sync(ctx, syncer.ModeSync, false)
		switch {
		case err != nil:
			fmt.Fprintf(os.Stderr, "[wb2api-sync] %s 同步出错: %v\n", time.Now().Format("15:04:05"), err)
		case plan.Changes() > 0 || res.Failed > 0:
			if opt.jsonOut {
				printJSON(plan, res, opt)
			} else {
				fmt.Printf("[wb2api-sync] %s 上传 %d / 下载 %d / 失败 %d\n",
					time.Now().Format("15:04:05"), res.Uploaded, res.Downloaded, res.Failed)
			}
		}
		select {
		case <-ctx.Done():
			return nil
		case <-t.C:
		}
	}
}

// ---------------------------------------------------------------------------
// 输出
// ---------------------------------------------------------------------------

func printHuman(plan *syncer.Plan, res *syncer.Result, opt options, eng *syncer.Engine) {
	cfg := eng.Config()
	mode := string(opt.mode)
	if opt.dryRun {
		mode += "（dry-run，未写入）"
	}
	if !opt.quiet {
		fmt.Printf("[wb2api-sync] 桶 %s  端点 %s\n", cfg.Bucket, cfg.Endpoint)
		fmt.Printf("  模式 %s   写入通道 %s   凭证目录 %s\n", mode, writerLabel(eng), cfg.AuthDir)
		if len(cfg.IncludeData) > 0 {
			fmt.Printf("  额外同步 data/: %s\n", strings.Join(cfg.IncludeData, ", "))
		} else {
			fmt.Printf("  额外同步 data/: 无（只同步 auths/ 凭证；需要时用 -include-data state.json 开启）\n")
		}
		fmt.Printf("  本地 %d 个文件 / 远端 %d 个文件\n\n", plan.LocalFiles, plan.RemoteFiles)
	}

	for _, a := range plan.Actions {
		if a.Op == syncer.OpNoop && !a.Conflict {
			continue // 一致的文件不刷屏
		}
		icon := "="
		switch a.Op {
		case syncer.OpUpload:
			icon = "↑"
		case syncer.OpDownload:
			icon = "↓"
		}
		tag := ""
		if a.Conflict {
			tag = " [冲突]"
		}
		status := ""
		if a.Err != "" {
			status = "  ✗ " + a.Err
		}
		fmt.Printf("  %s %-8s %s%s\n      %s%s\n", icon, a.Op, a.Path, tag, a.Reason, status)
	}

	if n := len(plan.Conflicts()); n > 0 && !opt.quiet {
		fmt.Printf("\n  本轮有 %d 个真冲突（两端都有且内容不同），已按「accessToken 新鲜度 → refreshToken 有无 → mtime」裁决。\n", n)
	}
	fmt.Printf("\n汇总: 上传 %d / 下载 %d / 跳过 %d / 失败 %d\n",
		res.Uploaded, res.Downloaded, res.Skipped, res.Failed)
	for _, e := range res.Errors {
		fmt.Fprintf(os.Stderr, "  ! %s\n", e)
	}
}

func writerLabel(eng *syncer.Engine) string {
	if n := eng.WriterName(); n != "" {
		return n
	}
	return "(只读：未配置写入通道)"
}

// jsonReport 是 -json 的输出结构。
type jsonReport struct {
	OK       bool            `json:"ok"`
	Mode     string          `json:"mode"`
	DryRun   bool            `json:"dry_run"`
	Bucket   string          `json:"bucket"`
	Writer   string          `json:"writer"`
	Local    int             `json:"local_files"`
	Remote   int             `json:"remote_files"`
	Changes  int             `json:"changes"`
	Uploaded int             `json:"uploaded"`
	Downld   int             `json:"downloaded"`
	Skipped  int             `json:"skipped"`
	Failed   int             `json:"failed"`
	Actions  []syncer.Action `json:"actions"`
	Errors   []string        `json:"errors,omitempty"`
}

func printJSON(plan *syncer.Plan, res *syncer.Result, opt options) {
	rep := jsonReport{
		OK: res.Failed == 0, Mode: string(opt.mode), DryRun: opt.dryRun,
		Bucket: plan.Bucket, Writer: plan.Writer,
		Local: plan.LocalFiles, Remote: plan.RemoteFiles, Changes: plan.Changes(),
		Uploaded: res.Uploaded, Downld: res.Downloaded, Skipped: res.Skipped, Failed: res.Failed,
		Actions: plan.Actions, Errors: res.Errors,
	}
	raw, _ := json.MarshalIndent(rep, "", "  ")
	fmt.Println(string(raw))
}

func printHelp(fs *flag.FlagSet) {
	fmt.Print(`wb2api-sync — 本地与 Hugging Face 中转之间的凭证/状态同步

用法:
  wb2api-sync status [标志]   只看差异（有差异退出码 2），不写任何东西
  wb2api-sync pull   [标志]   远端 → 本地
  wb2api-sync push   [标志]   本地 → 远端
  wb2api-sync sync   [标志]   双向对账（默认）
  wb2api-sync watch  [标志]   常驻，周期性 sync

退出码: 0 已对齐 / 1 出错 / 2 status 模式下存在差异

标志:
`)
	fs.PrintDefaults()
	fmt.Print(`
示例:
  # 本地启动前先拉一次远端凭证
  wb2api-sync pull -config config.json

  # 看两边差在哪（不写盘）
  wb2api-sync status -bucket a3216/wb2api-storage -token-file ./data/hf_token.txt

  # 常驻，每 5 分钟对账一次
  wb2api-sync watch -interval 5m
`)
}

// ---------------------------------------------------------------------------
// 小工具
// ---------------------------------------------------------------------------

// splitCSV 切分逗号列表并去掉空项。
func splitCSV(s string) []string {
	var out []string
	for _, part := range strings.Split(s, ",") {
		if p := strings.TrimSpace(part); p != "" {
			out = append(out, p)
		}
	}
	return out
}