package ru.yandex.qe.dispenser.ws.quota.request.owning_cost.formula;

import java.math.BigDecimal;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.Function;
import java.util.stream.Collectors;

import com.google.common.collect.ImmutableMap;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import ru.yandex.qe.dispenser.api.util.EnumUtils;
import ru.yandex.qe.dispenser.api.v1.DiUnit;
import ru.yandex.qe.dispenser.domain.QuotaChangeRequest;
import ru.yandex.qe.dispenser.ws.bot.Provider;
import ru.yandex.qe.dispenser.ws.quota.request.owning_cost.ChangeOwningCostContext;
import ru.yandex.qe.dispenser.ws.quota.request.owning_cost.pricing.PricingModel;
import ru.yandex.qe.dispenser.ws.quota.request.owning_cost.pricing.QuotaChangeOwningCostTariffManager;

/**
 * YP owning cost formula (DISPENSER-4375).
 *
 * @author Ruslan Kadriev <aqru@yandex-team.ru>
 */
@Component
public class YpOwningCostFormula implements ProviderOwningCostFormula {
    private static final Logger LOG = LoggerFactory.getLogger(YpOwningCostFormula.class);

    private static final Provider provider = Provider.YP;

    private static final String CPU_SKU = "yp.cpu.quota";
    private static final String RAM_SKU = "yp.memory.quota";
    private static final String HDD_SKU = "yp.hdd_storage.quota";
    private static final String SSD_SKU = "yp.ssd_storage.quota";
    private static final String GPU_1_SKU = "yt.arnold.gpu.tesla_a100_80g.strong_guarantee.gpu";
    private static final String GPU_2_SKU = "yt.arnold.gpu.tesla_a100_80g.usage.gpu";

    private final QuotaChangeOwningCostTariffManager quotaChangeOwningCostTariffManager;

    public YpOwningCostFormula(QuotaChangeOwningCostTariffManager quotaChangeOwningCostTariffManager) {
        this.quotaChangeOwningCostTariffManager = quotaChangeOwningCostTariffManager;
    }

    @Override
    public @NotNull String getProviderKey() {
        return provider.getServiceKey();
    }

    @Override
    public @NotNull Map<QuotaChangeRequest.ChangeKey, BigDecimal> calculateOwningCostFromContext(
            @NotNull Collection<ChangeOwningCostContext> changes) {
        Map<QuotaChangeRequest.ChangeKey, BigDecimal> result = new HashMap<>();
        Map<CampaignKey, List<ChangeOwningCostContext>> changesByCampaign = ProviderOwningCostFormula
                .groupByCampaign(changes, result);
        changesByCampaign.forEach((campaignKey, campaignChanges) -> {
            Map<String, ? extends PricingModel> pricingBySKU = quotaChangeOwningCostTariffManager
                    .getByProvidersCampaign(List.of(provider, Provider.YT), campaignKey.getKey()).stream()
                    .collect(Collectors.toMap(PricingModel::getSKU, Function.identity()));
            Map<QuotaChangeRequest.ChangeKey, BigDecimal> campaignCosts = campaignChanges.stream()
                    .map(ChangeOwningCostContext::getChange)
                    .collect(Collectors.toMap(QuotaChangeRequest.ChangeAmount::getKey,
                            change -> toOwningCost(change, pricingBySKU)));
            result.putAll(campaignCosts);
        });
        return result;
    }

    private BigDecimal toOwningCost(QuotaChangeRequest.Change change,
                                    Map<String, ? extends PricingModel> pricingBySKU) {
        if (pricingBySKU.isEmpty()) {
            return DEFAULT_OWNING_COST;
        }

        BigDecimal result = DEFAULT_OWNING_COST;
        String publicKey = change.getResource().getPublicKey();
        Optional<Resource> resourcesO = Resource.byKey(publicKey);
        if (resourcesO.isPresent()) {
            switch (resourcesO.get()) {
                case SSD:
                    result = calculateForSsd(change, pricingBySKU);
                    break;
                case HDD:
                    result = calculateForHdd(change, pricingBySKU);
                    break;
                case CPU:
                    result = calculateForCpu(change, pricingBySKU);
                    break;
                case RAM:
                    result = calculateForRam(change, pricingBySKU);
                    break;
                case GPU:
                    result = calculateForGpu(change, pricingBySKU);
                    break;
                case IO_HDD:
                case IO_SSD:
                    break;
            }
        }

        return result;
    }

