/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License
 * 2.0; you may not use this file except in compliance with the Elastic License
 * 2.0.
 */

package org.elasticsearch.xpack.autoscaling.capacity;

import org.elasticsearch.Version;
import org.elasticsearch.cluster.ClusterInfo;
import org.elasticsearch.cluster.ClusterName;
import org.elasticsearch.cluster.ClusterState;
import org.elasticsearch.cluster.DiskUsage;
import org.elasticsearch.cluster.metadata.Metadata;
import org.elasticsearch.cluster.node.DiscoveryNode;
import org.elasticsearch.cluster.node.DiscoveryNodeRole;
import org.elasticsearch.cluster.node.DiscoveryNodes;
import org.elasticsearch.common.collect.ImmutableOpenMap;
import org.elasticsearch.common.collect.Tuple;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.unit.ByteSizeValue;
import org.elasticsearch.snapshots.InternalSnapshotsInfoService;
import org.elasticsearch.snapshots.SnapshotShardSizeInfo;
import org.elasticsearch.xpack.autoscaling.AutoscalingMetadata;
import org.elasticsearch.xpack.autoscaling.AutoscalingTestCase;
import org.elasticsearch.xpack.autoscaling.capacity.memory.AutoscalingMemoryInfo;
import org.elasticsearch.xpack.autoscaling.policy.AutoscalingPolicy;
import org.elasticsearch.xpack.autoscaling.policy.AutoscalingPolicyMetadata;

import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.SortedMap;
import java.util.SortedSet;
import java.util.TreeMap;
import java.util.TreeSet;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.StreamSupport;

import static org.hamcrest.Matchers.containsString;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.nullValue;
import static org.hamcrest.Matchers.sameInstance;

public class AutoscalingCalculateCapacityServiceTests extends AutoscalingTestCase {
    public void testMultiplePoliciesFixedCapacity() {
        AutoscalingCalculateCapacityService service = new AutoscalingCalculateCapacityService(
            org.elasticsearch.common.collect.Set.of(new FixedAutoscalingDeciderService())
        );
        Set<String> policyNames = IntStream.range(0, randomIntBetween(1, 10))
            .mapToObj(i -> "test_ " + randomAlphaOfLength(10))
            .collect(Collectors.toSet());

        SortedMap<String, AutoscalingPolicyMetadata> policies = new TreeMap<>(
            policyNames.stream()
                .map(s -> Tuple.tuple(s, new AutoscalingPolicyMetadata(new AutoscalingPolicy(s, randomRoles(), randomFixedDeciders()))))
                .collect(Collectors.toMap(Tuple::v1, Tuple::v2))
        );
        ClusterState state = ClusterState.builder(ClusterName.DEFAULT)
            .metadata(Metadata.builder().putCustom(AutoscalingMetadata.NAME, new AutoscalingMetadata(policies)))
            .build();
        SortedMap<String, AutoscalingDeciderResults> resultsMap = service.calculate(
            state,
            ClusterInfo.EMPTY,
            null,
            AutoscalingMemoryInfo.EMPTY
        );
        assertThat(resultsMap.keySet(), equalTo(policyNames));
        for (Map.Entry<String, AutoscalingDeciderResults> entry : resultsMap.entrySet()) {
            AutoscalingDeciderResults results = entry.getValue();
            SortedMap<String, Settings> deciders = policies.get(entry.getKey()).policy().deciders();
            assertThat(deciders.size(), equalTo(1));
            Settings configuration = deciders.values().iterator().next();
            AutoscalingCapacity requiredCapacity = calculateFixedDeciderCapacity(configuration);
            assertThat(results.requiredCapacity(), equalTo(requiredCapacity));
            assertThat(results.results().size(), equalTo(1));
            AutoscalingDeciderResult deciderResult = results.results().get(deciders.firstKey());
            assertNotNull(deciderResult);
            assertThat(deciderResult.requiredCapacity(), equalTo(requiredCapacity));
            ByteSizeValue storage = configuration.getAsBytesSize(FixedAutoscalingDeciderService.STORAGE.getKey(), null);
            ByteSizeValue memory = configuration.getAsMemory(FixedAutoscalingDeciderService.MEMORY.getKey(), null);
            int nodes = FixedAutoscalingDeciderService.NODES.get(configuration);
            assertThat(deciderResult.reason(), equalTo(new FixedAutoscalingDeciderService.FixedReason(storage, memory, nodes)));
            assertThat(
                deciderResult.reason().summary(),
                equalTo("fixed storage [" + storage + "] memory [" + memory + "] nodes [" + nodes + "]")
            );

            // there is no nodes in any tier.
            assertThat(results.currentCapacity(), equalTo(AutoscalingCapacity.ZERO));
            assertThat(results.currentNodes(), equalTo(Collections.emptySortedSet()));
        }
    }

