/** The generic Autoscaler instance. */
public interface JobAutoScaler<KEY> {

    boolean scale(JobAutoScalerContext<KEY> context);

    /** Called when the job is deleted. */
    void cleanup(JobAutoScalerContext<KEY> context);

}

/**
 * The job autoscaler context.
 *
 * @param <KEY> The job key.
 */
 public interface JobAutoScalerContext<KEY> {
 
    // The identifier of each flink job.
    KEY getJobKey();
 
    JobID getJobID();
  
    Configuration getConfiguration();
 
    MetricGroup getMetricGroup();
 
    RestClusterClient<String> getRestClusterClient() throws Exception;

 }

public interface ScalingRealizer<KEY,
       Context extends JobAutoScalerContext<KEY>> {

   void realize(Context context,
                Map<String, String> parallelismOverrides);
}

/**
 * Handler all loggable events during scaling.
 *
 * @param <KEY>
 */
public interface AutoScalerEventHandler<KEY,
       Context extends JobAutoScalerContext<KEY>> {

   void handleEvent(Context context,
                    Type type,
                    String reason,
                    String message,
                    @Nullable String messageKey);

   enum Type {
       Normal,
       Warning
   }
}

/** It will be used store state during scaling. */
public interface AutoScalerStateStore<KEY,
       Context extends JobAutoScalerContext<KEY>> {

   void storeScalingHistory(Context jobContext, String decision);

   String getScalingHistory(Context jobContext);

   void removeScalingHistory(Context jobContext);

   void storeEvaluatedMetrics(Context jobContext, String metrics);

   String getEvaluatedMetrics(Context jobContext);

   void removeEvaluatedMetrics(Context jobContext);

   void storeParallelismOverrides(Context jobContext, String parallelismOverrides);

   String getParallelismOverrides(Context jobContext);

   void removeParallelismOverrides(Context jobContext);


   void removeInfoFromCache(Context jobContext);

   // The flush is needed because we just save data in cache for all store methods, and flush these data to the physical storage after the flush is called to improve the performance 
   void flush(Context jobContext);
}

public class KubernetesJobAutoScalerContext implements JobAutoScalerContext<ResourceID> {

   private final ResourceID resourceID;

   private final JobID jobID;

   private final Configuration configuration;

   private final MetricGroup metricGroup;

   private final SupplierWithException<RestClusterClient<String>, Exception> restClientSupplier;

   private final AbstractFlinkResource<?, ?> resource;

   public KubernetesJobAutoScalerContext(
           ResourceID resourceID,
           JobID jobID,
           Configuration configuration,
           MetricGroup metricGroup,
           SupplierWithException<RestClusterClient<String>, Exception> restClientSupplier,
           AbstractFlinkResource<?, ?> resource) {
       this.resourceID = resourceID;
       this.jobID = jobID;
       this.configuration = configuration;
       this.metricGroup = metricGroup;
       this.restClientSupplier = restClientSupplier;
       this.resource = resource;
   }

   @Override
   public ResourceID getJobKey() {
       return resourceID;
   }

   @Override
   public JobID getJobID() {
       return jobID;
   }

   @Override
   public Configuration getConfiguration() {
       return configuration;
   }

   @Override
   public MetricGroup getMetricGroup() {
       return metricGroup;
   }

   @Override
   public RestClusterClient<String> getRestClusterClient() throws Exception {
       return restClientSupplier.get();
   }

   public AbstractFlinkResource<?, ?> getResource() {
       return resource;
   }
}

public class KubernetesScalingRealizer
       implements ScalingRealizer<ResourceID, KubernetesJobAutoScalerContext> {
   @Override
   public void realize(
           KubernetesJobAutoScalerContext context,
           Map<String, String> parallelismOverrides) {
       spec.getFlinkConfiguration().put(
       PipelineOptions.PARALLELISM_OVERRIDES.key(),
       ConfigurationUtils.convertValue(parallelismOverrides, String.class));
       context.getResource().getStatus().setImmediateReconciliationNeeded(true);
   }
}

public class KubernetesAutoScalerEventHandler
       implements AutoScalerEventHandler<ResourceID, KubernetesJobAutoScalerContext> {

   private final EventRecorder eventRecorder;

   public KubernetesAutoScalerEventHandler(EventRecorder eventRecorder) {
       this.eventRecorder = eventRecorder;
   }

   @Override
   public void handleEvent(KubernetesJobAutoScalerContext context,
                           Type type,
                           String reason,
                           String message,
                           @Nullable String messageKey) {
       eventRecorder.triggerEvent(
               context.getResource(),
               EventRecorder.Type.valueOf(type.name()),
               reason,
               EventRecorder.Component.Operator,
               message,
               messageKey);
   }

}

