package cli

import (
	"os"
	"os/signal"
	"syscall"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/comfy"
)

// comfyKeyEnv carries the replica's key to the gate: the environment, which
// only the engine's user and root can read, rather than the command line,
// which every user on the node can.
const comfyKeyEnv = "LLMGW_ENGINE_KEY"

func newComfyUIGateCmd() *cobra.Command {
	var cfg comfy.GateConfig
	cmd := &cobra.Command{
		Use:    "comfyui-gate [flags] [-- main.py flags]",
		Short:  "Run ComfyUI as an engine: behind a key, on the loopback (started by the agent)",
		Hidden: true,
		Long: "Started by the node agent for a replica whose engine is comfyui. ComfyUI asks\n" +
			"nobody for a key, so it listens on the loopback only and this gate on the\n" +
			"replica's port, letting through, with the replica's key, only what the router\n" +
			"asks of it. It also lays the fetched weights out where ComfyUI looks for them.",
		RunE: func(cmd *cobra.Command, args []string) error {
			ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
			defer stop()
			cfg.Key = os.Getenv(comfyKeyEnv)
			_ = os.Unsetenv(comfyKeyEnv)
			cfg.Args = args
			cfg.Log = logger()
			return comfy.RunGate(ctx, cfg)
		},
	}
	f := cmd.Flags()
	f.StringVar(&cfg.Listen, "listen", "", "the replica's address")
	f.StringVar(&cfg.Model, "model", "", "the model's name, for /v1/models")
	f.StringVar(&cfg.Weights, "weights", "", "the fetched weights: a file or a directory")
	f.StringVar(&cfg.State, "state", "", "the replica's own directory")
	f.StringVar(&cfg.Python, "python", "python3", "the interpreter of ComfyUI's environment")
	f.StringVar(&cfg.Dir, "comfyui", "", "ComfyUI's directory, where main.py is")
	for _, n := range []string{"listen", "weights", "state", "comfyui"} {
		_ = cmd.MarkFlagRequired(n)
	}
	return cmd
}