/*
 * SPDX-FileCopyrightText: Copyright (c) 2022 Atalaya Tech. Inc
 * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * Modifications Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES
 */

package v1alpha1

import (
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"github.com/ai-dynamo/dynamo/deploy/operator/internal/consts"
)

// EDIT THIS FILE!  THIS IS SCAFFOLDING FOR YOU TO OWN!
// NOTE: json tags are required.  Any new fields you add must have json tags for the fields to be serialized.

// ComponentKind represents the type of underlying Kubernetes resource.
// +kubebuilder:validation:Enum=PodClique;PodCliqueScalingGroup;Deployment;LeaderWorkerSet
type ComponentKind string

const (
	// ComponentKindPodClique represents a PodClique resource.
	ComponentKindPodClique ComponentKind = "PodClique"
	// ComponentKindPodCliqueScalingGroup represents a PodCliqueScalingGroup resource.
	ComponentKindPodCliqueScalingGroup ComponentKind = "PodCliqueScalingGroup"
	// ComponentKindDeployment represents a Deployment resource.
	ComponentKindDeployment ComponentKind = "Deployment"
	// ComponentKindLeaderWorkerSet represents a LeaderWorkerSet resource.
	ComponentKindLeaderWorkerSet ComponentKind = "LeaderWorkerSet"
)

// +kubebuilder:validation:Enum=initializing;pending;successful;failed
type DGDState string

const (
	DGDStateInitializing DGDState = "initializing"
	DGDStatePending      DGDState = "pending"
	DGDStateSuccessful   DGDState = "successful"
	DGDStateFailed       DGDState = "failed"
)

// DynamoGraphDeploymentSpec defines the desired state of DynamoGraphDeployment.
type DynamoGraphDeploymentSpec struct {
	// Annotations to propagate to all child resources (PCS, DCD, Deployments, and pod templates).
	// Service-level annotations take precedence over these values.
	// +optional
	Annotations map[string]string `json:"annotations,omitempty"`
	// Labels to propagate to all child resources (PCS, DCD, Deployments, and pod templates).
	// Service-level labels take precedence over these values.
	// +optional
	Labels map[string]string `json:"labels,omitempty"`
	// PVCs defines a list of persistent volume claims that can be referenced by components.
	// Each PVC must have a unique name that can be referenced in component specifications.
	// +kubebuilder:validation:Optional
	// +kubebuilder:validation:MaxItems=100
	PVCs []PVC `json:"pvcs,omitempty"`
	// Services are the services to deploy as part of this deployment.
	// +kubebuilder:validation:Optional
	// +kubebuilder:validation:MaxProperties=25
	Services map[string]*DynamoComponentDeploymentSharedSpec `json:"services,omitempty"`
	// Envs are environment variables applied to all services in the deployment unless
	// overridden by service-specific configuration.
	// +kubebuilder:validation:Optional
	Envs []corev1.EnvVar `json:"envs,omitempty"`
	// BackendFramework specifies the backend framework (e.g., "sglang", "vllm", "trtllm").
	// +kubebuilder:validation:Enum=sglang;vllm;trtllm
	BackendFramework string `json:"backendFramework,omitempty"`

	// Restart specifies the restart policy for the graph deployment.
	// +kubebuilder:validation:Optional
	Restart *Restart `json:"restart,omitempty"`

	// TopologyConstraint is the deployment-level topology constraint.
	// When set, topologyProfile is required and names the ClusterTopology CR to use.
	// packDomain is optional here — it can be omitted when only services carry constraints.
	// Services without their own topologyConstraint inherit from this value.
	// +optional
	TopologyConstraint *SpecTopologyConstraint `json:"topologyConstraint,omitempty"`
}

type Restart struct {
	// ID is an arbitrary string that triggers a restart when changed.
	// Any modification to this value will initiate a restart of the graph deployment according to the strategy.
	// +kubebuilder:validation:Required
	// +kubebuilder:validation:MinLength=1
	ID string `json:"id"`

	// Strategy specifies the restart strategy for the graph deployment.
	// +kubebuilder:validation:Optional
	Strategy *RestartStrategy `json:"strategy,omitempty"`
}

type RestartStrategy struct {
	// Type specifies the restart strategy type.
	// +kubebuilder:validation:Enum=Sequential;Parallel
	// +kubebuilder:default=Sequential
	Type RestartStrategyType `json:"type,omitempty"`

	// Order specifies the order in which the services should be restarted.
	// +kubebuilder:validation:Optional
	Order []string `json:"order,omitempty"`
}

type RestartStrategyType string

const (
	RestartStrategyTypeSequential RestartStrategyType = "Sequential"
	RestartStrategyTypeParallel   RestartStrategyType = "Parallel"
)

