package io.milvus.pool;

import io.milvus.v2.exception.ErrorCode;
import io.milvus.v2.exception.MilvusClientException;
import org.apache.commons.pool2.BaseKeyedPooledObjectFactory;
import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.impl.DefaultPooledObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.lang.reflect.Constructor;
import java.lang.reflect.Method;

public class PoolClientFactory<C, T> extends BaseKeyedPooledObjectFactory<String, T> {
    protected static final Logger logger = LoggerFactory.getLogger(PoolClientFactory.class);
    private final C config;
    private Constructor<?> constructor;
    private Method closeMethod;
    private Method verifyMethod;

    public PoolClientFactory(C config, String clientClassName) throws ClassNotFoundException, NoSuchMethodException {
        this.config = config;
        try {
            Class<?> clientCls = Class.forName(clientClassName);
            Class<?> configCls = Class.forName(config.getClass().getName());
            constructor = clientCls.getConstructor(configCls);
            closeMethod = clientCls.getMethod("close", long.class);
            verifyMethod = clientCls.getMethod("clientIsReady");
        } catch (Exception e) {
            logger.error("Failed to create client pool factory, exception: ", e);
            throw e;
        }
    }

    @Override
    public T create(String key) throws Exception {
        try {
            T client = (T) constructor.newInstance(this.config);
            return client;
        } catch (Exception e) {
            logger.error("Failed to create client, exception: ", e);
            throw new MilvusClientException(ErrorCode.CLIENT_ERROR, e);
        }
    }

    @Override
    public PooledObject<T> wrap(T client) {
        return new DefaultPooledObject<>(client);
    }

    @Override
    public void destroyObject(String key, PooledObject<T> p) throws Exception {
        T client = p.getObject();
        closeMethod.invoke(client, 3L);
    }

    @Override
    public boolean validateObject(String key, PooledObject<T> p) {
        try {
            T client = p.getObject();
            return (boolean) verifyMethod.invoke(client);
        } catch (Exception e) {
            logger.error("Failed to validate client, exception: " + e);
            throw new MilvusClientException(ErrorCode.CLIENT_ERROR, e);
        }
    }

    @Override
    public void activateObject(String key, PooledObject<T> p) throws Exception {
        super.activateObject(key, p);
    }

    @Override
    public void passivateObject(String key, PooledObject<T> p) throws Exception {
        super.passivateObject(key, p);
    }
}
