/**
 * A wrapper class of {@link TableFunction} for synchronously lookup rows matching the lookup keys
 * from external system.
 *
 * <p>The output type of this table function is fixed as {@link RowData}.
 */
@PublicEvolving
public abstract class LookupFunction extends TableFunction<RowData> {

    /**
     * Synchronously lookup rows matching the lookup keys.
     *
     * @param keys - Keys to lookup.
     * @return A collections of all matching rows in the lookup table.
     */
    public abstract Collection<RowData> lookup(Object... keys) throws IOException;

    /** Invoke {@link #lookup} and handle exceptions. */
    public final void eval(Object... keys) {
        try {
            lookup(keys).forEach(this::collect);
        } catch (IOException e) {
            throw new RuntimeException("Failed to lookup values with given key", e);
        }
    }
}

/**
 * A wrapper class of {@link AsyncTableFunction} for asynchronously lookup rows matching the lookup
 * keys from external system.
 *
 * <p>The output type of this table function is fixed as {@link RowData}.
 */
@PublicEvolving
public abstract class AsyncLookupFunction extends AsyncTableFunction<RowData> {

    /**
     * Asynchronously lookup rows matching the lookup keys.
     *
     * @param keys - Keys to lookup.
     * @return A collections of all matching rows in the lookup table.
     */
    public abstract CompletableFuture<Collection<RowData>> asyncLookup(Object... keys);

    /** Invokes {@link #asyncLookup} and chains futures. */
    public final void eval(CompletableFuture<Collection<RowData>> future, Object... keys) {
        asyncLookup(keys)
                .whenCompleteAsync(
                        (result, exception) -> {
                            if (exception != null) {
                                future.completeExceptionally(exception);
                                return;
                            }
                            future.complete(result);
                        });
    }
}

/**
 * A semi-persistent mapping from keys to values for storing entries of lookup table.
 *
 * <p>The type of the caching key is a {@link RowData} with lookup key fields packed inside. The
 * type of value is a {@link Collection} of {@link RowData}, which are rows matching lookup key
 * fields.
 *
 * <p>Cache entries are manually added using {@link #put}, and are stored in the cache until either
 * evicted or manually invalidated.
 *
 * <p>Implementations of this interface are expected to be thread-safe, and can be safely accessed
 * by multiple concurrent threads.
 */
@PublicEvolving
public interface LookupCache {

    /**
     * Returns the value associated with key in this cache, or null if there is no cached value for
     * key.
     */
    @Nullable
    Collection<RowData> getIfPresent(RowData key);

    /**
     * Associates the specified value rows with the specified key row in the cache. If the cache
     * previously contained value associated with the key, the old value is replaced by the
     * specified value.
      *
     * @return the previous value rows associated with key, or null if there was no mapping for key.
     * @param key - key row with which the specified value is to be associated
     * @param value – value rows to be associated with the specified key
     */
    Collection<RowData> put(RowData key, Collection<RowData> value);

    /**
     * Copies all the mappings from the specified map to the cache. The effect of this call is
     * equivalent to that of calling {@code put(k, v)} on this map once for each mapping from key
     * {@code k} to value {@code v} in the specified map. The behavior of this operation is
     * undefined if the specified map is modified while the operation is in progress.
     */
    void putAll(Map<? extends RowData, ? extends Collection<RowData>> m);

    /** Discards any cached value for the specified key. */
    void invalidate(RowData key);

    /** Discards all entries in the cache. */
    void invalidateAll();

    /** Returns the number of key-value mappings in the cache. */
    long size();
}

/** Factory for creating an instance of {@link LookupCache}. */
@PublicEvolving
public interface LookupCacheFactory extends Serializable {

    /**
     * Create a {@link LookupCache}.
     *
     * @param metricGroup - The lookup cache metric group in which the cache register predefined and
     *     custom metrics.
     */
    LookupCache createCache(LookupCacheMetricGroup metricGroup);
}

/** Factory for creating instance of {@link DefaultLookupCache}. */
@PublicEvolving
public class DefaultLookupCacheFactory implements LookupCacheFactory {
    private final Duration expireAfterAccessDuration;
    private final Duration expireAfterWriteDuration;
    private final Integer initialCapacity;
    private final Long maximumSize;

