LingvAI/utils/replicate-translate.ts

import { readSettings } from "./settings-store";

type CogOutput = {
    status?: string;
    output?: unknown;
    error?: string;
};

type ReplicateOutput = string | string[] | { translation?: string; translated_text?: string; output?: string };

function extractText(output: unknown): string {
    if (typeof output === "string") return output;
    if (Array.isArray(output)) return output.join("");
    if (output && typeof output === "object") {
        const o = output as { translation?: string; translated_text?: string; output?: string };
        return o.translation ?? o.translated_text ?? o.output ?? JSON.stringify(output);
    }
    throw new Error("Unexpected output format from model");
}

/**
 * Parse modelVersion string into components.
 * Accepts:
 *   "jigsawstack/text-translate:454df4c..."  → { owner: "jigsawstack", model: "text-translate", hash: "454df4c..." }
 *   "jigsawstack/text-translate"             → { owner: "jigsawstack", model: "text-translate", hash: undefined }
 *   "454df4c..."                             → { owner: undefined, model: undefined, hash: "454df4c..." }
 */
function parseModelVersion(mv: string) {
    const [ownerModel, hash] = mv.split(":");
    if (ownerModel.includes("/")) {
        const [owner, model] = ownerModel.split("/");
        return { owner, model, hash };
    }
    // Bare hash
    return { owner: undefined, model: undefined, hash: ownerModel };
}

async function translateViaCloud(
    text: string,
    targetLanguage: string
): Promise<string> {
    const settings = readSettings();

    if (!settings.replicateApiToken) throw new Error("Replicate API token not configured");
    if (!settings.jigsawApiKey) throw new Error("JigsawStack API key not configured");

    const { owner, model, hash } = parseModelVersion(settings.modelVersion);

    let url: string;
    let body: Record<string, unknown>;

    if (owner && model && !hash) {
        // No version pinned — use the model's deployment endpoint (latest)
        url = `https://api.replicate.com/v1/models/${owner}/${model}/predictions`;
        body = { input: { text, api_key: settings.jigsawApiKey, target_language: targetLanguage } };
    } else if (owner && model && hash) {
        // Pinned version — use model endpoint with version (avoids 422 from bare /predictions)
        url = `https://api.replicate.com/v1/models/${owner}/${model}/predictions`;
        body = {
            version: hash,
            input: { text, api_key: settings.jigsawApiKey, target_language: targetLanguage }
        };
    } else {
        // Bare hash only — use the generic predictions endpoint
        url = "https://api.replicate.com/v1/predictions";
        body = {
            version: hash,
            input: { text, api_key: settings.jigsawApiKey, target_language: targetLanguage }
        };
    }

    const response = await fetch(url, {
        method: "POST",
        headers: {
            "Authorization": `Bearer ${settings.replicateApiToken}`,
            "Content-Type": "application/json",
            "Prefer": "wait"
        },
        body: JSON.stringify(body)
    });

    if (!response.ok) {
        const err = await response.text();
        throw new Error(`Replicate API error: ${response.status} ${err}`);
    }

    const data = await response.json();
    if (data.error) throw new Error(`Replicate model error: ${data.error}`);

    return extractText(data.output as ReplicateOutput);
}

async function translateViaLocal(
    text: string,
    targetLanguage: string
): Promise<string> {
    const settings = readSettings();

    if (!settings.jigsawApiKey) throw new Error("JigsawStack API key not configured");

    const endpoint = (settings.localEndpoint || "http://localhost:5030/predictions").replace(/\/$/, "");

    const response = await fetch(endpoint, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({
            input: {
                text,
                api_key: settings.jigsawApiKey,
                target_language: targetLanguage
            }
        })
    });

    if (!response.ok) {
        const err = await response.text();
        throw new Error(`Local model error: ${response.status} ${err}`);
    }

    const data: CogOutput = await response.json();

    if (data.error) throw new Error(`Local model error: ${data.error}`);

    // Cog wraps output in { status, output } or returns output directly
    const output = data.output !== undefined ? data.output : data;
    return extractText(output);
}

export async function replicateTranslate(
    text: string,
    targetLanguage: string
): Promise<string> {
    const { replicateMode } = readSettings();
    return replicateMode === "local"
        ? translateViaLocal(text, targetLanguage)
        : translateViaCloud(text, targetLanguage);
}

// Batch translate using separator trick to minimize API calls
const SEPARATOR = "\n{{SEP}}\n";

export async function replicateTranslateBatch(
    texts: string[],
    targetLanguage: string
): Promise<string[]> {
    if (texts.length === 0) return [];
    if (texts.length === 1) {
        return [await replicateTranslate(texts[0], targetLanguage)];
    }

    const joined = texts.join(SEPARATOR);
    const translated = await replicateTranslate(joined, targetLanguage);

    const parts = translated.split(SEPARATOR);
    if (parts.length === texts.length) return parts;

    // Fallback: translate individually if separator got mangled
    return Promise.all(texts.map(t => replicateTranslate(t, targetLanguage)));
}