/*
 * Copyright 2017-2019 EPAM Systems, Inc. (https://www.epam.com/)
 *
 * 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.
 */

package com.epam.pipeline.manager.cloud.azure;

import com.epam.pipeline.common.MessageConstants;
import com.epam.pipeline.common.MessageHelper;
import com.epam.pipeline.entity.cloud.CloudInstanceOperationResult;
import com.epam.pipeline.entity.cloud.azure.AzureVirtualMachineStats;
import com.epam.pipeline.entity.region.AzureRegion;
import com.epam.pipeline.exception.cloud.azure.AzureException;
import com.epam.pipeline.manager.cluster.KubernetesManager;
import com.epam.pipeline.manager.datastorage.providers.azure.AzureHelper;
import com.microsoft.azure.CloudException;
import com.microsoft.azure.Page;
import com.microsoft.azure.PagedList;
import com.microsoft.azure.management.Azure;
import com.microsoft.azure.management.compute.InstanceViewStatus;
import com.microsoft.azure.management.compute.VirtualMachine;
import com.microsoft.azure.management.compute.VirtualMachineDataDisk;
import com.microsoft.azure.management.compute.VirtualMachineInstanceView;
import com.microsoft.azure.management.compute.VirtualMachineScaleSet;
import com.microsoft.azure.management.compute.VirtualMachineScaleSetVM;
import com.microsoft.azure.management.network.NetworkInterface;
import com.microsoft.azure.management.resources.GenericResource;
import com.microsoft.azure.management.resources.fluentcore.arm.models.Resource;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Stream;

@RequiredArgsConstructor
@Slf4j
@Service
public class AzureVMService {

    private static final String VM_FAILED_STATE = "ProvisioningState/failed";
    private static final String TAG_NAME = "Name";
    private static final String VIRTUAL_MACHINE_PREFIX = "virtualmachine";
    private static final String VIRTUAL_MACHINES_TYPE = "virtualmachines";
    private static final String VIRTUAL_MACHINE_SCALE_SET_TYPE = "virtualmachinescalesets";
    private static final String NETWORK_INTERFACES = "networkinterfaces";
    private static final String RESOURCE_DELIMITER = "/";
    private static final String LOW_PRIORITY_INSTANCE_ID_TEMPLATE = "(az-[a-z0-9]{16})[0-9A-Z]{6}";
    private static final Pattern LOW_PRIORITY_VM_NAME_PATTERN = Pattern.compile(LOW_PRIORITY_INSTANCE_ID_TEMPLATE);
    private static final String SUCCEEDED = "Succeeded";
    private static final String KUBE_PREEMPTED_LABEL = "cloud-pipeline/preempted";
    private static final String PREEMPTED = "preempted";
    private static final String INSTANCE_WAS_PREEMPTED = "Low priority instance was preempted";
    private static final InstanceViewStatus SCALE_SET_FAILED_STATUS;

    static {
        SCALE_SET_FAILED_STATUS = new InstanceViewStatus();
        SCALE_SET_FAILED_STATUS.withCode(VM_FAILED_STATE + RESOURCE_DELIMITER + PREEMPTED);
        SCALE_SET_FAILED_STATUS.withMessage(INSTANCE_WAS_PREEMPTED);
    }

    private final MessageHelper messageHelper;

    private final KubernetesManager kubernetesManager;

    public CloudInstanceOperationResult startInstance(final AzureRegion region, final String instanceId) {
        getVmByName(region.getAuthFile(), region.getResourceGroup(), instanceId).start();
        return CloudInstanceOperationResult.success(
                messageHelper.getMessage(MessageConstants.INFO_INSTANCE_STARTED, instanceId)
        );
    }

    public void stopInstance(final AzureRegion region, final String instanceId) {
        getVmByName(region.getAuthFile(), region.getResourceGroup(), instanceId).powerOff();
    }

    public AzureVirtualMachineStats getRunningVMByRunId(final AzureRegion region, final String tagValue) {
        final AzureVirtualMachineStats virtualMachine = getVMStatsByTag(region, tagValue);
        if (!virtualMachine.hasRunningState()) {
            throw new AzureException(messageHelper.getMessage(
                    MessageConstants.ERROR_AZURE_INSTANCE_NOT_RUNNING, tagValue, virtualMachine.getPowerState()));
        }
        return virtualMachine;
    }

    public AzureVirtualMachineStats getAliveVMByRunId(final AzureRegion region, final String tagValue) {
        return getVMStatsByTag(region, tagValue);
    }