    private BigDecimal calculateForCpu(QuotaChangeRequest.Change change,
                                       Map<String, ? extends PricingModel> prisingModelsBySKU) {
        BigDecimal result = DEFAULT_OWNING_COST;

        if (prisingModelsBySKU.containsKey(CPU_SKU)) {
            PricingModel pricingModel = prisingModelsBySKU.get(CPU_SKU);
            result = convert(change, DiUnit.CORES, LOG)
                    .multiply(pricingModel.getPrice(), MATH_CONTEXT);
        }

        return result;
    }

    private BigDecimal calculateForRam(QuotaChangeRequest.Change change,
                                       Map<String, ? extends PricingModel> prisingModelsBySKU) {
        BigDecimal result = DEFAULT_OWNING_COST;

        if (prisingModelsBySKU.containsKey(RAM_SKU)) {
            PricingModel pricingModel = prisingModelsBySKU.get(RAM_SKU);
            result = convert(change, DiUnit.GIBIBYTE, LOG)
                    .multiply(pricingModel.getPrice(), MATH_CONTEXT);
        }

        return result;
    }

    private BigDecimal calculateForHdd(QuotaChangeRequest.Change change,
                                       Map<String, ? extends PricingModel> prisingModelsBySKU) {
        BigDecimal result = DEFAULT_OWNING_COST;

        if (prisingModelsBySKU.containsKey(HDD_SKU)) {
            PricingModel pricingModel = prisingModelsBySKU.get(HDD_SKU);
            result = convert(change, DiUnit.GIBIBYTE, LOG)
                    .multiply(pricingModel.getPrice(), MATH_CONTEXT);
        }

        return result;
    }

    private BigDecimal calculateForSsd(QuotaChangeRequest.Change change,
                                       Map<String, ? extends PricingModel> prisingModelsBySKU) {
        BigDecimal result = DEFAULT_OWNING_COST;

        if (prisingModelsBySKU.containsKey(SSD_SKU)) {
            PricingModel pricingModel = prisingModelsBySKU.get(SSD_SKU);
            result = convert(change, DiUnit.GIBIBYTE, LOG)
                    .multiply(pricingModel.getPrice(), MATH_CONTEXT);
        }

        return result;
    }

    private BigDecimal calculateForGpu(QuotaChangeRequest.Change change,
                                       Map<String, ? extends PricingModel> prisingModelsBySKU) {
        BigDecimal result = DEFAULT_OWNING_COST;

        if (prisingModelsBySKU.containsKey(GPU_1_SKU) && prisingModelsBySKU.containsKey(GPU_2_SKU)) {
            PricingModel pricingModel1 = prisingModelsBySKU.get(GPU_1_SKU);
            PricingModel pricingModel2 = prisingModelsBySKU.get(GPU_2_SKU);
            result = convert(change, DiUnit.COUNT, LOG)
                    .multiply(pricingModel1.getPrice().add(pricingModel2.getPrice(), MATH_CONTEXT), MATH_CONTEXT);
        }

        return result;
    }

    private enum Resource implements EnumUtils.StringKey {
        SSD("ssd_segmented"),
        HDD("hdd_segmented"),
        CPU("cpu_segmented"),
        RAM("ram_segmented"),
        GPU("gpu_segmented"),
        IO_HDD("io_hdd"),
        IO_SSD("io_ssd"),
        ;

        private static Map<String, Resource> resourceByKey;
        private final String key;

        Resource(String key) {
            this.key = key;
        }

        public String getKey() {
            return key;
        }

        public static Optional<Resource> byKey(String key) {
            if (Resource.resourceByKey == null) {
                Resource.resourceByKey = ImmutableMap.copyOf(EnumUtils.prepareKeysMap(Resource.values()));
            }

            return Optional.ofNullable(Resource.resourceByKey.get(key));
        }
    }
}
