oc-monitord/models/template.go

package models

import (
	"encoding/json"
	"fmt"
	"strings"

	w "cloud.o-forge.io/core/oc-lib/models/workflow"
	"cloud.o-forge.io/core/oc-lib/models/workflow/graph"

	"cloud.o-forge.io/core/oc-lib/models/common/enum"
	"cloud.o-forge.io/core/oc-lib/models/common/models"
	"cloud.o-forge.io/core/oc-lib/models/resources"
)

type Parameter struct {
	Name  string `yaml:"name,omitempty"`
	Value string `yaml:"value,omitempty"`
}

type Container struct {
	Image        string        `yaml:"image"`
	Command      []string      `yaml:"command,omitempty,flow"`
	Args         []string      `yaml:"args,omitempty,flow"`
	VolumeMounts []VolumeMount `yaml:"volumeMounts,omitempty"`
}

func (c *Container) AddVolumeMount(volumeMount VolumeMount, volumes []VolumeMount) []VolumeMount {
	for _, vm := range c.VolumeMounts {
		if vm.Name == volumeMount.Name {
			return volumes
		}
	}
	c.VolumeMounts = append(c.VolumeMounts, volumeMount)
	for _, vm := range c.VolumeMounts {
		for _, v := range volumes {
			if vm.Name == v.Name {
				return volumes
			}
		}
	}
	volumes = append(volumes, volumeMount)
	return volumes
}

type Task struct {
	Name         string   `yaml:"name"`
	Template     string   `yaml:"template"`
	Dependencies []string `yaml:"dependencies,omitempty"`
	Arguments    struct {
		Parameters []Parameter `yaml:"parameters,omitempty"`
	} `yaml:"arguments,omitempty"`
}

func NewTask(processingName string, graphItemID string) *Task {
	unique_name := GetArgoName(processingName, graphItemID)
	return &Task{
		Name:     unique_name,
		Template: unique_name,
	}
}

func (t *Task) BindToArgo(
	dag *Dag,
	graphItemID string,
	originWf *w.Workflow,
	processing *resources.ProcessingResource,
	firstItems, lastItems []string,
) (*Dag, []string, []string) {
	if instance := processing.GetSelectedInstance(); instance != nil {
		t.addParams(instance.(*resources.ProcessingInstance).Env)
		t.addParams(instance.(*resources.ProcessingInstance).Inputs)
		t.addParams(instance.(*resources.ProcessingInstance).Outputs)
	}
	t.Dependencies = TransformDepsToArgo(originWf.GetDependencies(graphItemID))
	name := ""
	if originWf.Graph.Items[graphItemID].Processing != nil {
		name = originWf.Graph.Items[graphItemID].Processing.GetName()
	}
	if originWf.Graph.Items[graphItemID].Workflow != nil {
		name = originWf.Graph.Items[graphItemID].Workflow.GetName()
	}
	if len(t.Dependencies) == 0 && name != "" {
		firstItems = append(firstItems, GetArgoName(name, graphItemID))
	}
	if deps := originWf.IsDependancy(graphItemID); len(deps) == 0 && name != "" {
		lastItems = append(lastItems, GetArgoName(name, graphItemID))
	}
	dag.Tasks = append(dag.Tasks, *t)
	return dag, firstItems, lastItems
}

func (t *Task) addParams(params []models.Param) {
	for _, value := range params {
		t.Arguments.Parameters = append(t.Arguments.Parameters, Parameter{
			Name:  value.Name,
			Value: value.Value,
		})
	}
}

func (t *Task) GetDeps(name string) (int, string) {
	for i, deps := range t.Dependencies {
		if strings.Contains(deps, name) {
			return i, deps
		}
	}
	return 0, ""
}

type Dag struct {
	Tasks []Task `yaml:"tasks,omitempty"`
}

func (d *Dag) GetTask(taskName string) *Task {
	for _, task := range d.Tasks {
		if strings.Contains(task.Name, taskName) {
			return &task
		}
	}
	return nil
}

type TemplateMetadata struct {
	Labels      map[string]string `yaml:"labels,omitempty"`
	Annotations map[string]string `yaml:"annotations,omitempty"`
}

type Secret struct {
	Name string `yaml:"name"`
	Key  string `yaml:"key"`
}

func NewSecret(name string, key string) *Secret {
	return &Secret{Name: name, Key: key + "-key"}
}

type Key struct {
	Key             string  `yaml:"key"`
	Bucket          string  `yaml:"bucket"`
	EndPoint        string  `yaml:"endpoint"`
	Insecure        bool    `yaml:"insecure"`
	AccessKeySecret *Secret `yaml accessKeySecret`
	SecretKeySecret *Secret `yaml secretKeySecret`
}

type Artifact struct {
	Name string `yaml:"name"`
	Path string `yaml:"path"`
	S3   *Key   `yaml:"s3,omitempty"`
}

func NewArtifact(name string, rw graph.StorageProcessingGraphLink, params []models.Param, template Template) *Artifact {
	if rw.Write {
		name += "-" + rw.Destination + "-input-write"
	} else {
		name = "-" + rw.Destination + "-input-read"
	}
	return &Artifact{
		Name: name,
		Path: template.ReplacePerEnv(rw.Source, params),
	}
}

func (a *Artifact) BindToArgo(storageType enum.StorageType, rw graph.StorageProcessingGraphLink, params []models.Param, template Template) {
	if rw.Write {
		template.Outputs.Artifacts = append(template.Inputs.Artifacts, *a)
	} else {
		template.Inputs.Artifacts = append(template.Outputs.Artifacts, *a)
	}
}