    public Optional<VirtualMachine> findVmByName(final AzureRegion region, final String instanceId) {
        return findVmByName(region.getAuthFile(), region.getResourceGroup(), instanceId);
    }

    public void terminateInstance(final AzureRegion region, final String instanceId) {
        final Azure azure = AzureHelper.buildClient(region.getAuthFile());
        final String instanceName = getInstanceResourceName(region, instanceId);
        resourcesByTag(region, instanceName)
                .sorted(this::resourcesTerminationOrder)
                .map(GenericResource::id)
                .forEach(resource -> azure.genericResources().deleteById(resource));
    }

    public NetworkInterface getVMNetworkInterface(final String authFile, final VirtualMachine vm) {
        final String interfaceId = vm.primaryNetworkInterfaceId();
        return AzureHelper.buildClient(authFile).networkInterfaces().getById(interfaceId);
    }

    public Optional<InstanceViewStatus> getFailingVMStatus(final AzureRegion region, final String vmName) {
        try {
            final Optional<String> scaleSetName = getScaleSetName(vmName);
            return scaleSetName.isPresent()
                    ? fetchFailingStatusFromScaleSet(region, scaleSetName.get(), vmName)
                    : fetchFailingStatusFromVm(region, vmName);
        } catch (CloudException e) {
            log.error(e.getMessage(), e);
            return Optional.empty();
        }
    }

    private String getInstanceResourceName(final AzureRegion region, final String instanceId) {
        return searchVmResource(region, instanceId)
                .map(Resource::tags)
                .map(Map::entrySet)
                .map(Set::stream)
                .orElseGet(Stream::empty)
                .filter(it -> it.getKey().equals(TAG_NAME))
                .map(Map.Entry::getValue)
                .findFirst()
                .orElseThrow(() -> new AzureException(messageHelper.getMessage(
                        MessageConstants.ERROR_AZURE_INSTANCE_NOT_FOUND, instanceId)));
    }

    public Optional<? extends Resource> searchVmResource(final AzureRegion region, final String instanceId) {
        final Optional<String> scaleSetName = getScaleSetName(instanceId);
        if(scaleSetName.isPresent()) {
            return findVmScaleSetByName(region, scaleSetName.get());
        } else {
            return findVmByName(region, instanceId);
        }
    }

    public void createAndAttachVolume(final String runId, final Long size, final AzureRegion region) {
        final Azure azure = AzureHelper.buildClient(region.getAuthFile());
        final PagedList<GenericResource> resources = azure.genericResources()
                .listByTag(region.getResourceGroup(), TAG_NAME, runId);
        createAndAttachAzureVolumeToVMContainer(
                azure,
                findVMContainerInPagedResult(resources.currentPage(), resources)
                        .orElseThrow(IllegalArgumentException::new),
                size
        );
    }

    private void createAndAttachAzureVolumeToVMContainer(final Azure azure,
                                                         final GenericResource vmc, final Long size) {
        final int sizeInGb = Math.toIntExact(size);
        if (vmc.resourceType().equals(VIRTUAL_MACHINE_SCALE_SET_TYPE)) {
            final VirtualMachineScaleSet scaleSet = azure.virtualMachineScaleSets().getById(vmc.id());
            final Optional<VirtualMachineDataDisk> dataDisk =
                    scaleSet.virtualMachines().list().stream().findFirst().flatMap(
                        vm -> vm.dataDisks().values().stream().findFirst()
                    );
            if (dataDisk.isPresent()) {
                final VirtualMachineDataDisk disk = dataDisk.get();
                scaleSet.update().withNewDataDisk(Math.toIntExact(sizeInGb),
                        disk.lun() + 1, disk.cachingType(), disk.storageAccountType()).apply();
            } else {
                scaleSet.update().withNewDataDisk(Math.toIntExact(sizeInGb)).apply();
            }
            scaleSet.virtualMachines().updateInstances(
                    scaleSet.virtualMachines().list().stream()
                            .map(VirtualMachineScaleSetVM::instanceId)
                            .toArray(String[]::new)
            );
        } else if (vmc.resourceType().equals(VIRTUAL_MACHINES_TYPE)) {
            final VirtualMachine virtualMachine = azure.virtualMachines().getById(vmc.id());
            final Optional<VirtualMachineDataDisk> dataDisk = virtualMachine.dataDisks().values().stream().findFirst();
            if (dataDisk.isPresent()) {
                final VirtualMachineDataDisk disk = dataDisk.get();
                virtualMachine.update().withNewDataDisk(Math.toIntExact(sizeInGb),
                       disk.lun() + 1, disk.cachingType(), disk.storageAccountType()).apply();
            } else {
                virtualMachine.update().withNewDataDisk(Math.toIntExact(sizeInGb)).apply();
            }
        } else {
            throw new AzureException(messageHelper.getMessage(
                    MessageConstants.ERROR_AZURE_RESOURCE_IS_NOT_VM_LIKE, vmc.id()));
        }
    }

