/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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.
 */

package org.apache.flink.kubernetes.operator.autoscaler;

import org.apache.flink.autoscaler.tuning.ConfigChanges;
import org.apache.flink.configuration.GlobalConfiguration;
import org.apache.flink.configuration.MemorySize;
import org.apache.flink.configuration.PipelineOptions;
import org.apache.flink.configuration.TaskManagerOptions;
import org.apache.flink.kubernetes.operator.api.FlinkDeployment;

import org.junit.jupiter.api.Test;

import java.util.LinkedHashMap;
import java.util.Map;

import static org.assertj.core.api.Assertions.assertThat;

/** Tests for KubernetesScalingRealizer. */
public class KubernetesScalingRealizerTest {

    @Test
    public void testApplyOverrides() {
        // Required to keep the test config context on legacy Flink YAML converters.
        GlobalConfiguration.setStandardYaml(false);

        KubernetesJobAutoScalerContext ctx =
                TestingKubernetesAutoscalerUtils.createContext("test", null);

        new KubernetesScalingRealizer()
                .realizeParallelismOverrides(ctx, Map.of("a", "1", "b", "2"));

        assertThat(
                        ctx.getResource()
                                .getSpec()
                                .getFlinkConfiguration()
                                .asFlatMap()
                                .get(PipelineOptions.PARALLELISM_OVERRIDES.key()))
                .satisfiesAnyOf(
                        // Currently no enforced order inside the overrides string
                        overrides -> assertThat(overrides).isEqualTo("a:1,b:2"),
                        overrides -> assertThat(overrides).isEqualTo("b:2,a:1"));
    }

    @Test
    public void testAutoscalerOverridesStringDoesNotChangeUnlessOverridesChange() {
        // Create an overrides map which returns the keys in a deterministic order
        LinkedHashMap<String, String> newOverrides = new LinkedHashMap<>();
        newOverrides.put("b", "2");
        newOverrides.put("a", "1");

        assertOverridesDoNotChange("a:1,b:2", newOverrides);
        assertOverridesDoNotChange("b:2,a:1", newOverrides);
    }

    @Test
    public void testRealizeParallelismOverridesIsNoOpWhenSpecAlreadyMatches() {
        // Required to keep the test config context on legacy Flink YAML converters.
        GlobalConfiguration.setStandardYaml(false);

        KubernetesJobAutoScalerContext ctx =
                TestingKubernetesAutoscalerUtils.createContext("test", null);
        FlinkDeployment resource = (FlinkDeployment) ctx.getResource();

        // Pre-populate the spec with the exact override string the autoscaler would produce
        // and persist it as the last reconciled spec so observeConfig is in sync.
        resource.getSpec()
                .getFlinkConfiguration()
                .put(PipelineOptions.PARALLELISM_OVERRIDES.key(), "a:1,b:2");
        resource.getStatus()
                .getReconciliationStatus()
                .serializeAndSetLastReconciledSpec(resource.getSpec(), resource);

        new KubernetesScalingRealizer()
                .realizeParallelismOverrides(ctx, Map.of("a", "1", "b", "2"));

        // Steady-state: realizer must short-circuit and leave the spec value byte-identical
        // (no NOOP rewrite that could re-order keys or churn lastReconciledSpec).
        assertThat(
                        resource.getSpec()
                                .getFlinkConfiguration()
                                .asFlatMap()
                                .get(PipelineOptions.PARALLELISM_OVERRIDES.key()))
                .isEqualTo("a:1,b:2");
    }

    @Test
    public void testRealizeParallelismOverridesReappliesWhenSpecDrifted() {
        // Required to keep the test config context on legacy Flink YAML converters.
        GlobalConfiguration.setStandardYaml(false);

        KubernetesJobAutoScalerContext ctx =
                TestingKubernetesAutoscalerUtils.createContext("test", null);
        FlinkDeployment resource = (FlinkDeployment) ctx.getResource();

        // Spec was clobbered with a stale value mid-upgrade (e.g. last-state recovery).
        resource.getSpec()
                .getFlinkConfiguration()
                .put(PipelineOptions.PARALLELISM_OVERRIDES.key(), "a:1");
        resource.getStatus()
                .getReconciliationStatus()
                .serializeAndSetLastReconciledSpec(resource.getSpec(), resource);

        new KubernetesScalingRealizer()
                .realizeParallelismOverrides(ctx, Map.of("a", "1", "b", "2"));

        // Drift was self-healed.
        assertThat(
                        resource.getSpec()
                                .getFlinkConfiguration()
                                .asFlatMap()
                                .get(PipelineOptions.PARALLELISM_OVERRIDES.key()))
                .isIn("a:1,b:2", "b:2,a:1");
    }

