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

import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import ru.yandex.bolts.collection.Tuple2;
import ru.yandex.qe.dispenser.domain.QuotaChangeRequest;
import ru.yandex.qe.dispenser.ws.quota.request.owning_cost.formula.ProviderOwningCostFormula;

import static ru.yandex.qe.dispenser.ws.quota.request.owning_cost.formula.ProviderOwningCostFormula.DEFAULT_FORMULA_KEY;
import static ru.yandex.qe.dispenser.ws.quota.request.owning_cost.formula.ProviderOwningCostFormula.DEFAULT_OWNING_COST;

/**
 * Manager for calculating owning cost of requested resources.
 *
 * @author Ruslan Kadriev <aqru@yandex-team.ru>
 */
@Component
public class QuotaChangeOwningCostManager {
    private static final Logger LOG = LoggerFactory.getLogger(QuotaChangeOwningCostManager.class);

    /**
     * Set of campaign keys, for which we need to calculate the cost of owning.
     */
    private final Set<Long> campaignIds;
    private final Map<String, ProviderOwningCostFormula> formulaByProviderKey;
    /**
     * Default provider owning cost formula, used if no provider formula found.
     */
    private final ProviderOwningCostFormula defaultFormula;
    private final QuotaChangeOwningCostPreCalculationManager quotaChangeOwningCostPreCalculationManager;

    public QuotaChangeOwningCostManager(@Value("#{${quota.request.owner.cost.campaign.ids}}") List<Long> campaignIdSet,
                                        List<ProviderOwningCostFormula> providerFormulas,
                                        QuotaChangeOwningCostPreCalculationManager quotaChangeOwningCostPreCalculationManager) {
        this.quotaChangeOwningCostPreCalculationManager = quotaChangeOwningCostPreCalculationManager;
        this.campaignIds = ImmutableSet.copyOf(campaignIdSet);
        this.defaultFormula = providerFormulas.stream()
                .filter(providerOwningCostFormula -> providerOwningCostFormula.getProviderKey()
                        .equals(DEFAULT_FORMULA_KEY))
                .findFirst()
                .orElseThrow();
        this.formulaByProviderKey = ImmutableMap.copyOf(providerFormulas.stream()
                .filter(providerOwningCostFormula -> !providerOwningCostFormula.getProviderKey()
                        .equals(DEFAULT_FORMULA_KEY))
                .collect(Collectors.toMap(ProviderOwningCostFormula::getProviderKey, Function.identity())));
    }

    /**
     * Get campaign ids, for which we can calculate owning cost.
     * @return campaign key set
     */
    @NotNull
    public Set<Long> getCampaignIds() {
        return campaignIds;
    }

    /**
     * Calculate owning cost of requested resources.
     * If no new owning cost, old will be used.
     *
     * @param quotaChangeRequestOwningCostContext with changes from one request
     * @return updated changes with calculated owning cost if campaign is valid, unmodified changes otherwise
     */
    @NotNull
    public List<QuotaChangeRequest.Change> getChangesWithCalculatedOwningCost(
            @NotNull QuotaChangeRequestOwningCostContext quotaChangeRequestOwningCostContext) {
        List<ChangeOwningCostContext> changes =
                quotaChangeRequestOwningCostContext.getChangeOwningCostContexts();
        List<Long> campaignIds = changes.stream()
                .map(changeOwningCostContext -> changeOwningCostContext.getCampaign().getId())
                .collect(Collectors.toList());
        if (!this.campaignIds.containsAll(campaignIds)) {
            return changes.stream()
                    .map(ChangeOwningCostContext::getChange)
                    .collect(Collectors.toList());
        }

        Tuple2<List<ChangeOwningCostContext>, QuotaChangeOwningCostPreCalculationManager.State> preparedChanges =
                prepareChangesForCalculation(quotaChangeRequestOwningCostContext);

        Map<QuotaChangeRequest.ChangeKey, BigDecimal> owningCostByKey;
        if (preparedChanges.get2() == QuotaChangeOwningCostPreCalculationManager.State.FAIL) {
            owningCostByKey = onFail(quotaChangeRequestOwningCostContext);
        } else {
            owningCostByKey = calculateOwningCost(preparedChanges.get1());
        }

        return changes.stream()
                .map(change -> owningCostByKey.containsKey(change.getChange().getKey())
                        ? change.getChange().copyBuilder().owningCost(owningCostByKey.get(change.getChange().getKey())).build()
                        : getOwningCostOnMode(change.getChange(), quotaChangeRequestOwningCostContext.getMode()))
                .collect(Collectors.toList());
    }

