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);
}

// The local Cog model handles exactly one prediction at a time.
// This queue ensures calls are strictly serialised regardless of
// how many batches / concurrent requests come in.
let localQueue: Promise<void> = Promise.resolve();

async function callLocalModel(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}`);

    const output = data.output !== undefined ? data.output : data;
    return extractText(output);
}

async function translateViaLocal(
    text: string,
    targetLanguage: string
): Promise<string> {
    // Enqueue — wait for any in-flight local call to finish first
    let resolve!: () => void;
    const slot = new Promise<void>(r => { resolve = r; });
    const turn = localQueue.then(() => callLocalModel(text, targetLanguage));
    localQueue = turn.then(resolve, resolve); // advance queue whether success or error
    return turn;
}

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

const SEPARATOR = "\n{{SEP}}\n";
const MAX_CHARS = 4800; // JigsawStack limit is 5000 — leave headroom

/**
 * Split a single text that exceeds MAX_CHARS into translatable chunks,
 * preferring paragraph/sentence boundaries.
 */
async function translateLongText(text: string, targetLanguage: string): Promise<string> {
    const chunks: string[] = [];
    let remaining = text;

    while (remaining.length > MAX_CHARS) {
        let splitAt = MAX_CHARS;
        const para = remaining.lastIndexOf("\n", MAX_CHARS);
        const sentence = remaining.lastIndexOf(". ", MAX_CHARS);
        if (para > MAX_CHARS * 0.5) splitAt = para + 1;
        else if (sentence > MAX_CHARS * 0.5) splitAt = sentence + 2;
        chunks.push(remaining.slice(0, splitAt));
        remaining = remaining.slice(splitAt);
    }
    if (remaining.trim()) chunks.push(remaining);

    const results = await Promise.all(chunks.map(c => replicateTranslate(c, targetLanguage)));
    return results.join(" ");
}

/**
 * Build groups of texts that fit within MAX_CHARS when joined with SEPARATOR.
 * Texts that individually exceed MAX_CHARS are kept alone for chunked translation.
 */
function buildBatches(texts: string[]): { indices: number[]; long: boolean }[] {
    const batches: { indices: number[]; long: boolean }[] = [];
    let current: number[] = [];
    let currentLen = 0;

    for (let i = 0; i < texts.length; i++) {
        const t = texts[i];

        if (t.length > MAX_CHARS) {
            // Flush current group first
            if (current.length > 0) { batches.push({ indices: current, long: false }); current = []; currentLen = 0; }
            batches.push({ indices: [i], long: true });
            continue;
        }

        const added = currentLen === 0 ? t.length : currentLen + SEPARATOR.length + t.length;
        if (added > MAX_CHARS && current.length > 0) {
            batches.push({ indices: current, long: false });
            current = [i];
            currentLen = t.length;
        } else {
            current.push(i);
            currentLen = added;
        }
    }
    if (current.length > 0) batches.push({ indices: current, long: false });
    return batches;
}

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

    const results: string[] = new Array(texts.length);
    const batches = buildBatches(texts);

    // Process batches sequentially to avoid hammering the model
    for (const batch of batches) {
        if (batch.long) {
            // Single oversized text — chunk it
            results[batch.indices[0]] = await translateLongText(texts[batch.indices[0]], targetLanguage);
        } else if (batch.indices.length === 1) {
            results[batch.indices[0]] = await replicateTranslate(texts[batch.indices[0]], targetLanguage);
        } else {
            // Multi-text batch within limit
            const joined = batch.indices.map(i => texts[i]).join(SEPARATOR);
            const translated = await replicateTranslate(joined, targetLanguage);
            const parts = translated.split(SEPARATOR);

            if (parts.length === batch.indices.length) {
                batch.indices.forEach((idx, i) => { results[idx] = parts[i]; });
            } else {
                // Separator got translated — fall back to individual calls
                const individual = await Promise.all(
                    batch.indices.map(i => replicateTranslate(texts[i], targetLanguage))
                );
                batch.indices.forEach((idx, i) => { results[idx] = individual[i]; });
            }
        }
    }

    return results;
}