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
	regWorkspaceRoot      string
	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: func(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 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")
			}

			var nodeDef *config.NodeDefinition
			exists := false
			for i, n := range cfg.Nodes {
				if n.ID == nodeID {
					nodeDef = &cfg.Nodes[i]
					exists = true
					break
				}
			}

			if exists {
				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 cmd.Flags().Changed("workspace-root") {
					nodeDef.Runtime.WorkspaceRoot = regWorkspaceRoot
				}

				if nodeDef.Token == "" {
					nodeDef.Token = tokenGenerator()
				}
			} else {
				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
				}

				workspaceRoot := ""
				if cmd.Flags().Changed("workspace-root") {
					workspaceRoot = regWorkspaceRoot
				}

				newDef := config.NodeDefinition{
					ID:    nodeID,
					Alias: alias,
					Token: token,
					Adapters: config.AdaptersConf{
						Ollama: ollama,
						CLI:    cli,
					},
					Runtime: config.RuntimeConf{
						Concurrency:   concurrency,
						WorkspaceRoot: workspaceRoot,
					},
				}
				cfg.Nodes = append(cfg.Nodes, newDef)
				nodeDef = &cfg.Nodes[len(cfg.Nodes)-1]
			}

			target := runtime.GOOS + "-" + runtime.GOARCH
			if cmd.Flags().Changed("target") {
				target = regTarget
			}

			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)
			}

			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)
			}

			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)
			}

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

			return nil
		},
	}

	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, "runtime concurrency limit")
	c.Flags().StringVar(&regWorkspaceRoot, "workspace-root", "", "workspace root path")
	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
}

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)
		}
	}
}