func (a *Artifact) bindS3(rw graph.StorageProcessingGraphLink, params []models.Param, template Template) {
	a.S3 = &Key{
		Key:      template.ReplacePerEnv(rw.Destination+"/"+rw.FileName, params),
		Insecure: true, // temporary
	}
	/* sel := storage.GetSelectedInstance()
	if sel != nil {
		if sel.(*resources.StorageResourceInstance).Credentials != nil {
			tool, err := tools2.NewService(conf.GetConfig().Mode)
			if err != nil || tool == nil {
				logger.Error().Msg("Could not create the access secret")
			} else {
				id, err := tool.CreateAccessSecret(namespace,
					sel.(*resources.StorageResourceInstance).Credentials.Login,
					sel.(*resources.StorageResourceInstance).Credentials.Pass)
				if err == nil {
					a.S3.AccessKeySecret = NewSecret(id, "access")
					a.S3.SecretKeySecret = NewSecret(id, "secret")
				}
			}
		}
		source := sel.(*resources.StorageResourceInstance).Source
		a.S3.Key = strings.ReplaceAll(strings.ReplaceAll(a.S3.Key, source+"/", ""), source, "")
		splits := strings.Split(a.S3.EndPoint, "/")
		if len(splits) > 1 {
			a.S3.Bucket = splits[0]
			a.S3.EndPoint = strings.Join(splits[1:], "/")
		} else {
			a.S3.Bucket = splits[0]
		}
	} */
}

type InOut struct {
	Parameters []Parameter `yaml:"parameters"`
	Artifacts  []Artifact  `yaml:"artifacts,omitempty"`
}

type Template struct {
	Name      string           `yaml:"name"`
	Inputs    InOut            `yaml:"inputs,omitempty"`
	Outputs   InOut            `yaml:"outputs,omitempty"`
	Container Container        `yaml:"container,omitempty"`
	Dag       *Dag             `yaml:"dag,omitempty"`
	Metadata  TemplateMetadata `yaml:"metadata,omitempty"`
	Resource  ServiceResource  `yaml:"resource,omitempty"`
}

func (template *Template) CreateContainer(processing *resources.ProcessingResource, dag *Dag) {
	instance := processing.GetSelectedInstance()
	if instance == nil {
		return
	}
	inst := instance.(*resources.ProcessingInstance)
	container := Container{Image: inst.Access.Container.Image}
	if container.Image == "" {
		return
	}
	container.Command = []string{"sh", "-c"} // all is bash
	for _, v := range inst.Env {
		template.Inputs.Parameters = append(template.Inputs.Parameters, Parameter{Name: v.Name})
	}
	for _, v := range inst.Inputs {
		template.Inputs.Parameters = append(template.Inputs.Parameters, Parameter{Name: v.Name})
	}
	for _, v := range inst.Inputs {
		template.Outputs.Parameters = append(template.Inputs.Parameters, Parameter{Name: v.Name})
	}
	cmd := strings.ReplaceAll(inst.Access.Container.Command, container.Image, "")

	for _, a := range strings.Split(cmd, " ") {
		container.Args = append(container.Args, template.ReplacePerEnv(a, inst.Env))
	}
	for _, a := range strings.Split(inst.Access.Container.Args, " ") {
		container.Args = append(container.Args, template.ReplacePerEnv(a, inst.Env))
	}
	container.Args = []string{strings.Join(container.Args, " ")}

	template.Container = container
}

func (template *Template) ReplacePerEnv(arg string, envs []models.Param) string {
	for _, v := range envs {
		if strings.Contains(arg, v.Name) {
			value := "{{ inputs.parameters." + v.Name + " }}"
			arg = strings.ReplaceAll(arg, v.Name, value)
			arg = strings.ReplaceAll(arg, "$"+v.Name, value)
			arg = strings.ReplaceAll(arg, "$", "")
		}
	}
	return arg
}

// Add the metadata that allow Admiralty to pick up an Argo Workflow that needs to be reparted
// The value of "clustername" is the peerId, which must be replaced by the node name's for this specific execution
func (t *Template) AddAdmiraltyAnnotations(peerID, namespace string) error {
	if t.Metadata.Annotations == nil {
		t.Metadata.Annotations = make(map[string]string)
	}

	const key = "admiralty.io/multi-cluster-scheduler"

	var annotation SchedulerAnnotation

	// Parse existing annotation if it exists
	if val, ok := t.Metadata.Annotations[key]; ok && val != "" {
		if err := json.Unmarshal([]byte(val), &annotation); err != nil {
			return fmt.Errorf("failed to parse existing scheduler annotation: %w", err)
		}
	}

	// Add new affinity
	annotation.Affinities = append(annotation.Affinities, affinity{
		Cluster:   "target-" + peerID + "-" + namespace,
		Namespace: namespace,
	})

	// Encode back to JSON
	bytes, err := json.Marshal(annotation)
	if err != nil {
		return fmt.Errorf("failed to encode scheduler annotation: %w", err)
	}

	t.Metadata.Annotations[key] = string(bytes)
	return nil
}

type affinity struct {
	Cluster   string `json:"cluster"`
	Namespace string `json:"namespace"`
}

type SchedulerAnnotation struct {
	Affinities []affinity `json:"affinities"`
}