    private Optional<String> getScaleSetName(final String instanceId) {
        final Matcher matcher = LOW_PRIORITY_VM_NAME_PATTERN.matcher(instanceId);
        return matcher.matches() ? Optional.ofNullable(matcher.group(1)) : Optional.empty();
    }

    private int resourcesTerminationOrder(final GenericResource r1, final GenericResource r2) {
        return resourceType(r1).startsWith(VIRTUAL_MACHINE_PREFIX)
                || resourceType(r1).equals(NETWORK_INTERFACES) && !resourceType(r2).startsWith(VIRTUAL_MACHINE_PREFIX)
                ? -1 : 0;
    }

    private String resourceType(final GenericResource resource) {
        return last(resource.resourceType().split(RESOURCE_DELIMITER)).toLowerCase();
    }

    private String last(final String[] items) {
        return items[items.length - 1];
    }

    private Optional<VirtualMachineScaleSet> findVmScaleSetByName(final AzureRegion region,
                                                                  final String scaleSetName) {
        return findVmScaleSetByName(region.getAuthFile(), region.getResourceGroup(), scaleSetName);
    }

    private Optional<InstanceViewStatus> fetchFailingStatusFromScaleSet(final AzureRegion region,
                                                                        final String scaleSetName,
                                                                        final String nodeName) {
        final Map<String, String> nodeLabels = kubernetesManager.findNodeByName(nodeName)
                .map(node -> node.getMetadata().getLabels())
                .orElse(Collections.emptyMap());
        if (nodeLabels.containsKey(KUBE_PREEMPTED_LABEL)) {
            return Optional.of(SCALE_SET_FAILED_STATUS);
        }
        final Optional<VirtualMachineScaleSet> scaleSet = findVmScaleSetByName(region, scaleSetName);
        if (scaleSet.isPresent() && scaleSet.get().inner().provisioningState().equals(SUCCEEDED)) {
            final PagedList<VirtualMachineScaleSetVM> scaleSetVMs = scaleSet.get().virtualMachines().list();
            return scaleSetVMs.size() > 0
                    ? findFailedStatus(scaleSetVMs.get(0).instanceView().statuses())
                    : Optional.of(SCALE_SET_FAILED_STATUS);
        }
        return Optional.empty();
    }

    private Optional<InstanceViewStatus> fetchFailingStatusFromVm(final AzureRegion region, final String vmName) {
        final Optional<VirtualMachine> virtualMachine = findVmByName(region, vmName);
        if (!virtualMachine.isPresent()) {
            return Optional.empty();
        }

        final List<InstanceViewStatus> statuses = virtualMachine
                .map(VirtualMachine::instanceView)
                .map(VirtualMachineInstanceView::statuses)
                .orElseGet(Collections::emptyList);
        if (CollectionUtils.isEmpty(statuses)) {
            log.debug("Virtual machine found, but status is not available");
            return Optional.empty();
        }

        return findFailedStatus(statuses);
    }

    private Optional<InstanceViewStatus> findFailedStatus(final List<InstanceViewStatus> statuses) {
        return statuses.stream()
                .filter(status -> status.code().startsWith(VM_FAILED_STATE))
                .findFirst();
    }

    private VirtualMachine getVmByName(final String authFile,
                                       final String resourceGroup,
                                       final String instanceId) {
        return findVmByName(authFile, resourceGroup, instanceId)
                .orElseThrow(() -> new AzureException(messageHelper.getMessage(
                        MessageConstants.ERROR_AZURE_INSTANCE_NOT_FOUND, instanceId)));
    }

    private Optional<VirtualMachine> findVmByName(final String authFile,
                                                  final String resourceGroup,
                                                  final String instanceId) {
        return Optional.of(AzureHelper.buildClient(authFile).virtualMachines())
                .map(client -> client.getByResourceGroup(resourceGroup, instanceId));
    }