// DynamoGraphDeploymentStatus defines the observed state of DynamoGraphDeployment.
type DynamoGraphDeploymentStatus struct {
	// ObservedGeneration is the most recent generation observed by the controller.
	// +optional
	ObservedGeneration int64 `json:"observedGeneration,omitempty"`
	// State is a high-level textual status of the graph deployment lifecycle.
	// +kubebuilder:default=initializing
	State DGDState `json:"state"`
	// Conditions contains the latest observed conditions of the graph deployment.
	// The slice is merged by type on patch updates.
	Conditions []metav1.Condition `json:"conditions,omitempty" patchStrategy:"merge" patchMergeKey:"type"`
	// Services contains per-service replica status information.
	// The map key is the service name from spec.services.
	// +optional
	Services map[string]ServiceReplicaStatus `json:"services,omitempty"`
	// Restart contains the status of the restart of the graph deployment.
	// +optional
	Restart *RestartStatus `json:"restart,omitempty"`
	// Checkpoints contains per-service checkpoint status information.
	// The map key is the service name from spec.services.
	// +optional
	Checkpoints map[string]ServiceCheckpointStatus `json:"checkpoints,omitempty"`
	// RollingUpdate tracks the progress of operator manged rolling updates.
	// Currently only supported for singl-node, non-Grove deployments (DCD/Deployment).
	// +optional
	RollingUpdate *RollingUpdateStatus `json:"rollingUpdate,omitempty"`
}

// ServiceCheckpointStatus contains checkpoint information for a single service.
type ServiceCheckpointStatus struct {
	// CheckpointName is the name of the associated Checkpoint CR
	// +optional
	CheckpointName string `json:"checkpointName,omitempty"`
	// IdentityHash is the computed hash of the checkpoint identity
	// +optional
	IdentityHash string `json:"identityHash,omitempty"`
	// Ready indicates if the checkpoint was visible to the worker at startup
	// +optional
	Ready bool `json:"ready,omitempty"`
}

// RestartStatus contains the status of the restart of the graph deployment.
type RestartStatus struct {
	// ObservedID is the restart ID that has been observed and is being processed.
	// Matches the Restart.ID field in the spec.
	ObservedID string `json:"observedID,omitempty"`
	// Phase is the phase of the restart.
	Phase RestartPhase `json:"phase,omitempty"`
	// InProgress contains the names of the services that are currently being restarted.
	// +optional
	InProgress []string `json:"inProgress,omitempty"`
}

type RestartPhase string

const (
	RestartPhasePending    RestartPhase = "Pending"
	RestartPhaseRestarting RestartPhase = "Restarting"
	RestartPhaseCompleted  RestartPhase = "Completed"
	RestartPhaseFailed     RestartPhase = "Failed"
	RestartPhaseSuperseded RestartPhase = "Superseded"
)

// RollingUpdatePhase represents the current phase of a rolling update.
// +kubebuilder:validation:Enum=Pending;InProgress;Completed;Failed;""
type RollingUpdatePhase string

const (
	RollingUpdatePhasePending    RollingUpdatePhase = "Pending"
	RollingUpdatePhaseInProgress RollingUpdatePhase = "InProgress"
	RollingUpdatePhaseCompleted  RollingUpdatePhase = "Completed"
	RollingUpdatePhaseNone       RollingUpdatePhase = ""
)

// RollingUpdateStatus tracks the progress of a rolling update.
type RollingUpdateStatus struct {
	// Phase indicates the current phase of the rolling update.
	// +optional
	Phase RollingUpdatePhase `json:"phase,omitempty"`

	// StartTime is when the rolling update began.
	// +optional
	StartTime *metav1.Time `json:"startTime,omitempty"`

	// EndTime is when the rolling update completed (successfully or failed).
	// +optional
	EndTime *metav1.Time `json:"endTime,omitempty"`

	// UpdatedServices is the list of services that have completed the rolling update.
	// A service is considered updated when its new replicas are all ready and old replicas are fully scaled down.
	// Only services of componentType Worker (or Prefill/Decode) are considered.
	// +optional
	UpdatedServices []string `json:"updatedServices,omitempty"`
}

// ServiceReplicaStatus contains replica information for a single service.
type ServiceReplicaStatus struct {
	// ComponentKind is the underlying resource kind (e.g., "PodClique", "PodCliqueScalingGroup", "Deployment", "LeaderWorkerSet").
	ComponentKind ComponentKind `json:"componentKind"`

	// ComponentName is the name of the primary underlying resource.
	// DEPRECATED: Use ComponentNames instead. This field will be removed in a future release.
	// During rolling updates, this reflects the new (target) component name.
	// +kubebuilder:deprecatedversion:warning="ComponentName is deprecated, view ComponentNames instead"
	ComponentName string `json:"componentName"`

	// ComponentNames is the list of underlying resource names for this service.
	// During normal operation, this contains a single name.
	// During rolling updates, this contains both old and new component names.
	// +optional
	ComponentNames []string `json:"componentNames,omitempty"`

	// Replicas is the total number of non-terminated replicas.
	// Required for all component kinds.
	// +kubebuilder:validation:Minimum=0
	Replicas int32 `json:"replicas"`

	// UpdatedReplicas is the number of replicas at the current/desired revision.
	// Required for all component kinds.
	// +kubebuilder:validation:Minimum=0
	UpdatedReplicas int32 `json:"updatedReplicas"`

	// ReadyReplicas is the number of ready replicas.
	// Populated for PodClique, Deployment, and LeaderWorkerSet.
	// Not available for PodCliqueScalingGroup.
	// When nil, the field is omitted from the API response.
	// +optional
	// +kubebuilder:validation:Minimum=0
	ReadyReplicas *int32 `json:"readyReplicas,omitempty"`

	// AvailableReplicas is the number of available replicas.
	// For Deployment: replicas ready for >= minReadySeconds.
	// For PodCliqueScalingGroup: replicas where all constituent PodCliques have >= MinAvailable ready pods.
	// Not available for PodClique or LeaderWorkerSet.
	// When nil, the field is omitted from the API response.
	// +optional
	// +kubebuilder:validation:Minimum=0
	AvailableReplicas *int32 `json:"availableReplicas,omitempty"`
}