    public void testDefaultDeciders() {
        FixedAutoscalingDeciderService defaultOn = new FixedAutoscalingDeciderService() {
            @Override
            public boolean defaultOn() {
                return true;
            }

            @Override
            public String name() {
                return "default_on";
            }
        };

        FixedAutoscalingDeciderService defaultOff = new FixedAutoscalingDeciderService();

        AutoscalingCalculateCapacityService service = new AutoscalingCalculateCapacityService(
            org.elasticsearch.common.collect.Set.of(defaultOn, defaultOff)
        );
        ClusterState state = ClusterState.builder(ClusterName.DEFAULT)
            .metadata(
                Metadata.builder()
                    .putCustom(
                        AutoscalingMetadata.NAME,
                        new AutoscalingMetadata(
                            new TreeMap<>(
                                org.elasticsearch.common.collect.Map.of(
                                    "test",
                                    new AutoscalingPolicyMetadata(new AutoscalingPolicy("test", randomRoles(), new TreeMap<>()))
                                )
                            )
                        )
                    )
            )
            .build();

        assertThat(
            service.calculate(state, ClusterInfo.EMPTY, SnapshotShardSizeInfo.EMPTY, AutoscalingMemoryInfo.EMPTY)
                .get("test")
                .results()
                .keySet(),
            equalTo(org.elasticsearch.common.collect.Set.of(defaultOn.name()))
        );
    }

    private SortedMap<String, Settings> randomFixedDeciders() {
        Settings.Builder settings = Settings.builder();
        if (randomBoolean()) {
            settings.put(FixedAutoscalingDeciderService.STORAGE.getKey(), randomByteSizeValue());
        }
        if (randomBoolean()) {
            settings.put(FixedAutoscalingDeciderService.MEMORY.getKey(), randomByteSizeValue());
        }
        settings.put(FixedAutoscalingDeciderService.NODES.getKey(), randomIntBetween(1, 10));
        return new TreeMap<>(org.elasticsearch.common.collect.Map.of(FixedAutoscalingDeciderService.NAME, settings.build()));
    }

    private AutoscalingCapacity calculateFixedDeciderCapacity(Settings configuration) {
        ByteSizeValue storage = configuration.getAsBytesSize(FixedAutoscalingDeciderService.STORAGE.getKey(), null);
        ByteSizeValue memory = configuration.getAsBytesSize(FixedAutoscalingDeciderService.MEMORY.getKey(), null);
        int nodes = FixedAutoscalingDeciderService.NODES.get(configuration);
        ByteSizeValue totalStorage = storage != null ? new ByteSizeValue(storage.getBytes() * nodes) : null;
        ByteSizeValue totalMemory = memory != null ? new ByteSizeValue(memory.getBytes() * nodes) : null;

        if (totalStorage == null && totalMemory == null) {
            return null;
        } else {
            return new AutoscalingCapacity(
                new AutoscalingCapacity.AutoscalingResources(totalStorage, totalMemory),
                new AutoscalingCapacity.AutoscalingResources(storage, memory)
            );
        }
    }