    private Optional<VirtualMachineScaleSet> findVmScaleSetByName(final String authFile,
                                                  final String resourceGroup,
                                                  final String scaleSetName) {
        return Optional.of(AzureHelper.buildClient(authFile).virtualMachineScaleSets())
                .map(client -> client.getByResourceGroup(resourceGroup, scaleSetName));
    }

    public AzureVirtualMachineStats getVMStatsByTag(final AzureRegion region, final String tagValue) {
        final Azure azure = AzureHelper.buildClient(region.getAuthFile());
        final PagedList<GenericResource> resources = azure.genericResources()
                .listByTag(region.getResourceGroup(), TAG_NAME, tagValue);
        return findVMContainerInPagedResult(resources.currentPage(), resources)
                .map(vmc -> getVmStatsByVmContainer(azure, vmc))
                .orElseThrow(() -> new AzureException(messageHelper.getMessage(
                        MessageConstants.ERROR_AZURE_INSTANCE_NOT_FOUND, tagValue)));
    }

    private AzureVirtualMachineStats getVmStatsByVmContainer(final Azure azure, final GenericResource vmc) {
        if (resourceType(vmc).equals(VIRTUAL_MACHINE_SCALE_SET_TYPE)) {
            final VirtualMachineScaleSetVM scaleSetVM = azure.virtualMachineScaleSets().getById(vmc.id())
                    .virtualMachines().list().stream().findFirst()
                    .orElseThrow(() -> new AzureException(messageHelper.getMessage(
                            MessageConstants.ERROR_AZURE_SCALE_SET_DOESNT_CONTAIN_VMS, vmc.id())));
            return AzureVirtualMachineStats.fromScaleSetVirtualMachine(scaleSetVM);
        } else if (resourceType(vmc).equals(VIRTUAL_MACHINES_TYPE)){
            return AzureVirtualMachineStats.fromVirtualMachine(azure.virtualMachines().getById(vmc.id()));
        }
        throw new AzureException(messageHelper.getMessage(
                MessageConstants.ERROR_AZURE_RESOURCE_IS_NOT_VM_LIKE, vmc.id()));
    }

    private Optional<GenericResource> findVMContainerInPagedResult(final Page<GenericResource> currentPage,
                                                                   final PagedList<GenericResource> resources) {
        if (currentPage == null || CollectionUtils.isEmpty(currentPage.items())) {
            return Optional.empty();
        }
        final Optional<GenericResource> virtualMachineContainer = currentPage.items().stream()
                .filter(r -> resourceType(r).startsWith(VIRTUAL_MACHINE_PREFIX) ||
                        resourceType(r).equals(VIRTUAL_MACHINE_SCALE_SET_TYPE))
                .findFirst();
        return virtualMachineContainer.isPresent() ? virtualMachineContainer : checkNextPage(currentPage, resources);
    }

    private Optional<GenericResource> checkNextPage(final Page<GenericResource> currentPage,
                                                    final PagedList<GenericResource> resources) {
        if (resources.hasNextPage()) {
            try {
                return findVMContainerInPagedResult(resources.nextPage(currentPage.nextPageLink()), resources);
            } catch (IOException e) {
                log.error(e.getMessage(), e);
            }
        }
        return Optional.empty();
    }

    private Stream<GenericResource> resourcesByTag(final AzureRegion region, final String tagValue) {
        final Azure azure = AzureHelper.buildClient(region.getAuthFile());
        final PagedList<GenericResource> resources = azure.genericResources()
                .listByTag(region.getResourceGroup(), TAG_NAME, tagValue);
        return resourcesInPagedResult(resources.currentPage(), resources);
    }

    private Stream<GenericResource> resourcesInPagedResult(final Page<GenericResource> currentPage,
                                                           final PagedList<GenericResource> resources) {
        if (currentPage == null || CollectionUtils.isEmpty(currentPage.items())) {
            return Stream.empty();
        }
        return Stream.concat(currentPage.items().stream(), resourcesFromNextPage(currentPage, resources));
    }

    private Stream<GenericResource> resourcesFromNextPage(final Page<GenericResource> currentPage,
                                                          final PagedList<GenericResource> resources) {
        if (resources.hasNextPage()) {
            try {
                return resourcesInPagedResult(resources.nextPage(currentPage.nextPageLink()), resources);
            } catch (IOException e) {
                log.error(e.getMessage(), e);
            }
        }
        return Stream.empty();
    }
}