    /**
     * Calculate owning cost of requested resources for update request.
     * On calculation fail will return changes with {@link ProviderOwningCostFormula#DEFAULT_OWNING_COST}
     * for entire changes in request.
     *
     * @param quotaChangeRequestOwningCostContext with updated changes and old changes from one request
     * @return updated changes with calculated owning cost if campaign is valid
     * @throws IllegalStateException on invalid {@link QuotaChangeRequestOwningCostContext.Mode}, allowed only
     * {@link QuotaChangeRequestOwningCostContext.Mode#UPDATE}.
     */
    @NotNull
    public ChangesWithOwningCostForUpdate getChangesWithCalculatedOwningCostForUpdate(
            @NotNull QuotaChangeRequestOwningCostContext quotaChangeRequestOwningCostContext) {
        if (quotaChangeRequestOwningCostContext.getMode() != QuotaChangeRequestOwningCostContext.Mode.UPDATE) {
            throw new IllegalStateException("'getChangesWithCalculatedOwningCostForUpdate' method only for update requests!");
        }

        List<ChangeOwningCostContext> changes =
                quotaChangeRequestOwningCostContext.getChangeOwningCostContexts();
        List<ChangeOwningCostContext> originalChanges =
                quotaChangeRequestOwningCostContext.getOriginalChangeOwningCostContexts()
                        .orElseThrow();

        List<Long> campaignIds = Stream.concat(changes.stream(), originalChanges.stream())
                .map(changeOwningCostContext -> changeOwningCostContext.getCampaign().getId())
                .collect(Collectors.toList());
        if (!this.campaignIds.containsAll(campaignIds)) {
            List<QuotaChangeRequest.Change> returnUnchanged = changes.stream()
                    .map(ChangeOwningCostContext::getChange)
                    .collect(Collectors.toList());
            return ChangesWithOwningCostForUpdate.builder()
                    .changesWithCalculatedOwningCost(returnUnchanged)
                    .changesWithCalculatedOwningCostForUpdateOwningCost(returnUnchanged)
                    .build();
        }

        Tuple2<List<ChangeOwningCostContext>, QuotaChangeOwningCostPreCalculationManager.State> preparedChanges =
                prepareChangesForCalculation(quotaChangeRequestOwningCostContext);

        Map<QuotaChangeRequest.ChangeKey, BigDecimal> owningCostByKey;
        if (preparedChanges.get2() == QuotaChangeOwningCostPreCalculationManager.State.FAIL) {
            owningCostByKey = onFail(quotaChangeRequestOwningCostContext);
        } else {
            List<ChangeOwningCostContext> preparedChanges1 = preparedChanges.get1();
            Map<QuotaChangeRequest.ChangeKey, ChangeOwningCostContext> entireChangesFromPrepared = Stream.concat(originalChanges.stream(),
                            preparedChanges1.stream())
                    .collect(Collectors.toMap(k -> k.getChange().getKey(), Function.identity(), (a,b) -> b));

            owningCostByKey = calculateOwningCost(entireChangesFromPrepared.values());
        }

        Map<QuotaChangeRequest.ChangeKey, ChangeOwningCostContext> entireChanges = Stream.concat(originalChanges.stream(),
                        changes.stream())
                .collect(Collectors.toMap(k -> k.getChange().getKey(), Function.identity(), (a, b) -> b));

        return ChangesWithOwningCostForUpdate.builder()
                .changesWithCalculatedOwningCost(changes.stream()
                        .map(change -> owningCostByKey.containsKey(change.getChange().getKey())
                                ? change.getChange().copyBuilder().owningCost(owningCostByKey.get(change.getChange().getKey())).build()
                                : getOwningCostOnMode(change.getChange(), QuotaChangeRequestOwningCostContext.Mode.UPDATE))
                        .collect(Collectors.toList()))
                .changesWithCalculatedOwningCostForUpdateOwningCost(
                        entireChanges.values().stream()
                                .map(change -> owningCostByKey.containsKey(change.getChange().getKey())
                                        ? change.getChange().copyBuilder().owningCost(owningCostByKey.get(change.getChange().getKey())).build()
                                        : getOwningCostOnMode(change.getChange(), QuotaChangeRequestOwningCostContext.Mode.UPDATE))
                                .collect(Collectors.toList())
                )
                .build();
    }

