/*
 * 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.amoro.client;

import com.github.benmanes.caffeine.cache.Caffeine;
import com.github.benmanes.caffeine.cache.LoadingCache;
import org.apache.amoro.Constants;
import org.apache.amoro.api.AmoroTableMetastore;
import org.apache.amoro.shade.thrift.org.apache.thrift.TException;
import org.apache.amoro.shade.thrift.org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.amoro.shade.thrift.org.apache.thrift.protocol.TMultiplexedProtocol;
import org.apache.amoro.shade.thrift.org.apache.thrift.protocol.TProtocol;

/** Client pool cache for different ams server, sharing in jvm. */
public class AmsClientPools {

  private static final LoadingCache<String, ThriftClientPool<AmoroTableMetastore.Client>>
      CLIENT_POOLS = Caffeine.newBuilder().build(AmsClientPools::buildClientPool);

  public static ThriftClientPool<AmoroTableMetastore.Client> getClientPool(String metastoreUrl) {
    return CLIENT_POOLS.get(metastoreUrl);
  }

  public static void cleanAll() {
    CLIENT_POOLS.cleanUp();
  }

  @SuppressWarnings("unchecked")
  private static ThriftClientPool<AmoroTableMetastore.Client> buildClientPool(String url) {
    PoolConfig<AmoroTableMetastore.Client> poolConfig =
        (PoolConfig<AmoroTableMetastore.Client>) PoolConfig.forUrl(url);
    return new ThriftClientPool<>(
        url,
        s -> {
          TProtocol protocol = new TBinaryProtocol(s);
          return new AmoroTableMetastore.Client(
              new TMultiplexedProtocol(protocol, Constants.THRIFT_TABLE_SERVICE_NAME));
        },
        c -> {
          try {
            ((AmoroTableMetastore.Client) c).ping();
          } catch (TException e) {
            return false;
          }
          return true;
        },
        poolConfig,
        Constants.THRIFT_TABLE_SERVICE_NAME);
  }
}