    public void testContext() {
        ClusterState state = ClusterState.builder(ClusterName.DEFAULT).build();
        ClusterInfo info = ClusterInfo.EMPTY;
        SortedSet<String> roleNames = randomRoles();
        boolean hasDataRole = roleNames.stream().anyMatch(r -> r.equals("data") || r.startsWith("data_"));

        AutoscalingCalculateCapacityService service = new AutoscalingCalculateCapacityService(
            org.elasticsearch.common.collect.Set.of(new FixedAutoscalingDeciderService())
        );
        SnapshotShardSizeInfo snapshotShardSizeInfo = new SnapshotShardSizeInfo(
            ImmutableOpenMap.<InternalSnapshotsInfoService.SnapshotShard, Long>builder().build()
        );
        AutoscalingDeciderContext context = service.createContext(
            roleNames,
            state,
            info,
            snapshotShardSizeInfo,
            n -> randomNonNegativeLong()
        );

        assertSame(state, context.state());
        assertThat(context.nodes(), equalTo(org.elasticsearch.common.collect.Set.of()));
        assertThat(context.currentCapacity(), equalTo(AutoscalingCapacity.ZERO));
        assertThat(context.info(), sameInstance(info));
        assertThat(context.snapshotShardSizeInfo(), sameInstance(snapshotShardSizeInfo));

        Set<DiscoveryNodeRole> roles = roleNames.stream().map(DiscoveryNode::getRoleFromRoleName).collect(Collectors.toSet());
        Set<DiscoveryNodeRole> otherRoles = mutateRoles(roleNames).stream()
            .map(DiscoveryNode::getRoleFromRoleName)
            .collect(Collectors.toSet());
        final long memory = between(0, 1000);
        state = ClusterState.builder(ClusterName.DEFAULT)
            .nodes(
                DiscoveryNodes.builder()
                    .add(
                        new DiscoveryNode(
                            "nodeId",
                            buildNewFakeTransportAddress(),
                            org.elasticsearch.common.collect.Map.of(),
                            roles,
                            Version.CURRENT
                        )
                    )
            )
            .build();
        context = new AutoscalingCalculateCapacityService.DefaultAutoscalingDeciderContext(roleNames, state, info, null, n -> memory);

        assertThat(context.nodes().size(), equalTo(1));
        assertThat(context.nodes(), equalTo(StreamSupport.stream(state.nodes().spliterator(), false).collect(Collectors.toSet())));
        if (hasDataRole) {
            assertNull(context.currentCapacity());
        } else {
            assertThat(context.currentCapacity().node().memory(), equalTo(new ByteSizeValue(memory)));
            assertThat(context.currentCapacity().total().memory(), equalTo(new ByteSizeValue(memory)));
            assertThat(context.currentCapacity().node().storage(), equalTo(ByteSizeValue.ZERO));
            assertThat(context.currentCapacity().total().storage(), equalTo(ByteSizeValue.ZERO));
        }

        ImmutableOpenMap.Builder<String, DiskUsage> leastUsagesBuilder = ImmutableOpenMap.builder();
        ImmutableOpenMap.Builder<String, DiskUsage> mostUsagesBuilder = ImmutableOpenMap.builder();
        DiscoveryNodes.Builder nodes = DiscoveryNodes.builder();
        Set<DiscoveryNode> expectedNodes = new HashSet<>();
        long sumTotal = 0;
        long maxTotal = 0;
        for (int i = 0; i < randomIntBetween(1, 5); ++i) {
            String nodeId = "nodeId" + i;
            boolean useOtherRoles = randomBoolean();
            DiscoveryNode node = new DiscoveryNode(
                nodeId,
                buildNewFakeTransportAddress(),
                org.elasticsearch.common.collect.Map.of(),
                useOtherRoles ? otherRoles : roles,
                Version.CURRENT
            );
            nodes.add(node);

            if (useOtherRoles == false) {
                long total = randomLongBetween(1, 1L << 40);
                DiskUsage diskUsage = new DiskUsage(nodeId, null, randomAlphaOfLength(5), total, randomLongBetween(0, total));
                leastUsagesBuilder.put(nodeId, diskUsage);
                if (randomBoolean()) {
                    diskUsage = new DiskUsage(nodeId, null, diskUsage.getPath(), total, diskUsage.getFreeBytes());
                }
                mostUsagesBuilder.put(nodeId, diskUsage);
                sumTotal += total;
                maxTotal = Math.max(total, maxTotal);
                expectedNodes.add(node);
            } else {
                long total1 = randomLongBetween(0, 1L << 40);
                leastUsagesBuilder.put(nodeId, new DiskUsage(nodeId, null, randomAlphaOfLength(5), total1, randomLongBetween(0, total1)));
                long total2 = randomLongBetween(0, 1L << 40);
                mostUsagesBuilder.put(nodeId, new DiskUsage(nodeId, null, randomAlphaOfLength(5), total2, randomLongBetween(0, total2)));
            }
        }
        state = ClusterState.builder(ClusterName.DEFAULT).nodes(nodes).build();
        ImmutableOpenMap<String, DiskUsage> leastUsages = leastUsagesBuilder.build();
        ImmutableOpenMap<String, DiskUsage> mostUsages = mostUsagesBuilder.build();
        info = new ClusterInfo(leastUsages, mostUsages, null, null, null, null);
        context = new AutoscalingCalculateCapacityService.DefaultAutoscalingDeciderContext(roleNames, state, info, null, n -> memory);

        assertThat(context.nodes(), equalTo(expectedNodes));
        if (hasDataRole) {
            assertThat(context.currentCapacity().node().storage(), equalTo(new ByteSizeValue(maxTotal)));
            assertThat(context.currentCapacity().total().storage(), equalTo(new ByteSizeValue(sumTotal)));
        } else {
            assertThat(context.currentCapacity().node().storage(), equalTo(ByteSizeValue.ZERO));
            assertThat(context.currentCapacity().total().storage(), equalTo(ByteSizeValue.ZERO));
        }
        assertThat(context.currentCapacity().node().memory(), equalTo(new ByteSizeValue(memory * Integer.signum(expectedNodes.size()))));
        assertThat(context.currentCapacity().total().memory(), equalTo(new ByteSizeValue(memory * expectedNodes.size())));

        if (expectedNodes.isEmpty() == false) {
            context = new AutoscalingCalculateCapacityService.DefaultAutoscalingDeciderContext(
                roleNames,
                state,
                info,
                null,
                AutoscalingMemoryInfo.EMPTY
            );
            assertThat(context.nodes(), equalTo(expectedNodes));
            assertThat(context.currentCapacity(), is(nullValue()));

            String multiPathNodeId = randomFrom(expectedNodes).getId();
            mostUsagesBuilder = ImmutableOpenMap.builder(mostUsages);
            DiskUsage original = mostUsagesBuilder.get(multiPathNodeId);
            mostUsagesBuilder.put(
                multiPathNodeId,
                new DiskUsage(
                    multiPathNodeId,
                    null,
                    randomValueOtherThan(original.getPath(), () -> randomAlphaOfLength(5)),
                    original.getTotalBytes(),
                    original.getFreeBytes()
                )
            );

            info = new ClusterInfo(leastUsages, mostUsagesBuilder.build(), null, null, null, null);
            context = new AutoscalingCalculateCapacityService.DefaultAutoscalingDeciderContext(roleNames, state, info, null, n -> memory);
            assertThat(context.nodes(), equalTo(expectedNodes));
            if (hasDataRole) {
                assertThat(context.currentCapacity(), is(nullValue()));
            } else {
                assertThat(context.currentCapacity().node().memory(), equalTo(new ByteSizeValue(memory)));
                assertThat(context.currentCapacity().total().memory(), equalTo(new ByteSizeValue(memory * expectedNodes.size())));
                assertThat(context.currentCapacity().node().storage(), equalTo(ByteSizeValue.ZERO));
                assertThat(context.currentCapacity().total().storage(), equalTo(ByteSizeValue.ZERO));
            }
        }
    }