    /**
     * Calculate owning cost of requested resources from requests.
     * If no new owning cost, old will be used.
     *
     * @param requests with changes for calculation.
     * @param mode of current request.
     * @return requests with updated changes with calculated owning cost if campaign is valid or unmodified otherwise.
     */
    @NotNull
    public List<QuotaChangeRequest> getRequestsWithCalculatedOwningCost(
            @NotNull List<QuotaChangeRequest> requests, @NotNull QuotaChangeRequestOwningCostContext.Mode mode) {
        List<QuotaChangeRequest> validRequests = requests.stream()
                .filter(request -> campaignIds.contains(request.getCampaignId()))
                .collect(Collectors.toList());

        if (validRequests.isEmpty()) {
            return requests;
        }

        Map<Long, QuotaChangeRequest> requestByChangeId = new HashMap<>();
        Map<Long, List<ChangeOwningCostContext>> validChangesByRequestId = new HashMap<>();
        List<ChangeOwningCostContext> validChanges = new ArrayList<>();

        validRequests.forEach(request -> request.getChanges().forEach(change -> {
            ChangeOwningCostContext changeOwningCostContext = getChangeOwningCostContext(request, change);
            validChanges.add(changeOwningCostContext);
            validChangesByRequestId.computeIfAbsent(request.getId(), k -> new ArrayList<>()).add(changeOwningCostContext);
            requestByChangeId.put(change.getId(), request);
        }));

        QuotaChangeRequestOwningCostContext quotaChangeRequestOwningCostContext =
                getQuotaChangeRequestOwningCostContext(mode, validChanges);

        Tuple2<List<ChangeOwningCostContext>, QuotaChangeOwningCostPreCalculationManager.State> preparedChanges =
                prepareChangesForCalculation(quotaChangeRequestOwningCostContext);

        if (preparedChanges.get2() == QuotaChangeOwningCostPreCalculationManager.State.FAIL) {
            Map<Long, Map<QuotaChangeRequest.ChangeKey, BigDecimal>> owningCostByChangeKeyByRequest = new HashMap<>();
            validRequests.forEach(request ->
                    owningCostByChangeKeyByRequest.put(request.getId(),
                            onFail(getQuotaChangeRequestOwningCostContext(mode, validChangesByRequestId.get(request.getId())))));

            return requests.stream()
                    .map(request -> request.copyBuilder().changes(
                            request.getChanges().stream()
                                    .map(change -> {
                                        Map<QuotaChangeRequest.ChangeKey, BigDecimal> owningCostByChangeKey =
                                                owningCostByChangeKeyByRequest.getOrDefault(request.getId(), Map.of());
                                        return owningCostByChangeKey.containsKey(change.getKey())
                                                ? change.copyBuilder()
                                                .owningCost(owningCostByChangeKey.get(change.getKey()))
                                                .build()
                                                : getOwningCostOnMode(change, quotaChangeRequestOwningCostContext.getMode());
                                    })
                                    .collect(Collectors.toList())
                    ).build())
                    .collect(Collectors.toList());
        }

        List<ChangeOwningCostContext> preparedValidChanges = preparedChanges.get1();

        Map<QuotaChangeRequest, List<ChangeOwningCostContext>> changesByRequest = preparedValidChanges.stream()
                .filter(c -> requestByChangeId.containsKey(c.getChange().getId()))
                .collect(Collectors.groupingBy(c -> requestByChangeId.get(c.getChange().getId())));

        Map<Long, List<QuotaChangeRequest.Change>> updatedRequestChangesByRequestId = changesByRequest.entrySet().stream()
                .map(e -> {
                    Map<QuotaChangeRequest.ChangeKey, BigDecimal> owningCostByKey = calculateOwningCost(e.getValue());

                    List<QuotaChangeRequest.Change> updatedRequestChanges = e.getKey().getChanges().stream()
                            .map(change -> owningCostByKey.containsKey(change.getKey())
                                    ? change.copyBuilder().owningCost(owningCostByKey.get(change.getKey())).build()
                                    : getOwningCostOnMode(change, mode))
                            .collect(Collectors.toList());
                    return Tuple2.tuple(e.getKey().getId(), updatedRequestChanges);
                })
                .collect(Collectors.toMap(Tuple2::get1, Tuple2::get2));

        return requests.stream()
                .map(request -> updatedRequestChangesByRequestId.containsKey(request.getId())
                        ? request.copyBuilder().changes(updatedRequestChangesByRequestId.get(request.getId())).build()
                        : request)
                .collect(Collectors.toList());
    }