// +kubebuilder:object:root=true
// +kubebuilder:subresource:status
// +kubebuilder:resource:shortName=dgd
// +kubebuilder:printcolumn:name="Ready",type="string",JSONPath=`.status.conditions[?(@.type=="Ready")].status`,description="Ready status of the graph deployment"
// +kubebuilder:printcolumn:name="Backend",type="string",JSONPath=`.spec.backendFramework`,description="Backend framework (sglang, vllm, trtllm)"
// +kubebuilder:printcolumn:name="Age",type="date",JSONPath=".metadata.creationTimestamp"
// DynamoGraphDeployment is the Schema for the dynamographdeployments API.
type DynamoGraphDeployment struct {
	metav1.TypeMeta   `json:",inline"`
	metav1.ObjectMeta `json:"metadata,omitempty"`

	// Spec defines the desired state for this graph deployment.
	Spec DynamoGraphDeploymentSpec `json:"spec,omitempty"`
	// Status reflects the current observed state of this graph deployment.
	Status DynamoGraphDeploymentStatus `json:"status,omitempty"`
}

func (s *DynamoGraphDeployment) SetState(state DGDState) {
	s.Status.State = state
}

// GetState returns the current lifecycle state
func (d *DynamoGraphDeployment) GetState() string {
	return string(d.Status.State)
}

// +kubebuilder:object:root=true

// DynamoGraphDeploymentList contains a list of DynamoGraphDeployment.
type DynamoGraphDeploymentList struct {
	metav1.TypeMeta `json:",inline"`
	metav1.ListMeta `json:"metadata,omitempty"`
	Items           []DynamoGraphDeployment `json:"items"`
}

func init() {
	SchemeBuilder.Register(&DynamoGraphDeployment{}, &DynamoGraphDeploymentList{})
}

func (s *DynamoGraphDeployment) GetSpec() any {
	return s.Spec
}

func (s *DynamoGraphDeployment) SetSpec(spec any) {
	s.Spec = spec.(DynamoGraphDeploymentSpec)
}

func (s *DynamoGraphDeployment) AddStatusCondition(condition metav1.Condition) {
	if s.Status.Conditions == nil {
		s.Status.Conditions = []metav1.Condition{}
	}
	// Check if condition with same type already exists
	for i, existingCondition := range s.Status.Conditions {
		if existingCondition.Type == condition.Type {
			// Replace the existing condition
			s.Status.Conditions[i] = condition
			return
		}
	}
	// If no matching condition found, append the new one
	s.Status.Conditions = append(s.Status.Conditions, condition)
}

// HasAnyTopologyConstraint reports whether any topology constraint is set at any level.
func (s *DynamoGraphDeployment) HasAnyTopologyConstraint() bool {
	if s.Spec.TopologyConstraint != nil {
		return true
	}
	for _, svc := range s.Spec.Services {
		if svc != nil && svc.TopologyConstraint != nil {
			return true
		}
	}
	return false
}

// HasAnyMultinodeService reports whether any service in the graph is configured with more than one node.
func (s *DynamoGraphDeployment) HasAnyMultinodeService() bool {
	for _, svc := range s.Spec.Services {
		if svc != nil && svc.GetNumberOfNodes() > 1 {
			return true
		}
	}
	return false
}

// HasEPPService returns true if any service in the DGD has EPP component type
func (dgd *DynamoGraphDeployment) HasEPPService() bool {
	for _, component := range dgd.Spec.Services {
		if component != nil && component.ComponentType == consts.ComponentTypeEPP {
			return true
		}
	}
	return false
}

// GetDynamoNamespaceForService returns the Dynamo namespace for a given service.
func (s *DynamoGraphDeployment) GetDynamoNamespaceForService(service *DynamoComponentDeploymentSharedSpec) string {
	return ComputeDynamoNamespace(service.GlobalDynamoNamespace, s.GetNamespace(), s.GetName())
}

// GetEPPService returns the EPP service name and spec if present
func (dgd *DynamoGraphDeployment) GetEPPService() (string, *DynamoComponentDeploymentSharedSpec, bool) {
	for serviceName, component := range dgd.Spec.Services {
		if component != nil && component.ComponentType == consts.ComponentTypeEPP {
			return serviceName, component, true
		}
	}
	return "", nil, false
}
