iop/apps/edge/internal/edgecmd/node_register.go

package edgecmd

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/google/uuid"
	"github.com/spf13/cobra"
	"gopkg.in/yaml.v3"

	"iop/packages/go/config"
	"iop/packages/go/version"
)

var (
	regAlias              string
	regToken              string
	regAdapter            string
	regOllamaBaseURL      string
	regOllamaContextSize  int
	regRuntimeConcurrency int
	regTarget             string
	regArtifactBaseURL    string
)

var tokenGenerator = func() string {
	return uuid.NewString()
}

func refreshBootstrapForExistingArtifact(cfg *config.EdgeConfig, target, artifactURL string) error {
	outputDir := ResolveBootstrapArtifactDir(cfg.Bootstrap.ArtifactDir)
	adv := resolveAdvertiseHost(cfg.Server.AdvertiseHost)
	bootstrapDir := filepath.Join(outputDir, "bootstrap")
	baseOpts := nodeBootstrapPackOptions{
		OutputDir:       outputDir,
		Version:         version.Version,
		ArtifactBaseURL: artifactURL,
		EdgeAddr:        addressFor(adv.Host, cfg.Server.Listen),
	}
	if err := writeUniversalBootstrapScript(bootstrapDir, baseOpts); err != nil {
		return err
	}
	if strings.HasPrefix(target, "windows-") {
		targetOpts := baseOpts
		targetOpts.Target = target
		psScriptPath := filepath.Join(bootstrapDir, "node-"+target+".ps1")
		if err := os.WriteFile(psScriptPath, []byte(nodeBootstrapPowerShellScript(targetOpts)), 0o644); err != nil {
			return fmt.Errorf("write powershell bootstrap script: %w", err)
		}
	}

	nodeBinary := filepath.Join(outputDir, target, "iop-node")
	if _, err := os.Stat(nodeBinary); err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return nil
		}
		return fmt.Errorf("stat node artifact: %w", err)
	}
	targetOpts := baseOpts
	targetOpts.NodeBinary = nodeBinary
	targetOpts.Target = target
	return packNodeBootstrap(targetOpts)
}

func nodeCmd() *cobra.Command {
	c := &cobra.Command{
		Use:   "node",
		Short: "Manage edge nodes",
		Long: `Manage node records stored in edge.yaml.

Use node register to create or update the Node definition and print the
one-line bootstrap command for the target host. After registration, run
iop-edge config check and then start iop-edge serve.`,
		Example: `  iop-edge node register node-silicon-ollama --adapter ollama
  iop-edge config check
  iop-edge serve`,
	}
	c.AddCommand(nodeRegisterCmd())
	return c
}

func nodeRegisterCmd() *cobra.Command {
	c := &cobra.Command{
		Use:   "register [node-id]",
		Short: "Register or update a node record",
		Long: `Register or update a node record in edge.yaml and print a one-line bootstrap command.

The command writes the node id, alias, token, adapter, and runtime settings into
edge.yaml. The printed bootstrap command is the user-facing command to run on
the Node host; the user should not create node.yaml or run iop-node serve
directly in the default flow.`,
		Example: `  iop-edge node register node-silicon-ollama --adapter ollama --ollama-base-url http://127.0.0.1:11434
  iop-edge node register node-linux-cli --adapter cli
  iop-edge config check`,
		Args: cobra.ExactArgs(1),
		RunE: runNodeRegister,
	}

	c.Flags().StringVar(&regAlias, "alias", "", "node alias (defaults to node-id)")
	c.Flags().StringVar(&regToken, "token", "", "node token (defaults to auto-generated uuid)")
	c.Flags().StringVar(&regAdapter, "adapter", "cli", "node adapter: ollama or cli")
	c.Flags().StringVar(&regOllamaBaseURL, "ollama-base-url", "", "ollama base URL")
	c.Flags().IntVar(&regOllamaContextSize, "ollama-context-size", 2048, "ollama context size")
	c.Flags().IntVar(&regRuntimeConcurrency, "runtime-concurrency", 1, "legacy runtime concurrency metadata")
	c.Flags().StringVar(&regTarget, "target", "", "bootstrap target platform (e.g. darwin-arm64, linux-amd64)")
	c.Flags().StringVar(&regArtifactBaseURL, "artifact-base-url", "", "artifact base URL for bootstrap download")

	return c
}