    public static DefaultLookupCacheFactory.Builder newBuilder() {
        return new Builder();
    }

    private DefaultLookupCacheFactory(
            Duration expireAfterAccessDuration,
            Duration expireAfterWriteDuration,
            Integer initialCapacity,
            Long maximumSize) {
        // Validation
        this.expireAfterAccessDuration = expireAfterAccessDuration;
        this.expireAfterWriteDuration = expireAfterWriteDuration;
        this.initialCapacity = initialCapacity;
        this.maximumSize = maximumSize;
    }

    @Override
    public LookupCache createCache(LookupCacheMetricGroup metricGroup) {
        // Create instance of DefaultLookupCache
    }

    /** Builder of {@link DefaultLookupCacheFactory}. */
    public static class Builder {
        private Duration expireAfterAccessDuration;
        private Duration expireAfterWriteDuration;
        private Integer initialCapacity;
        private Long maximumSize;

        public DefaultLookupCacheFactory.Builder expireAfterAccess(Duration duration) {
            expireAfterAccessDuration = duration;
            return this;
        }

        public DefaultLookupCacheFactory.Builder expireAfterWrite(Duration duration) {
            expireAfterWriteDuration = duration;
            return this;
        }

        public DefaultLookupCacheFactory.Builder initialCapacity(int initialCapacity) {
            this.initialCapacity = initialCapacity;
            return this;
        }

        public DefaultLookupCacheFactory.Builder maximumSize(long maximumSize) {
            this.maximumSize = maximumSize;
            return this;
        }

        public DefaultLookupCacheFactory build() {
            return new DefaultLookupCacheFactory(
                    expireAfterAccessDuration,
                    expireAfterWriteDuration,
                    initialCapacity,
                    maximumSize);
        }
    }
}

/** Pre-defined metrics for {@code LookupCache}. */
@PublicEvolving
public interface LookupCacheMetricGroup extends MetricGroup {
    /** The number of cache hits. */
    void setHitCounter(Counter hitCounter);

    /** The number of cache misses. */
    void setMissCounter(Counter missCounter);

    /** The number of times to load data into cache from external system. */
    void setLoadCounter(Counter loadCounter);

    /** The number of load failures. */
    void setNumLoadFailuresCounter(Counter numLoadFailuresCounter);

    /** The time spent for the latest load operation. */
    void setLatestLoadTimeGauge(Gauge<Long> latestLoadTimeGauge);

    /** The number of records in cache. */
    void setNumCachedRecordsGauge(Gauge<Long> numCachedRecordsGauge);

    /** The number of bytes used by cache. */
    void setNumCachedBytesGauge(Gauge<Long> numCachedBytesGauge);
}

/**
 * Provider for creating {@link LookupFunction} and {@link LookupCacheFactory} if caching should be
 * enabled for the lookup table.
 */
@PublicEvolving
public interface LookupFunctionProvider extends LookupTableSource.LookupRuntimeProvider {

    /**
     * Creates {@link LookupFunctionProvider} with the given {@link LookupFunction} and disable
     * lookup table caching.
     */
    static LookupFunctionProvider of(LookupFunction lookupFunction) {
        return new LookupFunctionProvider() {
            @Override
            public LookupFunction createLookupFunction() {
                return lookupFunction;
            }

            @Override
            public Optional<LookupCacheFactory> getCacheFactory() {
                return Optional.empty();
            }

            @Override
            public Optional<Boolean> cacheMissingKey() {
                return Optional.empty();
            }
        };
    }

    /**
     * Creates {@link LookupFunctionProvider} with the given {@link LookupFunction} and enable
     * caching with specified {@link LookupCacheFactory}.
     */
    static LookupFunctionProvider of(
            LookupFunction lookupFunction,
            LookupCacheFactory cacheFactory,
            boolean cacheMissingKey) {
        return new LookupFunctionProvider() {
            @Override
            public LookupFunction createLookupFunction() {
                return lookupFunction;
            }

            @Override
            public Optional<LookupCacheFactory> getCacheFactory() {
                return Optional.of(cacheFactory);
            }

            @Override
            public Optional<Boolean> cacheMissingKey() {
                return Optional.of(cacheMissingKey);
            }
        };
    }

