import { SupportedModalities } from '~/types/model';

export function mergeSets(setA: Set<string>, setB: Set<string>): string[] {
    return Array.from(new Set([...setA, ...setB]));
}

export const getModelData = (modelId: string, providerId: string, modelsDevData: any) => {
    const modelData = modelsDevData[providerId]?.models[modelId];

    if (modelData === undefined) return {
        cost: {},
        attributes: {
            inputModalities: ['text'],
            outputModalities: ['text'],
            capabilities: [],
        }
    };

    const capabilities = new Array<string>();
    if (modelData.reasoning) {
        capabilities.push('reasoning');
    }

    if (modelData.tool_call) {
        capabilities.push('tools');
    }

    let inputModalities = modelData.modalities.input.filter((m: string) => (SupportedModalities as Readonly<string[]>).includes(m));
    let outputModalities = modelData.modalities.output.filter((m: string) => (SupportedModalities as Readonly<string[]>).includes(m));
    let supportedParameters = [];
    if (modelData.temperature) {
        supportedParameters.push('temperature');
    }

    let contextWindow = modelData.limit?.context || null;
    let cost: Record<string, string> = {};
    for (const key in modelData.cost) {
        switch (key) {
            case 'input': {
                cost.prompt = formatMoney(modelData.cost[key].toString());
            } break;
            case 'output': {
                cost.completion = formatMoney(modelData.cost[key].toString());
            } break;
        }
    }

    let releasedAt = (new Date(modelData.release_date)).getTime();

    return {
        name: modelData.name,
        attributes: {
            inputModalities,
            outputModalities,
            capabilities: capabilities || [],
            contextWindow,
            supported_parameters: modelData.supportedParameters,
        },
        cost,
        releasedAt,
    }
}

export const lshDecimal = (number: string, shift: number) => {
    if (number[0] === '-') {
        return undefined;
    }

    if (number.length === 0) {
        return '0.00';
    }

    let value = '';
    let [integerPart, fractionalPart] = number.split('.');

    if (!fractionalPart) {
        fractionalPart = '';
    }

    if (shift <= fractionalPart.length) {
        integerPart += fractionalPart.substring(0, shift);
        fractionalPart = fractionalPart.substring(shift);
    } else if (shift > fractionalPart.length) {
        integerPart += fractionalPart;
        for (let i = 0; i < shift - fractionalPart.length; i++) {
            integerPart += '0';
        }
        fractionalPart = '';
    }

    integerPart = integerPart!.replace(/^0+/, '');
    if (integerPart.length === 0) {
        integerPart = '0';
    }

    value += integerPart;

    if (fractionalPart.length > 0) {
        if (fractionalPart.length < 2) {
            fractionalPart += '0';
        }

        value += '.' + fractionalPart;
    } else {
        value += '.00';
    }

    return value;
}

export const formatMoney = (value: string) => {
    const [integerPart, fractionalPart] = value.split('.');

    if (fractionalPart === undefined) {
        return `${integerPart}.00`;
    }

    if (fractionalPart.length === 1) {
        return `${integerPart}.${fractionalPart}0`;
    }

    return `${integerPart}.${fractionalPart}`;
}