/*
 * 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.addthis.hydra.job.auth;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;

import java.io.Closeable;
import java.io.IOException;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.attribute.PosixFilePermission;
import java.nio.file.attribute.PosixFilePermissions;

import com.addthis.codec.annotations.Time;

import com.google.common.collect.ImmutableSet;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class TokenCache implements Closeable {

    private static final ImmutableSet<PosixFilePermission> OWNER_READ_WRITE =
            ImmutableSet.of(PosixFilePermission.OWNER_READ, PosixFilePermission.OWNER_WRITE);

    public enum ExpirationPolicy {
        WRITE, ACCESS
    }

    private static final Logger log = LoggerFactory.getLogger(TokenCache.class);

    private static final ObjectMapper mapper = new ObjectMapper();

    /**
     * Expiration policy. Default is {@code WRITE}
     */
    @Nonnull
    public final ExpirationPolicy policy;

    /**
     * Expiration time in seconds.
     */
    public final int timeout;

    /**
     * Map each username to a map of tokens to expiration timestamps.
     */
    private final ConcurrentHashMap<String, ConcurrentHashMap<String, Long>> cache;

    @Nullable
    private final Path outputPath;

    @JsonCreator
    public TokenCache(@JsonProperty(value = "policy", required = true) ExpirationPolicy policy,
                      @JsonProperty(value = "timeout", required = true) @Time(TimeUnit.SECONDS) int timeout,
                      @JsonProperty("outputPath") Path outputPath) throws IOException {
        this.policy = policy;
        this.timeout = timeout;
        this.outputPath = outputPath;
        if ((outputPath != null) && (Files.isReadable(outputPath))) {
            log.info("Loading authentication tokens from disk.");
            TypeReference<ConcurrentHashMap<String, ConcurrentHashMap<String, Long>>> typeReference =
                    new TypeReference<ConcurrentHashMap<String, ConcurrentHashMap<String, Long>>>() {};
            this.cache = mapper.readValue(outputPath.toFile(), typeReference);
        } else {
            this.cache = new ConcurrentHashMap<>();
        }
    }

    private ConcurrentHashMap<String, Long>  buildCache() {
        return new ConcurrentHashMap<>();
    }

    public boolean get(@Nullable String name, @Nullable String secret) {
        long now = System.currentTimeMillis();
        if ((name == null) || (secret == null)) {
            return false;
        }
        ConcurrentHashMap<String, Long> userTokens = cache.computeIfAbsent(name, (k) -> buildCache());
        Long expiration = userTokens.compute(secret, (key, prev) -> {
            if (prev == null) {
                return null;
            } else if (prev < now) {
                return null;
            } else if (policy == ExpirationPolicy.ACCESS) {
                return (now + TimeUnit.SECONDS.toMillis(timeout));
            } else {
                return prev;
            }
        });
        return (expiration != null);
    }

    public void put(@Nonnull String name, @Nonnull String secret) {
        ConcurrentHashMap<String, Long> userTokens = cache.computeIfAbsent(name, (k) -> buildCache());
        userTokens.put(secret, System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(timeout));
    }

    public int remove(@Nonnull String name, @Nonnull String secret) {
        ConcurrentHashMap<String, Long> userTokens = cache.computeIfAbsent(name, (k) -> buildCache());
        userTokens.remove(secret);
        return userTokens.size();
    }

    public void evict(@Nonnull String name) {
        cache.remove(name);
    }

    @Override
    public void close() throws IOException {
        if (outputPath != null) {
            log.info("Persisting authentication tokens to disk.");
            if (!Files.exists(outputPath)) {
                Files.createFile(outputPath, PosixFilePermissions.asFileAttribute(OWNER_READ_WRITE));
            } else {
                Files.setPosixFilePermissions(outputPath, OWNER_READ_WRITE);
            }
            mapper.writeValue(outputPath.toFile(), cache);
        }
    }
}