// runNodeRegister loads edge.yaml, applies the requested node definition, resolves
// the bootstrap target/artifact URL, persists the config, and prints the one-line
// bootstrap command. Each step is delegated to a focused helper below.
func runNodeRegister(cmd *cobra.Command, args []string) error {
	nodeID := args[0]

	path, err := ResolveConfigPath(cmd)
	if err != nil {
		return err
	}

	cfgPtr, err := config.LoadEdge(path)
	if err != nil {
		return fmt.Errorf("load config: %w", err)
	}
	cfg := *cfgPtr

	data, err := os.ReadFile(path)
	if err != nil {
		return fmt.Errorf("read config file: %w", err)
	}

	if err := validateNodeRegisterFlags(cmd); err != nil {
		return err
	}

	nodeDef := applyNodeDefinition(cmd, &cfg, nodeID)

	target := resolveRegisterTarget(cmd)
	artifactURL := resolveRegisterArtifactURL(cmd, &cfg)

	if err := validateEdgeConfig(&cfg); err != nil {
		return fmt.Errorf("validate node register: %w", err)
	}
	if err := refreshBootstrapForExistingArtifact(&cfg, target, artifactURL); err != nil {
		return fmt.Errorf("refresh bootstrap artifacts: %w", err)
	}

	if err := writeEdgeConfigFile(path, data, &cfg); err != nil {
		return err
	}

	fmt.Fprintln(cmd.OutOrStdout(), bootstrapCommand(artifactURL, target, nodeDef.Token))

	return nil
}

// validateNodeRegisterFlags rejects invalid --adapter and --runtime-concurrency values.
func validateNodeRegisterFlags(cmd *cobra.Command) error {
	if cmd.Flags().Changed("adapter") {
		if regAdapter != "ollama" && regAdapter != "cli" {
			return fmt.Errorf("invalid adapter %q: must be ollama or cli", regAdapter)
		}
	}
	if cmd.Flags().Changed("runtime-concurrency") && regRuntimeConcurrency < 0 {
		return fmt.Errorf("runtime concurrency must be non-negative")
	}
	return nil
}

// applyNodeDefinition updates the existing node record for nodeID in place, or appends
// a new one, and returns a pointer to the affected definition inside cfg.Nodes.
func applyNodeDefinition(cmd *cobra.Command, cfg *config.EdgeConfig, nodeID string) *config.NodeDefinition {
	for i := range cfg.Nodes {
		if cfg.Nodes[i].ID == nodeID {
			nodeDef := &cfg.Nodes[i]
			updateExistingNodeDefinition(cmd, nodeDef)
			return nodeDef
		}
	}

	cfg.Nodes = append(cfg.Nodes, buildNewNodeDefinition(cmd, nodeID))
	return &cfg.Nodes[len(cfg.Nodes)-1]
}

// updateExistingNodeDefinition applies only the flags the caller changed onto an
// existing node record, backfilling a token when the record has none.
func updateExistingNodeDefinition(cmd *cobra.Command, nodeDef *config.NodeDefinition) {
	if cmd.Flags().Changed("alias") {
		nodeDef.Alias = regAlias
	}
	if cmd.Flags().Changed("token") {
		nodeDef.Token = regToken
	}

	if cmd.Flags().Changed("adapter") {
		if regAdapter == "ollama" {
			nodeDef.Adapters.Ollama.Enabled = true
			nodeDef.Adapters.CLI.Enabled = false
			if nodeDef.Adapters.Ollama.BaseURL == "" && regOllamaBaseURL == "" {
				nodeDef.Adapters.Ollama.BaseURL = "http://127.0.0.1:11434"
			}
		} else if regAdapter == "cli" {
			nodeDef.Adapters.CLI.Enabled = true
			nodeDef.Adapters.Ollama.Enabled = false
		}
	}

	if cmd.Flags().Changed("ollama-base-url") {
		nodeDef.Adapters.Ollama.BaseURL = regOllamaBaseURL
	}
	if cmd.Flags().Changed("ollama-context-size") {
		nodeDef.Adapters.Ollama.ContextSize = regOllamaContextSize
	}

	if cmd.Flags().Changed("runtime-concurrency") {
		nodeDef.Runtime.Concurrency = regRuntimeConcurrency
	}

	if nodeDef.Token == "" {
		nodeDef.Token = tokenGenerator()
	}
}

// buildNewNodeDefinition constructs a fresh node record from defaults and changed flags.
func buildNewNodeDefinition(cmd *cobra.Command, nodeID string) config.NodeDefinition {
	alias := nodeID
	if cmd.Flags().Changed("alias") {
		alias = regAlias
	}

	token := regToken
	if token == "" {
		token = tokenGenerator()
	}

	adapter := "cli"
	if cmd.Flags().Changed("adapter") {
		adapter = regAdapter
	}

	var ollama config.OllamaConf
	var cli config.CLIConf

	if adapter == "ollama" {
		ollama.Enabled = true
		ollama.BaseURL = "http://127.0.0.1:11434"
		if cmd.Flags().Changed("ollama-base-url") {
			ollama.BaseURL = regOllamaBaseURL
		}
		ollama.ContextSize = 2048
		if cmd.Flags().Changed("ollama-context-size") {
			ollama.ContextSize = regOllamaContextSize
		}
	} else {
		cli.Enabled = true
	}

	concurrency := 1
	if cmd.Flags().Changed("runtime-concurrency") {
		concurrency = regRuntimeConcurrency
	}

	return config.NodeDefinition{
		ID:    nodeID,
		Alias: alias,
		Token: token,
		Adapters: config.AdaptersConf{
			Ollama: ollama,
			CLI:    cli,
		},
		Runtime: config.RuntimeConf{
			Concurrency: concurrency,
		},
	}
}