    public void testValidateDeciderName() {
        AutoscalingCalculateCapacityService service = new AutoscalingCalculateCapacityService(
            org.elasticsearch.common.collect.Set.of(new FixedAutoscalingDeciderService())
        );
        String badDeciderName = randomValueOtherThan(FixedAutoscalingDeciderService.NAME, () -> randomAlphaOfLength(8));
        AutoscalingPolicy policy = new AutoscalingPolicy(
            randomAlphaOfLength(8),
            randomRoles(),
            new TreeMap<>(org.elasticsearch.common.collect.Map.of(badDeciderName, Settings.EMPTY))
        );
        IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> service.validate(policy));
        assertThat(exception.getMessage(), equalTo("unknown decider [" + badDeciderName + "]"));
    }

    public void testValidateDeciderRoles() {
        Set<String> roles = randomRoles();
        AutoscalingCalculateCapacityService service = new AutoscalingCalculateCapacityService(
            org.elasticsearch.common.collect.Set.of(new FixedAutoscalingDeciderService() {

                @Override
                public List<DiscoveryNodeRole> roles() {
                    return roles.stream().map(DiscoveryNode::getRoleFromRoleName).collect(Collectors.toList());
                }

                @Override
                public boolean appliesToEmptyRoles() {
                    return false;
                }

            })
        );
        SortedSet<String> badRoles = new TreeSet<>(randomRoles());
        badRoles.removeAll(roles);
        AutoscalingPolicy policy = new AutoscalingPolicy(
            FixedAutoscalingDeciderService.NAME,
            badRoles,
            new TreeMap<>(org.elasticsearch.common.collect.Map.of(FixedAutoscalingDeciderService.NAME, Settings.EMPTY))
        );
        IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> service.validate(policy));
        assertThat(
            exception.getMessage(),
            equalTo("decider [" + FixedAutoscalingDeciderService.NAME + "] not applicable to policy with roles [ " + badRoles + "]")
        );
    }

    public void testValidateNotEmptyDeciders() {
        AutoscalingCalculateCapacityService service = new AutoscalingCalculateCapacityService(
            org.elasticsearch.common.collect.Set.of(new FixedAutoscalingDeciderService())
        );
        String policyName = randomAlphaOfLength(8);
        AutoscalingPolicy policy = new AutoscalingPolicy(
            policyName,
            new TreeSet<>(randomBoolean() ? org.elasticsearch.common.collect.Set.of() : org.elasticsearch.common.collect.Set.of("master")),
            new TreeMap<>()
        );
        IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> service.validate(policy));
        assertThat(
            exception.getMessage(),
            equalTo("no default nor user configured deciders for policy [" + policyName + "] with roles [" + policy.roles() + "]")
        );
    }

    public void testValidateSettingName() {
        AutoscalingCalculateCapacityService service = new AutoscalingCalculateCapacityService(
            org.elasticsearch.common.collect.Set.of(new FixedAutoscalingDeciderService())
        );
        Set<String> legalNames = org.elasticsearch.common.collect.Set.of(
            FixedAutoscalingDeciderService.STORAGE.getKey(),
            FixedAutoscalingDeciderService.MEMORY.getKey(),
            FixedAutoscalingDeciderService.NODES.getKey()
        );
        String badSettingName = randomValueOtherThanMany(legalNames::contains, () -> randomAlphaOfLength(8));
        AutoscalingPolicy policy = new AutoscalingPolicy(
            randomAlphaOfLength(8),
            randomRoles(),
            new TreeMap<>(
                org.elasticsearch.common.collect.Map.of(
                    FixedAutoscalingDeciderService.NAME,
                    Settings.builder().put(badSettingName, randomAlphaOfLength(1)).build()
                )
            )
        );
        IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> service.validate(policy));
        assertThat(
            exception.getMessage(),
            equalTo("unknown setting [" + badSettingName + "] for decider [" + FixedAutoscalingDeciderService.NAME + "]")
        );
    }

    public void testValidateSettingValue() {
        AutoscalingCalculateCapacityService service = new AutoscalingCalculateCapacityService(
            org.elasticsearch.common.collect.Set.of(new FixedAutoscalingDeciderService())
        );
        String value = randomValueOtherThanMany(s -> Character.isDigit(s.charAt(0)), () -> randomAlphaOfLength(5));
        AutoscalingPolicy policy = new AutoscalingPolicy(
            randomAlphaOfLength(8),
            randomRoles(),
            new TreeMap<>(
                org.elasticsearch.common.collect.Map.of(
                    FixedAutoscalingDeciderService.NAME,
                    Settings.builder().put(FixedAutoscalingDeciderService.STORAGE.getKey(), value).build()
                )
            )
        );
        IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> service.validate(policy));
        assertThat(exception.getMessage(), containsString("[" + value + "]"));
        assertThat(exception.getMessage(), containsString("[" + FixedAutoscalingDeciderService.STORAGE.getKey() + "]"));
    }
}