    private ChangeOwningCostContext getChangeOwningCostContext(QuotaChangeRequest request, QuotaChangeRequest.Change change) {
        return ChangeOwningCostContext.builder()
                .change(change)
                .abcServiceIdFromRequest(request)
                .campaignFromRequest(request)
                .build();
    }

    private QuotaChangeRequestOwningCostContext getQuotaChangeRequestOwningCostContext(
            QuotaChangeRequestOwningCostContext.Mode mode,
            List<ChangeOwningCostContext> changes) {
        return QuotaChangeRequestOwningCostContext.builder()
                .changeOwningCostContexts(changes)
                .mode(mode)
                .build();
    }

    private QuotaChangeRequest.Change getOwningCostOnMode(QuotaChangeRequest.Change change,
                                                          QuotaChangeRequestOwningCostContext.Mode mode) {
        QuotaChangeRequest.Change updatedChange = change;
        switch (mode) {
            case REFRESH:
                break;
            case UPDATE:
            case CREATE:
                updatedChange = updatedChange.copyBuilder().owningCost(DEFAULT_OWNING_COST).build();
                break;
        }
        return updatedChange;
    }

    private Map<QuotaChangeRequest.ChangeKey, BigDecimal> onFail(QuotaChangeRequestOwningCostContext quotaChangeRequestOwningCostContext) {
        QuotaChangeRequestOwningCostContext.Mode mode = quotaChangeRequestOwningCostContext.getMode();
        switch (mode) {
            case REFRESH:
                // use old values on refresh
                return toMap(quotaChangeRequestOwningCostContext.getChangeOwningCostContexts());
            case UPDATE:
            case CREATE:
                // in modify modes reset owning cost to default
                return toMapWithDuplicates(Stream.concat(quotaChangeRequestOwningCostContext.getOriginalChangeOwningCostContexts()
                                        .orElse(List.of())
                                        .stream(),
                                quotaChangeRequestOwningCostContext.getChangeOwningCostContexts().stream())
                        .map(changeOwningCostContexts -> changeOwningCostContexts.copyBuilder()
                                .change(changeOwningCostContexts
                                        .getChange()
                                        .copyBuilder()
                                        .owningCost(DEFAULT_OWNING_COST)
                                        .build())
                                .build()));
        }
        throw new IllegalStateException("Unsupported mode!");
    }