// resolveRegisterTarget returns the bootstrap target platform, defaulting to the host.
func resolveRegisterTarget(cmd *cobra.Command) string {
	target := runtime.GOOS + "-" + runtime.GOARCH
	if cmd.Flags().Changed("target") {
		target = regTarget
	}
	return target
}

// resolveRegisterArtifactURL resolves the artifact base URL, persisting an explicit
// --artifact-base-url override into cfg and falling back to the advertised default.
func resolveRegisterArtifactURL(cmd *cobra.Command, cfg *config.EdgeConfig) string {
	artifactURL := cfg.Bootstrap.ArtifactBaseURL
	if cmd.Flags().Changed("artifact-base-url") {
		artifactURL = regArtifactBaseURL
		cfg.Bootstrap.ArtifactBaseURL = regArtifactBaseURL
	}
	if artifactURL == "" {
		artifactURL = defaultArtifactBaseURL(resolveAdvertiseHost(cfg.Server.AdvertiseHost).Host, cfg.Bootstrap.Listen)
	}
	return artifactURL
}

// writeEdgeConfigFile renders the updated config back onto the original YAML document
// with preserving-patch semantics and writes it with 0600 permissions.
func writeEdgeConfigFile(path string, data []byte, cfg *config.EdgeConfig) error {
	outData, err := patchYAML(data, cfg)
	if err != nil {
		return fmt.Errorf("patch config: %w", err)
	}
	if err := os.WriteFile(path, outData, 0o600); err != nil {
		return fmt.Errorf("write config file: %w", err)
	}
	return nil
}

func patchYAML(data []byte, cfg *config.EdgeConfig) ([]byte, error) {
	var doc yaml.Node
	if err := yaml.Unmarshal(data, &doc); err != nil {
		return nil, fmt.Errorf("unmarshal yaml node: %w", err)
	}

	if len(doc.Content) == 0 || doc.Content[0].Kind != yaml.MappingNode {
		return nil, errors.New("invalid yaml structure: root is not a map")
	}

	rootMap := doc.Content[0]

	nodesBytes, err := yaml.Marshal(cfg.Nodes)
	if err != nil {
		return nil, fmt.Errorf("marshal nodes: %w", err)
	}
	var tempDoc yaml.Node
	if err := yaml.Unmarshal(nodesBytes, &tempDoc); err != nil {
		return nil, fmt.Errorf("unmarshal nodes temp doc: %w", err)
	}
	var nodesValNode *yaml.Node
	if len(tempDoc.Content) > 0 {
		nodesValNode = tempDoc.Content[0]
	} else {
		nodesValNode = &yaml.Node{Kind: yaml.SequenceNode, Tag: "!!seq"}
	}

	foundNodes := false
	for i := 0; i < len(rootMap.Content); i += 2 {
		if rootMap.Content[i].Value == "nodes" {
			rootMap.Content[i+1] = nodesValNode
			foundNodes = true
			break
		}
	}
	if !foundNodes {
		keyNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: "nodes"}
		rootMap.Content = append(rootMap.Content, keyNode, nodesValNode)
	}

	fieldsToPatch := make(map[string]string)
	if cfg.Bootstrap.ArtifactBaseURL != "" {
		fieldsToPatch["artifact_base_url"] = cfg.Bootstrap.ArtifactBaseURL
	}
	patchBootstrapStringFields(rootMap, fieldsToPatch)

	out, err := yaml.Marshal(&doc)
	if err != nil {
		return nil, fmt.Errorf("marshal patched doc: %w", err)
	}
	return out, nil
}

func patchBootstrapStringFields(rootMap *yaml.Node, fields map[string]string) {
	if len(fields) == 0 {
		return
	}

	var bootstrapVal *yaml.Node
	foundBootstrap := false
	for i := 0; i < len(rootMap.Content); i += 2 {
		if rootMap.Content[i].Value == "bootstrap" {
			bootstrapVal = rootMap.Content[i+1]
			foundBootstrap = true
			break
		}
	}

	if !foundBootstrap {
		bootstrapKey := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: "bootstrap"}
		bootstrapVal = &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}
		rootMap.Content = append(rootMap.Content, bootstrapKey, bootstrapVal)
	}

	if bootstrapVal.Kind != yaml.MappingNode {
		return
	}

	keys := []string{"artifact_base_url"}
	for _, k := range keys {
		v, ok := fields[k]
		if !ok || v == "" {
			continue
		}
		foundField := false
		for j := 0; j < len(bootstrapVal.Content); j += 2 {
			if bootstrapVal.Content[j].Value == k {
				bootstrapVal.Content[j+1].Value = v
				bootstrapVal.Content[j+1].Tag = "!!str"
				foundField = true
				break
			}
		}
		if !foundField {
			keyNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: k}
			valNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: v}
			bootstrapVal.Content = append(bootstrapVal.Content, keyNode, valNode)
		}
	}
}