gemini-compat-endpoint

import type { Chat } from "@lmstudio/sdk";
import type { GenerationContext } from "./strategies/ModelStrategy";
import { loadThoughtState, saveThoughtState, computeContentHash, type ThoughtState } from "./thought-state";

/**
 * Minimal, modellunabhängiges Interface für Reasoning-/Signature-Handling.
 *
 * Strategien (Text, Image, ...) arbeiten nur noch gegen dieses Interface
 * und müssen die interne Logik (Hashing, Pruning, Modell-Spezifika) nicht kennen.
 */
export interface ReasoningStatePolicy {
    /**
     * Baut die Gemini-contents aus History + vorhandenem ThoughtState.
     * Implementierungen injizieren ggf. thoughtSignature-Felder in Parts.
     */
    buildContents(params: {
        history: Chat;
        context: GenerationContext;
        baseContentsBuilder: () => any[]; // z.B. bestehendes toGeminiMessages(history, ...)
    }): any[];

    /**
     * Nimmt die finale Modellantwort (inkl. Streaming-aggregiertem Zustand) entgegen
     * und aktualisiert den ThoughtState für den Turn.
     */
    updateFromResponse(params: {
        response: any;
        history: Chat;
        context: GenerationContext;
        collectedFullText: string;          // inkl. thoughts
        collectedToolCalls: Array<{ name: string; args: any }>; // alle ToolCalls des Turns
        capturedSignature?: string;         // vom Strategy-Code extrahierte letzte thoughtSignature
    }): void;
}

/**
 * No-op-Policy für Modelle ohne Thinking/Signatures.
 */
export class NoopReasoningPolicy implements ReasoningStatePolicy {
    buildContents(params: { history: Chat; context: GenerationContext; baseContentsBuilder: () => any[] }): any[] {
        return params.baseContentsBuilder();
    }

    updateFromResponse(): void {
        // absichtlich leer
    }
}

/**
 * Leichtgewichtige Policy speziell für gemini-3-pro-image-preview.
 *
 * - ignoriert Bild-Entwürfe in der Signatur-Logik
 * - speichert nur Signaturen für den aktuellen Turn (Text/Tools)
 * - pruned ThoughtState nach jedem Turn aggressiv auf ein Minimal-Set
 */
export class ImageLightweightReasoningPolicy implements ReasoningStatePolicy {
    buildContents(params: {
        history: Chat;
        context: GenerationContext;
        baseContentsBuilder: () => any[];
    }): any[] {
        const contents = params.baseContentsBuilder();

        const chatWd = params.context.ctl.getWorkingDirectory();
        const state = loadThoughtState(chatWd);
        const { debugChunks } = params.context;

        if (debugChunks) {
            const keys = Object.keys(state.signatures);
            console.info("[ImageReasoningPolicy] buildContents - signatures keys:", keys.slice(0, 10), "total=", keys.length);
        }

        // Aktuell keine Injection; der State wird nur beobachtet.
        return contents;
    }

    updateFromResponse(params: {
        response: any;
        history: Chat;
        context: GenerationContext;
        collectedFullText: string;
        collectedToolCalls: Array<{ name: string; args: any }>;
        capturedSignature?: string;
    }): void {
        const { capturedSignature, collectedFullText, collectedToolCalls, context } = params;
        if (!capturedSignature) return;

        const chatWd = context.ctl.getWorkingDirectory();
        const state: ThoughtState = loadThoughtState(chatWd);
        const { debugChunks } = context;

        // 1) Volltext (inkl. thoughts) hashen und auf Signatur mappen
        if (collectedFullText && collectedFullText.trim()) {
            const fullHash = computeContentHash(collectedFullText);
            state.signatures[fullHash] = capturedSignature;

            // optional: getrimmte Variante für Robustheit
            const trimmed = collectedFullText.trim();
            if (trimmed !== collectedFullText) {
                const trimmedHash = computeContentHash(trimmed);
                state.signatures[trimmedHash] = capturedSignature;
            }
        }

        // 2) Fallback-Schlüssel für "letzter Text-Turn"
        state.signatures["LATEST_TEXT_SIG"] = capturedSignature;

        // 3) ToolCalls des Turns auf dieselbe Signatur mappen
        for (const tc of collectedToolCalls || []) {
            const id = `${tc.name}:${JSON.stringify(tc.args ?? {})}`;
            const hash = computeContentHash(id);
            state.signatures[hash] = capturedSignature;
        }

        // 4) Aggressives Pruning: nur noch das Minimal-Set behalten.
        const minimal: ThoughtState = { signatures: {} };
        for (const [key, value] of Object.entries(state.signatures)) {
            // Wir behalten nur Einträge aus diesem Turn plus LATEST_TEXT_SIG.
            // Für das Skeleton gehen wir pragmatisch vor und behalten alles,
            // was gerade auf capturedSignature zeigt.
            if (value === capturedSignature || key === "LATEST_TEXT_SIG") {
                minimal.signatures[key] = value;
            }
        }

        saveThoughtState(chatWd, minimal);

        if (debugChunks) {
            const keys = Object.keys(minimal.signatures);
            console.info("[ImageReasoningPolicy] updateFromResponse - kept signature keys:", keys);
        }
    }
}

/**
 * Platzhalter für ein vollwertiges Text-Reasoning-Profil
 * (z.B. gemini-3-pro-preview). Aktuell identisch zu Noop,
 * kann später gezielt ausgebaut werden.
 */
export class TextFullReasoningPolicy implements ReasoningStatePolicy {
    buildContents(params: { history: Chat; context: GenerationContext; baseContentsBuilder: () => any[] }): any[] {
        return params.baseContentsBuilder();
    }

    updateFromResponse(params: {
        response: any;
        history: Chat;
        context: GenerationContext;
        collectedFullText: string;
        collectedToolCalls: Array<{ name: string; args: any }>;
        capturedSignature?: string;
    }): void {
        // Noch keine abweichende Logik – kann später
        // analog zur Image-Policy, aber weniger aggressiv,
        // implementiert werden.
        void params;
    }
}

export type ReasoningProfile = "none" | "text_full" | "image_lightweight";

export function getReasoningPolicy(profile: ReasoningProfile | undefined): ReasoningStatePolicy {
    switch (profile) {
        case "image_lightweight":
            return new ImageLightweightReasoningPolicy();
        case "text_full":
            return new TextFullReasoningPolicy();
        case "none":
        default:
            return new NoopReasoningPolicy();
    }
}