/*
 * 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.kafka.controller;

import org.apache.kafka.common.config.ConfigDef;
import org.apache.kafka.common.errors.InvalidRequestException;
import org.apache.kafka.common.metadata.ClientQuotaRecord;
import org.apache.kafka.common.metadata.ClientQuotaRecord.EntityData;
import org.apache.kafka.common.protocol.Errors;
import org.apache.kafka.common.quota.ClientQuotaAlteration;
import org.apache.kafka.common.quota.ClientQuotaEntity;
import org.apache.kafka.common.requests.ApiError;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.server.common.ApiMessageAndVersion;
import org.apache.kafka.server.config.QuotaConfig;
import org.apache.kafka.server.mutable.BoundedList;
import org.apache.kafka.timeline.SnapshotRegistry;
import org.apache.kafka.timeline.TimelineHashMap;

import org.slf4j.Logger;

import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.function.Supplier;
import java.util.stream.Collectors;

import static org.apache.kafka.controller.QuorumController.MAX_RECORDS_PER_USER_OP;


public class ClientQuotaControlManager {
    static class Builder {
        private LogContext logContext = null;
        private SnapshotRegistry snapshotRegistry = null;

        Builder setLogContext(LogContext logContext) {
            this.logContext = logContext;
            return this;
        }

        Builder setSnapshotRegistry(SnapshotRegistry snapshotRegistry) {
            this.snapshotRegistry = snapshotRegistry;
            return this;
        }

        ClientQuotaControlManager build() {
            if (logContext == null) logContext = new LogContext();
            if (snapshotRegistry == null) snapshotRegistry = new SnapshotRegistry(logContext);
            return new ClientQuotaControlManager(logContext, snapshotRegistry);
        }
    }

    private final Logger log;

    private final SnapshotRegistry snapshotRegistry;

    final TimelineHashMap<ClientQuotaEntity, TimelineHashMap<String, Double>> clientQuotaData;

    ClientQuotaControlManager(
        LogContext logContext,
        SnapshotRegistry snapshotRegistry
    ) {
        this.log = logContext.logger(ClientQuotaControlManager.class);
        this.snapshotRegistry = snapshotRegistry;
        this.clientQuotaData = new TimelineHashMap<>(snapshotRegistry, 0);
    }

    /**
     * Determine the result of applying a batch of client quota alteration.  Note
     * that this method does not change the contents of memory.  It just generates a
     * result, that you can replay later if you wish using replay().
     *
     * @param quotaAlterations  List of client quota alterations to evaluate
     * @return                  The result.
     */
    ControllerResult<Map<ClientQuotaEntity, ApiError>> alterClientQuotas(
            Collection<ClientQuotaAlteration> quotaAlterations) {
        List<ApiMessageAndVersion> outputRecords =
                BoundedList.newArrayBacked(MAX_RECORDS_PER_USER_OP);
        Map<ClientQuotaEntity, ApiError> outputResults = new HashMap<>();

        quotaAlterations.forEach(quotaAlteration -> {
            // Note that the values in this map may be null
            Map<String, Double> alterations = new HashMap<>(quotaAlteration.ops().size());
            quotaAlteration.ops().forEach(op -> {
                if (alterations.containsKey(op.key())) {
                    outputResults.put(quotaAlteration.entity(), ApiError.fromThrowable(
                            new InvalidRequestException("Duplicate quota key " + op.key() +
                                " not updating quota for this entity " + quotaAlteration.entity())));
                } else {
                    alterations.put(op.key(), op.value());
                }
            });
            if (outputResults.containsKey(quotaAlteration.entity())) {
                outputResults.put(quotaAlteration.entity(), ApiError.fromThrowable(
                        new InvalidRequestException("Ignoring duplicate entity " + quotaAlteration.entity())));
            } else {
                alterClientQuotaEntity(quotaAlteration.entity(), alterations, outputRecords, outputResults);
            }
        });

        return ControllerResult.atomicOf(outputRecords, outputResults);
    }

    /**
     * Apply a quota record to the in-memory state.
     *
     * @param record    A ClientQuotaRecord instance.
     */
    public void replay(ClientQuotaRecord record) {
        Map<String, String> entityMap = new HashMap<>(2);
        record.entity().forEach(entityData -> entityMap.put(entityData.entityType(), entityData.entityName()));
        ClientQuotaEntity entity = new ClientQuotaEntity(entityMap);
        TimelineHashMap<String, Double> quotas = clientQuotaData.get(entity);
        if (quotas == null) {
            quotas = new TimelineHashMap<>(snapshotRegistry, 0);
            clientQuotaData.put(entity, quotas);
        }
        if (record.remove()) {
            quotas.remove(record.key());
            if (quotas.isEmpty()) {
                clientQuotaData.remove(entity);
            }
            log.info("Replayed ClientQuotaRecord for {} removing {}.", entity, record.key());
        } else {
            quotas.put(record.key(), record.value());
            log.info("Replayed ClientQuotaRecord for {} setting {} to {}.",
                    entity, record.key(), record.value());
        }
    }

    private void alterClientQuotaEntity(
        ClientQuotaEntity entity,
        Map<String, Double> newQuotaConfigs,
        List<ApiMessageAndVersion> outputRecords,
        Map<ClientQuotaEntity, ApiError> outputResults
    ) {
        // Check entity types and sanitize the names
        Map<String, String> validatedEntityMap = new HashMap<>(3);
        ApiError error = validateEntity(entity, validatedEntityMap);
        if (error.isFailure()) {
            outputResults.put(entity, error);
            return;
        }

        // Check the combination of entity types and get the config keys
        Map<String, ConfigDef.ConfigKey> configKeys = new HashMap<>(4);
        error = configKeysForEntityType(validatedEntityMap, configKeys);
        if (error.isFailure()) {
            outputResults.put(entity, error);
            return;
        }

        // Don't share objects between different records
        Supplier<List<EntityData>> recordEntitySupplier = () ->
                validatedEntityMap.entrySet().stream().map(mapEntry -> new EntityData()
                        .setEntityType(mapEntry.getKey())
                        .setEntityName(mapEntry.getValue()))
                        .collect(Collectors.toList());

        List<ApiMessageAndVersion> newRecords = new ArrayList<>(newQuotaConfigs.size());
        Map<String, Double> currentQuotas = clientQuotaData.containsKey(entity) ?
                clientQuotaData.get(entity) : Map.of();
        for (Map.Entry<String, Double> entry : newQuotaConfigs.entrySet()) {
            String key = entry.getKey();
            Double newValue = entry.getValue();
            if (newValue == null) {
                if (currentQuotas.containsKey(key)) {
                    // Null value indicates removal
                    newRecords.add(new ApiMessageAndVersion(new ClientQuotaRecord()
                            .setEntity(recordEntitySupplier.get())
                            .setKey(key)
                            .setRemove(true),
                        (short) 0));
                }
            } else {
                ApiError validationError = validateQuotaKeyValue(configKeys, key, newValue);
                if (validationError.isFailure()) {
                    outputResults.put(entity, validationError);
                    return;
                } else {
                    final Double currentValue = currentQuotas.get(key);
                    if (!Objects.equals(currentValue, newValue)) {
                        // Only record the new value if it has changed
                        newRecords.add(new ApiMessageAndVersion(new ClientQuotaRecord()
                                .setEntity(recordEntitySupplier.get())
                                .setKey(key)
                                .setValue(newValue),
                            (short) 0));
                    }
                }
            }
        }

        outputRecords.addAll(newRecords);
        outputResults.put(entity, ApiError.NONE);
    }

    static ApiError configKeysForEntityType(Map<String, String> entity, Map<String, ConfigDef.ConfigKey> output) {
        // We only allow certain combinations of quota entity types. Which type is in use determines which config
        // keys are valid
        boolean hasUser = entity.containsKey(ClientQuotaEntity.USER);
        boolean hasClientId = entity.containsKey(ClientQuotaEntity.CLIENT_ID);
        boolean hasIp = entity.containsKey(ClientQuotaEntity.IP);

        final Map<String, ConfigDef.ConfigKey> configKeys;
        if (hasIp) {
            if (hasUser || hasClientId) {
                return new ApiError(Errors.INVALID_REQUEST, "Invalid quota entity combination, IP entity should" +
                    "not be combined with User or ClientId");
            } else {
                if (isValidIpEntity(entity.get(ClientQuotaEntity.IP))) {
                    configKeys = QuotaConfig.ipConfigs().configKeys();
                } else {
                    return new ApiError(Errors.INVALID_REQUEST, entity.get(ClientQuotaEntity.IP) + " is not a valid IP or resolvable host.");
                }
            }
        } else if (hasUser || hasClientId) {
            configKeys = QuotaConfig.userAndClientQuotaConfigs().configKeys();
        } else {
            return new ApiError(Errors.INVALID_REQUEST, "Invalid empty client quota entity");
        }

        output.putAll(configKeys);
        return ApiError.NONE;
    }

    static ApiError validateQuotaKeyValue(
        Map<String, ConfigDef.ConfigKey> validKeys,
        String key,
        double value
    ) {
        // Ensure we have an allowed quota key
        ConfigDef.ConfigKey configKey = validKeys.get(key);
        if (configKey == null) {
            return new ApiError(Errors.INVALID_REQUEST, "Invalid configuration key " + key);
        }
        if (value <= 0.0) {
            return new ApiError(Errors.INVALID_REQUEST, "Quota " + key + " must be greater than 0");
        }

        // Ensure the quota value is valid
        return switch (configKey.type()) {
            case DOUBLE -> ApiError.NONE;
            case SHORT -> {
                if (value > Short.MAX_VALUE) {
                    yield new ApiError(Errors.INVALID_REQUEST,
                        "Proposed value for " + key + " is too large for a SHORT.");
                }
                yield getErrorForIntegralQuotaValue(value, key);
            }
            case INT -> {
                if (value > Integer.MAX_VALUE) {
                    yield new ApiError(Errors.INVALID_REQUEST,
                        "Proposed value for " + key + " is too large for an INT.");
                }
                yield getErrorForIntegralQuotaValue(value, key);
            }
            case LONG -> {
                if (value > Long.MAX_VALUE) {
                    yield new ApiError(Errors.INVALID_REQUEST,
                        "Proposed value for " + key + " is too large for a LONG.");
                }
                yield getErrorForIntegralQuotaValue(value, key);
            }
            default -> new ApiError(Errors.UNKNOWN_SERVER_ERROR,
                    "Unexpected config type " + configKey.type() + " should be Long or Double");
        };
    }

    static ApiError getErrorForIntegralQuotaValue(double value, String key) {
        double remainder = Math.abs(value % 1.0);
        if (remainder > 1e-6) {
            return new ApiError(Errors.INVALID_REQUEST, key + " cannot be a fractional value.");
        }
        return ApiError.NONE;
    }

    static boolean isValidIpEntity(String ip) {
        if (ip == null) return true;
        try {
            InetAddress.getByName(ip);
            return true;
        } catch (UnknownHostException e) {
            return false;
        }
    }

    private ApiError validateEntity(ClientQuotaEntity entity, Map<String, String> validatedEntityMap) {
        // Given a quota entity (which is a mapping of entity type to entity name), validate its types
        if (entity.entries().isEmpty()) {
            return new ApiError(Errors.INVALID_REQUEST, "Invalid empty client quota entity");
        }

        for (Entry<String, String> entityEntry : entity.entries().entrySet()) {
            String entityType = entityEntry.getKey();
            String entityName = entityEntry.getValue();
            if (validatedEntityMap.containsKey(entityType)) {
                return new ApiError(Errors.INVALID_REQUEST, "Invalid client quota entity, duplicate entity entry " + entityType);
            }
            if (Objects.equals(entityType, ClientQuotaEntity.USER)) {
                validatedEntityMap.put(ClientQuotaEntity.USER, entityName);
            } else if (Objects.equals(entityType, ClientQuotaEntity.CLIENT_ID)) {
                validatedEntityMap.put(ClientQuotaEntity.CLIENT_ID, entityName);
            } else if (Objects.equals(entityType, ClientQuotaEntity.IP)) {
                validatedEntityMap.put(ClientQuotaEntity.IP, entityName);
            } else {
                return new ApiError(Errors.INVALID_REQUEST, "Unhandled client quota entity type: " + entityType);
            }

            if (entityName != null && entityName.isEmpty()) {
                return new ApiError(Errors.INVALID_REQUEST, "Empty " + entityType + " not supported");
            }
        }

        return ApiError.NONE;
    }
}
