utils

type AWSConfig struct {
	SSHKeyPath              string   `json:"ssh_key_path" mapstructure:"ssh_key_path"` // path to AWS identity key
	EnabledRegions          []string `json:"enabled_regions" mapstructure:"enabled_regions"`
	EnabledInstanceFamilies []string `json:"enabled_instance_families" mapstructure:"enabled_instance_families"`
	LaunchTemplateName      string   `json:"launch_template" mapstructure:"launch_template"`
	ImageId                 string   `json:"image_id" mapstructure:"image_id"` // AMI image id
	User                    string   `json:"user" mapstructure:"user"`         // user if using a custom AMI
}

type ActionScripts struct {
	PreDump    string `json:"pre_dump" mapstructure:"pre_dump"`
	PostDump   string `json:"post_dump" mapstructure:"post_dump"`
	PreRestore string `json:"pre_restore" mapstructure:"pre_restore"`
}

type CedanaConfig struct {
	SelfServe        bool     `json:"self_serve" mapstructure:"self_serve"`
	EnabledProviders []string `json:"enabed_providers" mapstructure:"enabled_providers"`
	// configures cedana to spin up a new instance and restore on instance failure
	KeepRunning bool `json:"keep_running" mapstructure:"keep_running"`

	// provider specific configurations
	AWSConfig        AWSConfig        `json:"aws" mapstructure:"aws"`
	PaperspaceConfig PaperspaceConfig `json:"paperspace" mapstructure:"paperspace"`

	Checkpoint    Checkpoint          `json:"checkpoint" mapstructure:"checkpoint"`
	SharedStorage SharedStorageConfig `json:"shared_storage" mapstructure:"shared_storage"`
	Connection    Connection          `json:"connection" mapstructure:"connection"`
}

type Checkpoint struct {
	HeartbeatEnabled  bool `json:"heartbeat_enabled" mapstructure:"heartbeat_enabled"`
	HeartbeatInterval int  `json:"heartbeat_interval_seconds" mapstructure:"heartbeat_interval_seconds"`
}

type Client struct {
	ProcessName          string `json:"process_name" mapstructure:"process_name"`
	LeaveRunning         bool   `json:"leave_running" mapstructure:"leave_running"`
	SignalProcessPreDump bool   `json:"signal_process_pre_dump" mapstructure:"signal_process_pre_dump"`
	SignalProcessTimeout int    `json:"signal_process_timeout" mapstructure:"signal_process_timeout"`
}

type ConfigClient struct {
	CedanaManaged bool          `json:"cedana_managed" mapstructure:"cedana_managed"`
	Client        Client        `json:"client" mapstructure:"client"`
	ActionScripts ActionScripts `json:"action_scripts" mapstructure:"action_scripts"`
	Connection    Connection    `json:"connection" mapstructure:"connection"`
	Docker        Docker        `json:"docker" mapstructure:"docker"`
	SharedStorage SharedStorage `json:"shared_storage" mapstructure:"shared_storage"`
}

type Connection struct {
	NATSUrl   string `json:"nats_url" mapstructure:"nats_url"`
	NATSPort  int    `json:"nats_port" mapstructure:"nats_port"`
	AuthToken string `json:"auth_token" mapstructure:"auth_token"`
}

type Docker struct {
	LeaveRunning  bool   `json:"leave_running" mapstructure:"leave_running"`
	ContainerName string `json:"container_name" mapstructure:"container_name"`
	CheckpointID  string `json:"checkpoint_id" mapstructure:"checkpoint_id"`
}

type PaperspaceConfig struct {
	APIKey         string   `json:"api_key" mapstructure:"api_key"`
	SSHKeyPath     string   `json:"ssh_key_path" mapstructure:"ssh_key_path"`
	EnabledRegions []string `json:"enabled_regions" mapstructure:"enabled_regions"`
	TemplateId     string   `json:"template_id" mapstructure:"template_id"`
	User           string   `json:"user" mapstructure:"user"`
}

type SharedStorage struct {
	// only useful for multi-machine checkpoint/restore
	MountPoint     string `json:"mount_point" mapstructure:"mount_point"`
	DumpStorageDir string `json:"dump_storage_dir" mapstructure:"dump_storage_dir"`
}

type SharedStorageConfig struct {
	DumpStorageDir string `json:"dump_storage_dir" mapstructure:"dump_storage_dir"`
	MountPoint     string `json:"mount_point" mapstructure:"mount_point"`
}