/** The kubernetes auto scaler state store, it's based on the config map. */ public class KubernetesAutoScalerStateStore
       implements AutoScalerStateStore<ResourceID, KubernetesJobAutoScalerContext> {

   private static final Logger LOG = LoggerFactory.getLogger(KubernetesAutoScalerStateStore.class);

   private static final String LABEL_COMPONENT_AUTOSCALER = "autoscaler";
private static final String COLLECTED_METRICS_KEY = "collectedMetrics";
private static final String SCALING_HISTORY_KEY = "scalingHistory";
private static final String PARALLELISM_OVERRIDES_KEY = "parallelismOverrides";

   private final KubernetesClient kubernetesClient;

   private final ConcurrentHashMap<ResourceID, ConfigMap> cache = new ConcurrentHashMap<>();

   public KubernetesAutoScalerStateStore(KubernetesClient kubernetesClient) {
       this.kubernetesClient = kubernetesClient;
   }

   // COLLECTED_METRICS_KEY and PARALLELISM_OVERRIDES_KEY is similar to SCALING_HISTORY_KEY
   @Override
   public void storeScalingHistory(KubernetesJobAutoScalerContext jobContext, String decision) {
     getState(jobContext).put(SCALING_HISTORY_KEY, decision);
   }

   @Override
   public String getScalingHistory(KubernetesJobAutoScalerContext jobContext) {
    getState(jobContext).get(SCALING_HISTORY_KEY);
   }

   @Override
   public void removeScalingHistory(KubernetesJobAutoScalerContext jobContext) {
      getState(jobContext).remove(SCALING_HISTORY_KEY);
   }

   private Map<String, String> getState(KubernetesJobAutoScalerContext jobContext) {
       return cache.computeIfAbsent(jobContext.getJobKey(), jobKey -> getOrCreateConfigMap(jobContext)).getData();
   }

   @Override
   public void flush(KubernetesJobAutoScalerContext jobContext) {
       ConfigMap configMap = cache.get(jobContext.getJobKey());
       Preconditions.checkState(configMap != null, "The configMap shouldn't be null.");
       try {
           cache.put(jobContext.getJobKey(), kubernetesClient.resource(configMap).update());
       } catch (Exception e) {
           LOG.error(
                   "Error while updating autoscaler info configmap, invalidating to clear the cache",
                   e);
           removeInfoFromCache(jobContext);
           throw e;
       }
   }

   @Override
   public void removeInfoFromCache(KubernetesJobAutoScalerContext jobContext) {
       cache.remove(jobContext.getJobKey());
   }

   private ConfigMap getOrCreateConfigMap(KubernetesJobAutoScalerContext jobContext) {
       AbstractFlinkResource<?, ?> cr = jobContext.getResource();
       var meta = createCmObjectMeta(ResourceID.fromResource(cr));
       return getScalingInfoConfigMap(meta).orElseGet(() -> createConfigMap(cr, meta));
   }

   private ConfigMap createConfigMap(HasMetadata cr, ObjectMeta meta) {
       LOG.info("Creating scaling info config map");
       var cm = new ConfigMap();
       cm.setMetadata(meta);
       cm.addOwnerReference(cr);
       cm.setData(new HashMap<>());
       return kubernetesClient.resource(cm).create();
   }

   private ObjectMeta createCmObjectMeta(ResourceID uid) {
       var objectMeta = new ObjectMeta();
       objectMeta.setName("autoscaler-" + uid.getName());
       uid.getNamespace().ifPresent(objectMeta::setNamespace);
       objectMeta.setLabels(
               Map.of(
                       Constants.LABEL_COMPONENT_KEY,
                       LABEL_COMPONENT_AUTOSCALER,
                       Constants.LABEL_APP_KEY,
                       uid.getName()));
       return objectMeta;
   }

   private Optional<ConfigMap> getScalingInfoConfigMap(ObjectMeta objectMeta) {
       return Optional.ofNullable(
               kubernetesClient
                       .configMaps()
                       .inNamespace(objectMeta.getNamespace())
                       .withName(objectMeta.getName())
                       .get());
   }
}