/**
 * 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 keyRow - A {@link RowData} that wraps keys to lookup.
     * @return A collections of all matching rows in the lookup table.
     */
    public abstract Collection<RowData> lookup(RowData keyRow) throws IOException;

    /** Invoke {@link #lookup} and handle exceptions. */
    public final void eval(Object... keys) {
        try {
            lookup(GenericRowData.of(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 keyRow - A {@link RowData} that wraps keys to lookup.
     * @return A collections of all matching rows in the lookup table.
     */
    public abstract CompletableFuture<Collection<RowData>> asyncLookup(RowData keyRow);

    /** Invokes {@link #asyncLookup} and chains futures. */
    public final void eval(CompletableFuture<Collection<RowData>> future, Object... keys) {
         asyncLookup(GenericRowData.of(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 extends AutoClosable {

    /**
     * Initialize the cache.
     *
     * @param metricGroup the metric group to register cache related metrics.
     */
    void open(CacheMetricGroup metricGroup);

    /**
     * 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);

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

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

/** Default implementation of {@link LookupCache}. */
@PublicEvolving
public class DefaultLookupCache implements LookupCache {
    private final Duration expireAfterAccessDuration;
    private final Duration expireAfterWriteDuration;
    private final Long maximumSize;
    private final boolean cacheMissingKey;
    
    private DefaultLookupCache(
            Duration expireAfterAccessDuration,
            Duration expireAfterWriteDuration,
            Long maximumSize,
			boolean CacheMissingKey) {
        this.expireAfterAccessDuration = expireAfterAccessDuration;
        this.expireAfterWriteDuration = expireAfterWriteDuration;
        this.initialCapacity = initialCapacity;
        this.maximumSize = maximumSize;
		this.cacheMissingKey = cacheMissingKeyl
    }
    
    public static Builder newBuilder() {
        return new Builder();
    } 

   public static class Builder {         
        private Duration expireAfterAccessDuration;
        private Duration expireAfterWriteDuration;
        private Long maximumSize;
        private Boolean cacheMissingKey;

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

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

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

        public Builder cacheMissingKey(boolean cacheMissingKey) {
            this.cacheMissingKey = cacheMissingKey;
            return this;
        }          

        public DefaultLookupCache build() {
            return new DefaultLookupCache(
                    expireAfterAccessDuration,
                    expireAfterWriteDuration,
                    maximumSize,
					cacheMissingKey);
        }
    }     
}

/**
 * Pre-defined metrics for cache.
 *
 * <p>Please note that these methods should only be invoked once. Registering a metric with same
 * name for multiple times would lead to an undefined behavior.
 */
@PublicEvolving
public interface CacheMetricGroup extends MetricGroup {
    /** The number of cache hits. */
    void hitCounter(Counter hitCounter);

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

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

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

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

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

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

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

    /** Creates a builder of {@link PartialCachingLookupProvider}. */
    static Builder newBuilder() {
        return new Builder();
    }

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

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

    /** Builder class for {@link PartialCachingLookupProvider}. */
    class Builder {

        private LookupFunction lookupFunction;
        private LookupCache cache;
        private Boolean cacheMissingKey;

        /** Sets lookup function. */
        public Builder withLookupFunction(LookupFunction lookupFunction) {
            this.lookupFunction = lookupFunction;
            return this;
        }

        /** Enables caching and sets the cache factory. */
        public Builder withCache(LookupCache cache, boolean cacheMissingKey) {
            this.cache = cache;
            this.cacheMissingKey = cacheMissingKey;
            return this;
        }

        public PartialCachingLookupProvider build() {
           ...
        }
    }
}

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

    /** Creates a builder of {@link AsyncPartialCachingLookupProvider}. */
    static AsyncPartialCachingLookupProvider.Builder newBuilder() {
        return new AsyncPartialCachingLookupProvider.Builder();
    }

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

    /**
     * Gets the {@link LookupCache} 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 LookupCache}, or an empty {@link Optional} if caching
     *     shouldn't be applies to the lookup table.
     */
    Optional<LookupCache> getCache();

    /** Builder class for {@link AsyncPartialCachingLookupProvider}. */
    class Builder {

        private AsyncLookupFunction asyncLookupFunction;
        private LookupCache cache;
        private Boolean cacheMissingKey;

        /** Sets lookup function. */
        public AsyncPartialCachingLookupProvider.Builder withLookupFunction(
                AsyncLookupFunction asyncLookupFunction) {
            this.asyncLookupFunction = asyncLookupFunction;
            return this;
        }

        /** Enables caching and sets the cache factory. */
        public AsyncPartialCachingLookupProvider.Builder withCache(
                LookupCache cache, boolean cacheMissingKey) {
            this.cache = cache;
            this.cacheMissingKey = cacheMissingKey;
            return this;
        }

        public AsyncPartialCachingLookupProvider build() {
           ...
        }
    }
}

/**
 * Runtime provider for fully loading and periodically reloading all entries of the lookup table and
 * storing the table locally for lookup.
 *
 * <p>Implementations should provide a {@link ScanTableSource.ScanRuntimeProvider} in order to reuse
 * the ability of scanning for loading all entries from the lookup table.
 */
@PublicEvolving
public interface FullCachingLookupProvider extends LookupTableSource.LookupRuntimeProvider {

    static FullCachingLookupProvider of(
            ScanTableSource.ScanRuntimeProvider scanRuntimeProvider,
            FullCachingReloadTrigger fullCachingReloadTrigger) {
        return new FullCachingLookupProvider() {
            @Override
            public ScanTableSource.ScanRuntimeProvider getScanRuntimeProvider() {
                return scanRuntimeProvider;
            }

            @Override
            public FullCachingReloadTrigger getReloadTrigger() {
                return fullCachingReloadTrigger;
            }
        };
    }

    /**
     * Gets the {@link ScanTableSource.ScanRuntimeProvider} for executing the periodically reload.
     */
    ScanTableSource.ScanRuntimeProvider getScanRuntimeProvider();

    /** Get the {@link FullCachingReloadTrigger} for triggering a full caching reload operation. */
    FullCachingReloadTrigger getReloadTrigger();
}

/** Customized trigger for reloading all lookup table entries in full caching mode. */
public interface FullCachingReloadTrigger extends AutoCloseable, Serializable {

    /** Open the trigger. */
    void open(Context context) throws Exception;

    /**
     * Context of {@link FullCachingReloadTrigger} for getting information about times and
     * triggering reload.
     */
    interface Context {

        /** Get current processing time. */
        long currentProcessingTime();

        /** Get current watermark on the main stream. */
        long currentWatermark();

        /** Trigger a reload operation on the full cache. */
        CompletableFuture<Void> triggerReload();
    }
}

/** A trigger that reloads all entries periodically with specified interval or delay. */
public class PeriodicFullCachingReloadTrigger implements FullCachingReloadTrigger {

    private final Duration reloadInterval;
    private final ScheduleMode scheduleMode;

    private ScheduledExecutorService scheduledExecutor;

    public PeriodicFullCachingReloadTrigger(Duration reloadInterval, ScheduleMode scheduleMode) {
        this.reloadInterval = reloadInterval;
        this.scheduleMode = scheduleMode;
    }

    @Override
    public void open(FullCachingReloadTrigger.Context context) {
        scheduledExecutor = Executors.newSingleThreadScheduledExecutor();
        switch (scheduleMode) {
            case FIXED_RATE:
                scheduledExecutor.scheduleAtFixedRate(
                        context::triggerReload,
                        0,
                        reloadInterval.toMillis(),
                        TimeUnit.MILLISECONDS);
                break;
            case FIXED_DELAY:
                scheduledExecutor.scheduleWithFixedDelay(
                        () -> {
                            try {
                                context.triggerReload().get();
                            } catch (Exception e) {
                                throw new RuntimeException(
                                        "Uncaught exception during the reload", e);
                            }
                        },
                        0,
                        reloadInterval.toMillis(),
                        TimeUnit.MILLISECONDS);
                break;
            default:
                throw new IllegalArgumentException(
                        String.format("Unrecognized schedule mode \"%s\"", scheduleMode));
        }
    }

    @Override
    public void close() throws Exception {
        if (scheduledExecutor != null) {
            scheduledExecutor.shutdownNow();
        }
    }

    public enum ScheduleMode {
        FIXED_DELAY,
        FIXED_RATE
    }
}