gemini-compat-endpoint

import { type Chat, type GeneratorController } from "@lmstudio/sdk";
import { sanitizeToolName, sanitizeToolResponseForModel } from "./tools";
import { computeContentHash } from "./thought-signatures";

export function stableJsonStringify(obj: any): string {
    if (typeof obj !== "object" || obj === null) return JSON.stringify(obj);
    if (Array.isArray(obj)) return "[" + obj.map(stableJsonStringify).join(",") + "]";
    const keys = Object.keys(obj).sort();
    const parts = keys.map(key => JSON.stringify(key) + ":" + stableJsonStringify(obj[key]));
    return "{" + parts.join(",") + "}";
}

export function safeStringify(obj: any, secrets: string[] = []): string {
    const secretSet = new Set(secrets.filter(Boolean));
    const redactor = (key: string, value: any) => {
        if (typeof value === "string") {
            // Field-based redaction
            if (/^(x-goog-api-key|authorization|apiKey|api_key)$/i.test(key)) {
                return "***";
            }
        }
        return value;
    };
    let s = "";
    try { s = JSON.stringify(obj, redactor); } catch {
        try { s = String(obj); } catch { s = "[unserializable]"; }
    }
    // Content-based redaction
    for (const sec of secretSet) {
        if (!sec) continue;
        try {
            const esc = sec.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
            s = s.replace(new RegExp(esc, "g"), "***");
        } catch { /* ignore */ }
    }
    return s;
}

export function normalizeFunctionResponsePayload(payload: any): any {
    if (payload === null || payload === undefined) return {};
    if (Array.isArray(payload)) return { result: payload };
    if (typeof payload !== "object") return { result: payload };
    return payload;
}

export function toGeminiMessages(history: Chat, nameMap?: Map<string, string>, signatures: Array<{ contentHash: string; signature: string }> = []): any[] {
    const contents: any[] = [];
    const signatureMap = new Map(signatures.map(s => [s.contentHash, s.signature]));

    function getSignatureForText(text: string): string | undefined {
        const hash = computeContentHash(text);
        if (signatureMap.has(hash)) return signatureMap.get(hash);
        if (signatureMap.has("LATEST_TEXT_SIG")) return signatureMap.get("LATEST_TEXT_SIG");
        return undefined;
    }

    function getSignatureForFunctionCall(name: string, args: any): string | undefined {
        const id = `${name}:${stableJsonStringify(args ?? {})}`;
        const hash = computeContentHash(id);
        if (signatureMap.has(hash)) return signatureMap.get(hash);
        // Fallback to LATEST_TEXT_SIG if available. This is critical for retries where the exact hash might be missed,
        // or if the tool call signature wasn't explicitly saved under the tool-hash key in previous versions.
        // While this might technically map a newer signature to an older call in rare cases, it prevents 400 errors.
        if (signatureMap.has("LATEST_TEXT_SIG")) return signatureMap.get("LATEST_TEXT_SIG");
        return undefined;
    }

    for (const message of history) {
        const parts: any[] = [];

        switch (message.getRole()) {
            case "system":
                // System messages will be handled via systemInstruction; skip here
                break;

            case "user": {
                const parsed = parseAttachmentWrappers(message.getText());
                if (parsed.text.trim().length > 0) {
                    parts.push({ text: parsed.text });
                }
                // Handle attachments nur textuell; wir erzeugen KEINE echten Bild-Parts
                // in der Reasoning-History, um thought_signature-Pflichten zu vermeiden.
                // NOTE: We intentionally do NOT add text for image URLs to avoid leaking
                // absolute paths (like /Users/helge/...) to the model. Images are promoted
                // separately via the vision promotion system.
                for (const p of parsed.parts) {
                    if (p.kind === "image" && p.url) {
                        // Skip - images are handled by vision promotion, don't leak paths
                    } else if (p.kind === "text" && p.text) {
                        parts.push({ text: p.text });
                    } else if (p.kind === "text_link" && p.url) {
                        // Only show basename for text links to avoid leaking paths
                        const basename = p.url.split('/').pop() || p.url;
                        parts.push({ text: `Attached file: ${basename}` });
                    }
                }
                contents.push({ role: "user", parts });
                break;
            }

            case "assistant": {
                // Include assistant text and any functionCall requests from history
                const text = message.getText();
                if (text && text.trim()) {
                    const part: any = { text };
                    const sig = getSignatureForText(text);
                    if (sig) part.thought_signature = sig;
                    parts.push(part);
                }
                const calls = (message as any).getToolCallRequests?.() as any[] | undefined;
                if (Array.isArray(calls) && calls.length) {
                    for (const c of calls) {
                        const orig = String(c?.name || "tool");
                        const safe = nameMap?.get(orig) || sanitizeToolName(orig);
                        const args = c?.arguments ?? {};
                        const part: any = { functionCall: { name: safe, args } };
                        const sig = getSignatureForFunctionCall(safe, args);
                        if (sig) part.thought_signature = sig;
                        parts.push(part);
                    }
                }
                if (parts.length) contents.push({ role: "model", parts });
                break;
            }

            case "tool": {
                // Map tool call results to Gemini functionResponse parts
                const results = (message as any).getToolCallResults() as any[];
                for (const r of results) {
                    const rName = (r?.toolName || r?.name || r?.functionName || r?.tool || r?.id || "tool");
                    const safeName = nameMap?.get(rName) || sanitizeToolName(String(rName));
                    let payload: any = r?.content ?? r?.result ?? r?.output ?? null;
                    if (typeof payload === "string") {
                        try { payload = JSON.parse(payload); } catch { /* keep as string */ }
                    }
                    // Sanitize tool result for model context (avoid double-render of tool-provided image objects)
                    payload = sanitizeToolResponseForModel(payload, "B");
                    const responsePayload = normalizeFunctionResponsePayload(payload);
                    const frPart: any = { functionResponse: { name: safeName, response: responsePayload } };

                    const sig = getSignatureForFunctionCall(safeName, (r?.arguments ?? {}));
                    if (sig) frPart.thought_signature = sig;

                    parts.push(frPart);
                }
                if (parts.length) contents.push({ role: "user", parts });
                break;
            }
        }
    }

    return contents;
}