    /** Creates an {@link LookupFunction} instance. */
    LookupFunction createLookupFunction();

    /**
     * Gets the {@link LookupCacheFactory} for creating lookup cache.
     *
     * <p>This factory will be used for creating an instance of cache during runtime execution for
     * optimizing the access to external lookup table.
     *
     * @return an {@link Optional} of {@link LookupCacheFactory}, or an empty {@link Optional} if
     *     caching shouldn't be applies to the lookup table.
     */
    Optional<LookupCacheFactory> getCacheFactory();

    /**
     * Whether the missing key (key fields without any matching value rows) should be stored in the
     * cache.
     *
     * <p>Please note that this option is required if {@link #getCacheFactory()} returns a non-empty
     * instance. If the cache factory is empty, the return value of this function will be ignored.
     *
     * @return true if a null or empty value should be stored in the cache.
     */
    Optional<Boolean> cacheMissingKey();
}

@PublicEvolving
public interface AsyncLookupFunctionProvider extends LookupTableSource.LookupRuntimeProvider {

    /**
     * Creates {@link AsyncLookupFunctionProvider} with the given {@link AsyncLookupFunction} and
     * disable lookup table caching.
     */
    static AsyncLookupFunctionProvider of(AsyncLookupFunction asyncLookupFunction) {
        return new AsyncLookupFunctionProvider() {
            @Override
            public AsyncLookupFunction createAsyncLookupFunction() {
                return asyncLookupFunction;
            }

            @Override
            public Optional<LookupCacheFactory> getCacheFactory() {
                return Optional.empty();
            }

            @Override
            public Optional<Boolean> cacheMissingKey() {
                return Optional.empty();
            }
        };
    }

    /**
     * Creates {@link AsyncLookupFunctionProvider} with the given {@link AsyncLookupFunction} and
     * enable caching with specified {@link LookupCacheFactory}.
     */
    static AsyncLookupFunctionProvider of(
            AsyncLookupFunction asyncLookupFunction,
            LookupCacheFactory cacheBuilder,
            boolean cacheMissingKey) {
        return new AsyncLookupFunctionProvider() {
            @Override
            public AsyncLookupFunction createAsyncLookupFunction() {
                return asyncLookupFunction;
            }

            @Override
            public Optional<LookupCacheFactory> getCacheFactory() {
                return Optional.of(cacheBuilder);
            }

            @Override
            public Optional<Boolean> cacheMissingKey() {
                return Optional.of(cacheMissingKey);
            }
        };
    }

    /** Creates an {@link AsyncLookupFunction} instance. */
    AsyncLookupFunction createAsyncLookupFunction();

    /**
     * Gets the {@link LookupCacheFactory} for creating lookup cache.
     *
     * <p>This factory will be used for creating an instance of cache during runtime execution for
     * optimizing the access to external lookup table.
     *
     * @return an {@link Optional} of {@link LookupCacheFactory}, or an empty {@link Optional} if
     *     caching shouldn't be applies to the lookup table.
     */
    Optional<LookupCacheFactory> getCacheFactory();

    /**
     * Whether the missing key (key fields without any matching value rows) should be stored in the
     * cache.
     *
     * <p>Please note that this option is required if {@link #getCacheFactory()} returns a non-empty
     * instance. If the cache factory is empty, the return value of this function will be ignored.
     *
     * @return true if a null or empty value should be stored in the cache.
     */
    Optional<Boolean> cacheMissingKey();
}

/**
 * Runtime provider for periodically re-scanning all entries of the lookup table and storing the
 * table locally for lookup.
 */
@PublicEvolving
public interface RescanRuntimeProvider extends LookupTableSource.LookupRuntimeProvider {

    /**
     * Gets the {@link org.apache.flink.table.connector.source.ScanTableSource.ScanRuntimeProvider}
     * for executing the periodically re-scan.
     */
    ScanTableSource.ScanRuntimeProvider getScanRuntimeProvider();

    /** Gets the interval between two re-scans. */
    Duration getRescanInterval();
}