import { isAbortError } from 'abort-controller-x';
import ConnectionGQL from './gql.js';
import { getClientVersionHeader } from './http.js';
import { ChannelCredentials, createChannel, createClientFactory, Metadata, Status, } from 'nice-grpc';
import { retryMiddleware } from 'nice-grpc-client-middleware-retry';
import { HealthCheckResponse_ServingStatus, HealthDefinition } from '../proto/google/health/v1/health.js';
import { WeaviateDefinition } from '../proto/v1/weaviate.js';
import Batcher from '../grpc/batcher.js';
import Searcher from '../grpc/searcher.js';
import TenantsManager from '../grpc/tenantsManager.js';
import { DbVersionSupport, initDbVersionProvider } from '../utils/dbVersion.js';
import { WeaviateGRPCUnavailableError, WeaviateUnsupportedFeatureError } from '../errors.js';
import Aggregator from '../grpc/aggregator.js';
const clientFactory = createClientFactory().use(retryMiddleware);
const MAX_GRPC_MESSAGE_LENGTH = 104858000; // 10mb, needs to be synchronized with GRPC server
// Must extend from ConnectionGQL so that it can be passed to all the builder methods,
// which are tightly coupled to ConnectionGQL
export default class ConnectionGRPC extends ConnectionGQL {
    grpc;
    constructor(params) {
        super(params);
        this.grpc = grpcClient(params);
    }
    static use = (params) => {
        const rest = new ConnectionGQL(params);
        const dbVersionProvider = initDbVersionProvider(rest);
        const dbVersionSupport = new DbVersionSupport(dbVersionProvider);
        if (params.skipInitChecks) {
            return {
                connection: new ConnectionGRPC({
                    ...params,
                    grpcMaxMessageLength: MAX_GRPC_MESSAGE_LENGTH,
                }),
                dbVersionProvider,
                dbVersionSupport,
            };
        }
        return Promise.all([
            ConnectionGRPC.connect(params, rest.get('/meta', true).then((res) => res.grpcMaxMessageSize || MAX_GRPC_MESSAGE_LENGTH)),
            dbVersionSupport.supportsCompatibleGrpcService().then((check) => {
                if (!check.supports) {
                    throw new WeaviateUnsupportedFeatureError(`Checking for gRPC compatibility failed with message: ${check.message}`);
                }
            }),
        ]).then(([connection]) => {
            return { connection, dbVersionProvider, dbVersionSupport };
        });
    };
    static async connect(params, grpcMaxLengthPromise) {
        const connection = await grpcMaxLengthPromise.then((grpcMaxMessageLength) => new ConnectionGRPC({
            ...params,
            grpcMaxMessageLength,
        }));
        const isHealthy = await connection.grpc.health();
        if (!isHealthy) {
            await connection.close();
            throw new WeaviateGRPCUnavailableError(params.grpcAddress);
        }
        return connection;
    }
    batch = (collection, consistencyLevel, tenant) => {
        if (this.authEnabled) {
            return this.login().then((token) => this.grpc.batch(collection, consistencyLevel, tenant, `Bearer ${token}`));
        }
        return new Promise((resolve) => resolve(this.grpc.batch(collection, consistencyLevel, tenant)));
    };
    aggregate = (collection, consistencyLevel, tenant) => {
        if (this.authEnabled) {
            return this.login().then((token) => this.grpc.aggregate(collection, consistencyLevel, tenant, `Bearer ${token}`));
        }
        return new Promise((resolve) => resolve(this.grpc.aggregate(collection, consistencyLevel, tenant)));
    };
    search = (collection, consistencyLevel, tenant) => {
        if (this.authEnabled) {
            return this.login().then((token) => this.grpc.search(collection, consistencyLevel, tenant, `Bearer ${token}`));
        }
        return new Promise((resolve) => resolve(this.grpc.search(collection, consistencyLevel, tenant)));
    };
    tenants = (collection) => {
        if (this.authEnabled) {
            return this.login().then((token) => this.grpc.tenants(collection, `Bearer ${token}`));
        }
        return new Promise((resolve) => resolve(this.grpc.tenants(collection)));
    };
    close = () => {
        this.grpc.close();
        this.http.close();
    };
}
export const grpcClient = (config) => {
    const channelOptions = {
        'grpc.max_send_message_length': config.grpcMaxMessageLength,
        'grpc.max_receive_message_length': config.grpcMaxMessageLength,
    };
    if (config.grpcProxyUrl) {
        // grpc.http_proxy is not used by grpc.js under-the-hood
        // only uses the env var and whether http_proxy is enabled
        process.env.grpc_proxy = config.grpcProxyUrl;
        channelOptions['grpc.enabled_http_proxy'] = true;
    }
    const channel = createChannel(config.grpcAddress, config.grpcSecure ? ChannelCredentials.createSsl() : ChannelCredentials.createInsecure(), channelOptions);
    const client = clientFactory.create(WeaviateDefinition, channel);
    const health = clientFactory.create(HealthDefinition, channel);
    return {
        aggregate: (collection, consistencyLevel, tenant, bearerToken) => Aggregator.use(client, collection, getMetadataWithEmbeddingServiceAuth(config, bearerToken), config.timeout?.query || 30, consistencyLevel, tenant),
        batch: (collection, consistencyLevel, tenant, bearerToken) => Batcher.use(client, collection, getMetadataWithEmbeddingServiceAuth(config, bearerToken), config.timeout?.insert || 90, consistencyLevel, tenant),
        close: () => channel.close(),
        health: () => {
            const controller = new AbortController();
            const timeoutId = setTimeout(() => controller.abort(), (config.timeout?.init || 2) * 1000);
            return health
                .check({ service: '/grpc.health.v1.Health/Check' }, {
                signal: controller.signal,
                retry: true,
                retryMaxAttempts: 1,
                retryableStatuses: [Status.UNAVAILABLE],
                onRetryableError(error, attempt, delayMs) {
                    console.warn(error, `Healthcheck ${attempt} failed. Retrying in ${delayMs}ms.`);
                },
            })
                .then((res) => res.status === HealthCheckResponse_ServingStatus.SERVING)
                .catch((err) => {
                if (isAbortError(err)) {
                    throw new WeaviateGRPCUnavailableError(config.grpcAddress);
                }
                throw err;
            })
                .finally(() => clearTimeout(timeoutId));
        },
        search: (collection, consistencyLevel, tenant, bearerToken) => Searcher.use(client, collection, getMetadataWithEmbeddingServiceAuth(config, bearerToken), config.timeout?.query || 30, consistencyLevel, tenant),
        tenants: (collection, bearerToken) => TenantsManager.use(client, collection, new Metadata(bearerToken ? { ...config.headers, authorization: bearerToken } : config.headers), config.timeout?.query || 30),
    };
};
const getMetadataWithEmbeddingServiceAuth = (config, bearerToken) => new Metadata({
    ...(bearerToken
        ? {
            ...config.headers,
            authorization: bearerToken,
            'X-Weaviate-Cluster-Url': config.host,
        }
        : config.headers),
    ...getClientVersionHeader(),
});