export function parseAttachmentWrappers(
    text: string,
): { text: string; parts: Array<{ kind: "image" | "text" | "text_link"; url?: string; text?: string }> } {
    const parts: Array<{ kind: "image" | "text" | "text_link"; url?: string; text?: string }> = [];
    if (!text) return { text, parts };

    const re = /\[\[LMSTUDIO_ATTACHMENT:\s*(\{[\s\S]*?\})\s*\]\]/g;
    let cleaned = text;
    let m: RegExpExecArray | null;
    while ((m = re.exec(text)) !== null) {
        try {
            const obj = JSON.parse(m[1]);
            if (obj && obj.kind === "image" && typeof obj.url === "string") {
                parts.push({ kind: "image", url: obj.url });
            } else if (obj && obj.kind === "text" && typeof obj.text === "string") {
                parts.push({ kind: "text", text: obj.text });
            } else if (obj && obj.kind === "text_link" && typeof obj.url === "string") {
                parts.push({ kind: "text_link", url: obj.url });
            }
            cleaned = cleaned.replace(m[0], "");
        } catch {
            // ignore malformed wrappers
        }
    }
    // Also detect plain file:// URIs and absolute image paths; do not remove from text
    try {
        const seen = new Set<string>();
        const pushOnce = (u: string) => { if (!seen.has(u)) { parts.push({ kind: 'image', url: u }); seen.add(u); } };
        let m2: RegExpExecArray | null;
        const reFile = /file:\/\/[^\s)]+/gi;
        while ((m2 = reFile.exec(text)) !== null) pushOnce(m2[0]);
        const reAbs = /(^|\s)(\/[\w@#$%&+.,:\-\/]+?\.(?:png|jpe?g|webp|gif|bmp|tiff?|heic))(\s|$|[)])/gi;
        while ((m2 = reAbs.exec(text)) !== null) pushOnce(m2[2]);
    } catch { /* ignore */ }
    return { text: cleaned, parts };
}

export function getLastUserText(history: Chat): string {
    let last = "";
    for (const msg of history) {
        if (msg.getRole() === "user") last = msg.getText() || "";
    }
    return last;
}

export function collectSystemText(history: Chat): string {
    const parts: string[] = [];
    for (const msg of history) {
        if (msg.getRole() === "system") {
            const t = msg.getText();
            if (t && t.trim()) parts.push(t.trim());
        }
    }
    return parts.join("\n\n");
}

export function flattenToolCallsToText(contents: any[]) {
    for (const msg of contents) {
        if (!Array.isArray(msg.parts)) continue;
        msg.parts = msg.parts.map((p: any) => {
            if (p.functionCall) {
                return { text: `\n[System: The model previously called tool '${p.functionCall.name}' with arguments: ${JSON.stringify(p.functionCall.args)}]\n` };
            }
            if (p.functionResponse) {
                const resp = p.functionResponse.response;
                let text = "[System: Tool execution completed.]";
                // Optional: Handle image markdown specially if your tool returns it
                if (resp && Array.isArray(resp.images)) {
                    const mds = resp.images.map((img: any) => img.markdown || "").join("\n");
                    text = `[System: The tool successfully generated an image...]\n\n${mds}`;
                } else {
                    try { text = `[System: Tool returned: ${JSON.stringify(resp)}]`; } catch { }
                }
                return { text };
            }
            return p;
        });
    }
}

export function pad2(n: number) { return n.toString().padStart(2, "0"); }

export function formatTimestamp(d: Date) {
    const yyyy = d.getFullYear();
    const MM = pad2(d.getMonth() + 1);
    const DD = pad2(d.getDate());
    const hh = pad2(d.getHours());
    const mm = pad2(d.getMinutes());
    const ss = pad2(d.getSeconds());
    return `${yyyy}${MM}${DD}${hh}${mm}${ss}`;
}

export async function streamTextFragments(ctl: GeneratorController, text: string) {
    const chunkSize = 512;
    for (let i = 0; i < text.length; i += chunkSize) {
        const chunk = text.slice(i, i + chunkSize);
        ctl.fragmentGenerated(chunk);
        // Small yield to keep UI responsive; adjust or remove if undesired
        // eslint-disable-next-line no-await-in-loop
        await new Promise(res => setTimeout(res, 10));
    }
}