import { type StreamTextTransform } from "ai";
import type { Provider as ProviderDrizzle, Model as ModelDrizzle } from '~/composables/useModels';
import { type Result, Err, Ok } from "~~/types/result";
import { getProvider, getProviderBaseUrl } from './providers';

export type ModelGateway = any;

export interface Gateway {
    gateway: ModelGateway;
    streamTransformer: StreamTextTransform<{}> | StreamTextTransform<{}>[] | undefined;
    textTransformer: ((text: string) => string) | ((text: string) => string)[] | undefined;
}

export interface Provider {
    gateway: Gateway | null;
    endpoint: {
        baseURL: string;
        modelsEndpoint: string | null;
        headers: Record<string, string>;
    }
}

export enum GatewayFetchError {
    NoProviderApiKey = 0,
    NoProviderBaseUrl,
}

export async function getProviderDetails(provider: ProviderDrizzle, providerApiKey?: string, model?: ModelDrizzle): Promise<Result<Provider, GatewayFetchError>> {
    const providerModule = getProvider(provider.type);

    if (!providerModule) {
        return Err(GatewayFetchError.NoProviderBaseUrl);
    }

    let baseURL = undefined;
    if (provider.config.apiProxyUrl && provider.config.apiProxyUrl.trim() !== '') {
        baseURL = provider.config.apiProxyUrl;
    } else {
        baseURL = providerModule.baseUrl;
    }
    baseURL = baseURL.replace(/\/$/, '');

    const gatewayResult = providerModule.createGateway({
        apiKey: providerApiKey,
        baseURL,
        model,
    });

    if (!gatewayResult.ok) {
        return gatewayResult;
    }

    const rawHeaders = providerModule.getAuthHeaders(providerApiKey);
    const headers: Record<string, string> = {};
    for (const [key, value] of Object.entries(rawHeaders)) {
        if (value !== undefined) {
            headers[key] = value;
        }
    }

    return Ok({
        gateway: gatewayResult.data,
        endpoint: {
            baseURL,
            modelsEndpoint: providerModule.modelsEndpoint,
            headers,
        }
    });
}

export function isRerankingProvider(gateway: any): gateway is { reranking: (modelId: string) => any } {
    return (
        gateway !== null &&
        'reranking' in gateway &&
        typeof (gateway as any).reranking === 'function'
    );
}