    private Map<QuotaChangeRequest.ChangeKey, BigDecimal> toMap(List<ChangeOwningCostContext> changes) {
        return changes.stream()
                .collect(Collectors.toMap(k -> k.getChange().getKey(), v -> v.getChange().getOwningCost()));
    }

    private Map<QuotaChangeRequest.ChangeKey, BigDecimal> toMapWithDuplicates(Stream<ChangeOwningCostContext> changes) {
        return changes
                .collect(Collectors.toMap(k -> k.getChange().getKey(), v -> v.getChange().getOwningCost(), (a,b) -> a));
    }

    private Tuple2<List<ChangeOwningCostContext>, QuotaChangeOwningCostPreCalculationManager.State> prepareChangesForCalculation(QuotaChangeRequestOwningCostContext quotaChangeRequestOwningCostContext) {
        return checkPreCalculationChain(
                quotaChangeOwningCostPreCalculationManager.apply(quotaChangeRequestOwningCostContext),
                quotaChangeRequestOwningCostContext
        );
    }

    private Map<QuotaChangeRequest.ChangeKey, BigDecimal> calculateOwningCost(Collection<ChangeOwningCostContext> changes) {
        return groupingByProvider(changes).entrySet().stream()
                .flatMap(e ->
                        checkCalculationResult(
                                formulaByProviderKey.getOrDefault(e.getKey(), defaultFormula)
                                        .calculateOwningCostFromContext(e.getValue())
                                , e.getValue())
                                .entrySet().stream())
                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
    }

    private Map<String, List<ChangeOwningCostContext>> groupingByProvider(Collection<ChangeOwningCostContext> changes) {
        return changes.stream()
                .collect(Collectors.groupingBy(c -> c.getChange().getResource().getService().getKey()));
    }

    private Map<QuotaChangeRequest.ChangeKey, BigDecimal> checkCalculationResult(
            Map<QuotaChangeRequest.ChangeKey, BigDecimal> calculateOwningCost,
            List<ChangeOwningCostContext> changes) {
        if (changes.size() != calculateOwningCost.size()) {
            changes.stream()
                    .filter(change -> !calculateOwningCost.containsKey(change.getChange().getKey()))
                    .forEach(this::logMissingOwningCost);
        }
        return calculateOwningCost;
    }

    private void logMissingOwningCost(ChangeOwningCostContext change) {
        LOG.warn("Missing calculating of owning cost for change: " + change);
    }

    private Tuple2<List<ChangeOwningCostContext>, QuotaChangeOwningCostPreCalculationManager.State> checkPreCalculationChain(
            Tuple2<List<ChangeOwningCostContext>, QuotaChangeOwningCostPreCalculationManager.State> apply,
            QuotaChangeRequestOwningCostContext quotaChangeRequestOwningCostContext) {
        if (apply.get1().size() != quotaChangeRequestOwningCostContext.getChangeOwningCostContexts().size()) {
            Map<QuotaChangeRequest.ChangeKey, ChangeOwningCostContext> owningCostContextMap = apply.get1().stream()
                    .collect(Collectors.toMap(k -> k.getChange().getKey(), Function.identity()));

            quotaChangeRequestOwningCostContext.getChangeOwningCostContexts().stream()
                    .filter(changeOwningCostContext -> !owningCostContextMap.containsKey(changeOwningCostContext.getChange().getKey()))
                    .forEach(this::logMissingChangeAfterPreCalculation);
        }
        return apply;
    }

    private void logMissingChangeAfterPreCalculation(ChangeOwningCostContext change) {
        LOG.warn("Missing change after pre calculation: " + change);
    }
}