    @Test
    public void testRealizeConfigOverridesIsNoOpWhenSpecAlreadyMatches() {
        KubernetesJobAutoScalerContext ctx =
                TestingKubernetesAutoscalerUtils.createContext("test", null);
        FlinkDeployment resource = (FlinkDeployment) ctx.getResource();

        // Pre-populate the spec so the override is already present and the removal target
        // is already absent.
        resource.getSpec()
                .getFlinkConfiguration()
                .put(TaskManagerOptions.MANAGED_MEMORY_FRACTION.key(), "0.42");
        String initialTmMemory = resource.getSpec().getTaskManager().getResource().getMemory();

        ConfigChanges overrides = new ConfigChanges();
        overrides.addOverride(TaskManagerOptions.MANAGED_MEMORY_FRACTION, 0.42f);
        overrides.addRemoval(TaskManagerOptions.TASK_HEAP_MEMORY);
        new KubernetesScalingRealizer().realizeConfigOverrides(ctx, overrides);

        assertThat(
                        resource.getSpec()
                                .getFlinkConfiguration()
                                .asFlatMap()
                                .get(TaskManagerOptions.MANAGED_MEMORY_FRACTION.key()))
                .isEqualTo("0.42");
        assertThat(resource.getSpec().getFlinkConfiguration().asFlatMap())
                .doesNotContainKey(TaskManagerOptions.TASK_HEAP_MEMORY.key());
        // TM memory must not be churned in steady-state (the value can only change if total
        // memory tuning produced a new target; here the spec already reflects intent).
        assertThat(resource.getSpec().getTaskManager().getResource().getMemory())
                .isEqualTo(initialTmMemory);
    }

    @Test
    public void testRealizeConfigOverridesReappliesWhenSpecDrifted() {
        KubernetesJobAutoScalerContext ctx =
                TestingKubernetesAutoscalerUtils.createContext("test", null);
        FlinkDeployment resource = (FlinkDeployment) ctx.getResource();

        // Spec was clobbered with a stale value.
        resource.getSpec()
                .getFlinkConfiguration()
                .put(TaskManagerOptions.MANAGED_MEMORY_FRACTION.key(), "0.10");

        ConfigChanges overrides = new ConfigChanges();
        overrides.addOverride(TaskManagerOptions.MANAGED_MEMORY_FRACTION, 0.42f);
        new KubernetesScalingRealizer().realizeConfigOverrides(ctx, overrides);

        // Drift was self-healed.
        assertThat(
                        resource.getSpec()
                                .getFlinkConfiguration()
                                .asFlatMap()
                                .get(TaskManagerOptions.MANAGED_MEMORY_FRACTION.key()))
                .isEqualTo("0.42");
    }

    @Test
    public void testApplyMemoryOverrides() {
        KubernetesJobAutoScalerContext ctx =
                TestingKubernetesAutoscalerUtils.createContext("test", null);

        ConfigChanges overrides = new ConfigChanges();
        MemorySize memoryOverride = MemorySize.ofMebiBytes(4096);
        overrides.addOverride(TaskManagerOptions.TOTAL_PROCESS_MEMORY, memoryOverride);
        new KubernetesScalingRealizer().realizeConfigOverrides(ctx, overrides);

        assertThat(ctx.getResource()).isInstanceOf(FlinkDeployment.class);
        assertThat(
                        ((FlinkDeployment) ctx.getResource())
                                .getSpec()
                                .getTaskManager()
                                .getResource()
                                .getMemory())
                .isEqualTo(String.valueOf(memoryOverride.getBytes()));
    }

    private void assertOverridesDoNotChange(
            String currentOverrides, LinkedHashMap<String, String> newOverrides) {

        KubernetesJobAutoScalerContext ctx =
                TestingKubernetesAutoscalerUtils.createContext("test", null);
        FlinkDeployment resource = (FlinkDeployment) ctx.getResource();

        // Create resource with existing parallelism overrides
        resource.getSpec()
                .getFlinkConfiguration()
                .put(PipelineOptions.PARALLELISM_OVERRIDES.key(), currentOverrides);
        resource.getStatus()
                .getReconciliationStatus()
                .serializeAndSetLastReconciledSpec(resource.getSpec(), resource);
        resource.getSpec()
                .getFlinkConfiguration()
                .remove(PipelineOptions.PARALLELISM_OVERRIDES.key());

        new KubernetesScalingRealizer().realizeParallelismOverrides(ctx, newOverrides);

        assertThat(
                        ctx.getResource()
                                .getSpec()
                                .getFlinkConfiguration()
                                .asFlatMap()
                                .get(PipelineOptions.PARALLELISM_OVERRIDES.key()))
                .isEqualTo(currentOverrides);
    }
}
