workloads

type DeploymentFlags struct {
	// The model to use
	Model string
}

type ExecFlags struct {
	// Run without modifying resources
	DryRun bool

	// Create namespace if it doesn't already exist
	CreateNamespace bool

	// Custom template, if any
	Template string

	// Path to workload folder
	Path string

	// OverlayPath contains specific files that override files in Path
	OverlayPath string

	// The key used to store the GPU count per node in the resource flavor
	ResourceFlavorGpuNodeLabelKey string

	// WorkloadFiles list the files that are considered to be part of the workload after merging Path and OverlayPath
	// The map is from the workload path (how the workload would see it) to the true relative path (how the CLI client sees it)
	WorkloadFiles map[string]string
}

type JobFlags struct {
	// The Kueue queue to use
	Queue string
}

type MetaFlags struct {
	// The name of the resource
	Name string

	// The namespace of the resource
	Namespace string

	// The base image to use
	Image string

	ImagePullSecret string

	Version string

	User string

	// Environment variables
	EnvVars []corev1.EnvVar

	// Secret volumes
	SecretVolumes []k8s.SecretVolume

	// Whether there is associated config map data
	HasConfigMap bool
}

type SchedulingFlags struct {
	// GPUsAvailablePerNode refers to the Kueue resource flavor's GPU count
	GPUsAvailablePerNode int

	RequestedGPUsPerReplica int

	// TotalRequestedGPUs refers to the total number of GPUs requested for the workload
	TotalRequestedGPUs int

	RequestedReplicas int

	// CalculatedGPUsPerReplica refers to the number of GPUs per replica, calculated from the available GPUs per node
	CalculatedGPUsPerReplica int

	// CalculatedNumReplicas refers to the number of replicas, calculated from the available GPUs per node
	CalculatedNumReplicas int

	Storage *StorageSchedulingFlags
}

type StorageSchedulingFlags struct {
	Quantity         string
	StorageClassName string
}

type Workload interface {
	// GenerateTemplateContext creates the workload-specific context that can be passed to the template
	GenerateTemplateContext(ExecFlags) (any, error)

	// DefaultTemplate returns a default template to use for this workload
	DefaultTemplate() ([]byte, error)

	SetFromTemplate([]byte, WorkloadTemplateConfig) error

	// IgnoreFiles lists the files that should be ignored in the ConfigMap
	IgnoreFiles() []string

	// GenerateAdditionalResourceManifests returns additional manifests that are required for creating the workload
	GenerateAdditionalResourceManifests(context.Context, client.Client, WorkloadTemplateConfig) ([]client.Object, error)

	// Reload loads the current state from Kubernetes again
	Reload(context.Context, client.Client) error

	// LoadFromObjectKey loads the state from Kubernetes given an object key (name and namespace)
	LoadFromObjectKey(context.Context, client.Client, client.ObjectKey) error

	// ResolveStructure loads the current expanded structure from k8s to make it easier to interact with the workload
	ResolveStructure(ctx context.Context, k8sClient client.Client) error

	// ListKnownPods returns the pods that the reference is currently aware of
	ListKnownPods() []WorkloadPod

	// GetStatus returns a human-readable status of the workload
	GetStatus() string

	GetName() string

	GetObject() client.Object

	GetNamespace() string

	GetKaiwoUser() string

	GetServices(ctx context.Context, k8sClient client.Client) ([]corev1.Service, error)
}

type WorkloadBase struct{}

type WorkloadPod struct {
	Pod          corev1.Pod
	LogicalGroup string
}

type WorkloadTemplateConfig struct {
	// Workload-specific config
	Workload any

	// Workload type specific config
	WorkloadMeta any

	// Custom configuration from the user
	Custom any

	// Scheduling configuration for Kubernetes
	Scheduling SchedulingFlags

	// Meta configuration for Kubernetes
	Meta MetaFlags
}