executor

type ExecutionContext struct {
	JobID         string
	LeaseID       int64      // Carried from K8Decision for fencing validation
	LocalDecision K8Decision // From Layer 6 (Local Scheduler)
	PodName       string
	Namespace     string
	StartTime     time.Time
	Timeout       time.Duration
	Status        JobStatus
	LastUpdated   time.Time
	CurrentPhase  PodPhase
	NodeID        string                 // From LocalSchedulingDecision
	GPUIndices    []int                  // From LocalSchedulingDecision
	Logs          string                 // Job logs (if collected)
	Metrics       map[string]interface{} // Pod metrics
}

type ExecutionResult struct {
	JobID        string
	PodName      string
	Status       JobStatus
	Phase        PodPhase
	StartTime    time.Time
	EndTime      time.Time
	Duration     time.Duration
	ExitCode     int
	ErrorMessage string
	Logs         string
	Metrics      map[string]interface{}
	CompletedAt  time.Time
}

type Executor struct {
	ClusterID    string
	ControlPlane string
	Log          *logger.Logger
	Config       *ExecutorConfig

	JobStore     job.JobStore
	LeaseManager *lease.LeaseManager

	// Job tracking
	JobsMu        sync.RWMutex
	ActiveJobs    map[string]*ExecutionContext // JobID -> ExecutionContext
	CompletedJobs map[string]*ExecutionResult  // JobID -> ExecutionResult

	// Metrics (atomic for thread-safety)
	TotalJobs       uint64
	TotalSuccessful uint64
	TotalFailed     uint64
	TotalCancelled  uint64
	TotalDuration   int64 // nanoseconds, atomic

	// Pod management
	PodMu       sync.RWMutex
	PodRegistry map[string]*PodInfo // PodName -> PodInfo

	K8sClient K8sClient

	OnJobRunning func(jobID string)
	// Callback when job completes (for resource release)
	OnJobComplete func(jobID string, nodeID string, gpuCount int, memoryMB int)
}

type ExecutorConfig struct {
	ClusterID                string        // e.g., "cluster-us-west-2a"
	Namespace                string        // e.g., "default" or "ares-jobs"
	DefaultTimeout           time.Duration // Pod timeout (default: 1 hour)
	DefaultMemoryMB          int           // Pod default memory
	DefaultCPUMillis         int           // Pod default CPU (millicores)
	HealthCheckInterval      time.Duration // Pod health check frequency
	MaxConcurrentJobs        int           // Max concurrent Pods
	ImageRegistry            string        // e.g., "docker.io", "ghcr.io"
	DefaultJobImage          string        // e.g., "ares-job:latest"
	RestartPolicy            string        // Always, OnFailure, Never
	ImagePullPolicy          string        // Always, IfNotPresent, Never
	EnableGPUSupport         bool          // Enable GPU requests
	LogCollectionEnabled     bool          // Collect Pod logs
	MetricsCollectionEnabled bool          // Collect Pod metrics
}

type JobStatus string

type K8Decision struct {
	JobID             string
	NodeID            string
	GPUIndices        []int
	NodeScore         float64
	GPUAffinityScore  float64
	PlacementReasons  []string
	ScheduledAt       time.Time
	Command           []string
	Args              []string
	Image             string
	LeaseID           int64  // Fencing: Prevents split-brain at pod level
	FencingToken      string // Fencing: Human-readable fencing token for validation
	CheckpointEnabled bool
	CheckpointPath    string // Where to write checkpoints
	CheckpointRestore string // Last checkpoint to restore from (empty = fresh)
	CheckpointMeta    string // Metadata from last checkpoint

	// Gang scheduling: rank/world-size for distributed frameworks (torchrun,
	// Horovod, DeepSpeed). The gang manager already knows each member's index
	// and the gang size; these surface that into the pod env so the training
	// entrypoint can self-configure. Zero-valued for non-gang jobs.
	// NOTE: MASTER_ADDR for rendezvous is intentionally NOT here yet — it needs
	// a per-gang headless Service so rank-0's pod resolves via stable DNS.
	GangID        string // Gang identifier (empty = not a gang member)
	GangSize      int    // Total members in the gang (world size)
	GangMemberIdx int    // 0-based rank of this member within the gang
}

type K8sClient interface {
	CreatePod(ctx context.Context, pod *PodSpec) (podName string, err error)
	GetPod(ctx context.Context, podName string) (*PodInfo, error)
	DeletePod(ctx context.Context, podName string, gracePeriodSeconds int64) error
	ListPods(ctx context.Context) ([]*PodInfo, error)
	GetPodLogs(ctx context.Context, podName string) (string, error)
	GetPodMetrics(ctx context.Context, podName string) (map[string]interface{}, error)
	WatchPod(ctx context.Context, podName string, callback func(*PodInfo)) error
}

type MockK8sClient struct {
	// contains filtered or unexported fields
}

type PodInfo struct {
	PodName       string
	Namespace     string
	JobID         string
	NodeID        string
	Phase         PodPhase
	CreatedAt     time.Time
	StartedAt     time.Time
	FinishedAt    time.Time
	ContainerID   string
	GPUIndices    []int
	ResourceUsage map[string]interface{}
	Ready         bool
}

type PodPhase string

type PodSpec struct {
	PodName         string
	Namespace       string
	Image           string
	Command         []string
	Args            []string
	ImagePullPolicy string
	EnvVars         map[string]string
	MemoryMB        int
	CPUMillis       int
	GPUCount        int
	GPUIndices      []int
	Timeout         time.Duration
	RestartPolicy   string
	NodeID          string
	Labels          map[string]string
}