tsgames/src/games/storywriter/utils/prompt.ts

import { type AppState, CharacterRole, type ChatMessage } from "../contexts/state";
import Chapters from "./chapters";
import LLM from "./llm";
import { Tools } from "./tools";

namespace Prompt {
    function approxTokens(text: string): number {
        return text.length / 3;
    }

    const KEEP_RECENT_CHUNKS = 2;

    interface ChunkSlot {
        header: string;
        body: string;
        summary: string | null;
        mode: 'full' | 'summary' | 'omitted';
    }

    function buildSlots(text: string, chapters: Chapters.Chapter[]): ChunkSlot[] {
        const parsed = Chapters.parseText(text);
        const slots: ChunkSlot[] = [];

        for (const parsedChapter of parsed) {
            const cachedChapter = chapters.find(c => c.header === parsedChapter.header)
                ?? Chapters.emptyChapter(parsedChapter.header);
            const chunks = Chapters.splitIntoChunks(parsedChapter.body);

            for (const body of chunks) {
                const { summary } = Chapters.lookupSummary(cachedChapter, body);
                slots.push({ header: parsedChapter.header, body, summary, mode: 'full' });
            }
        }

        return slots;
    }

    function countSlotTokens(slots: ChunkSlot[]): number {
        let total = 0;
        const countedHeaders = new Set<string>();

        for (const slot of slots) {
            if (slot.mode === 'omitted') continue;
            if (slot.header && !countedHeaders.has(slot.header)) {
                total += approxTokens(slot.header);
                countedHeaders.add(slot.header);
            }
            total += approxTokens(slot.mode === 'summary' ? (slot.summary ?? '') : slot.body);
        }

        return total;
    }

    function renderSlots(slots: ChunkSlot[]): string {
        const parts: string[] = [];
        const shownHeaders = new Set<string>();

        for (const slot of slots) {
            const lines: string[] = [];

            if (slot.header && !shownHeaders.has(slot.header)) {
                lines.push(slot.header);
                shownHeaders.add(slot.header);
            }

            const content = slot.mode === 'omitted' ? '[...]'
                : slot.mode === 'summary' ? `[Summary: ${slot.summary}]`
                    : slot.body;
            lines.push(content);
            parts.push(lines.join('\n\n'));
        }

        return parts.join('\n\n');
    }

    export function formatStoryChunks(
        text: string,
        chapters: Chapters.Chapter[],
        tokenBudget: number,
    ): string {
        if (!text) return '';

        const slots = buildSlots(text, chapters);
        if (slots.length === 0) return '';

        if (countSlotTokens(slots) <= tokenBudget) {
            return renderSlots(slots);
        }

        const recentStart = Math.max(0, slots.length - KEEP_RECENT_CHUNKS);

        // Phase 1: summarize non-recent chunks, stop as soon as we fit
        for (let i = 0; i < recentStart; i++) {
            if (slots[i].summary) {
                slots[i].mode = 'summary';
                if (countSlotTokens(slots) <= tokenBudget) return renderSlots(slots);
            }
        }

        // Phase 2: delete from middle outward, never delete last slot
        const middle = recentStart / 2;
        const deletable = Array.from({ length: recentStart }, (_, i) => i)
            .sort((a, b) => Math.abs(a - middle) - Math.abs(b - middle));

        for (const i of deletable) {
            slots[i].mode = 'omitted';
            if (countSlotTokens(slots) <= tokenBudget) break;
        }

        return renderSlots(slots);
    }

    /**
     * Character detail configuration based on role.
     * Determines how much information to include in the system prompt for each character.
     * AI agents can retrieve full character details on-demand via the get_character tool.
     */
    interface CharacterDetailConfig {
        includeRole: boolean;
        includeNicknames: boolean;
        includeShortDescription: boolean;
        includeFullDescription: boolean;
        includeRelations: boolean;
    }

    const CHARACTER_ROLE_DETAIL: Record<CharacterRole, CharacterDetailConfig> = {
        [CharacterRole.Protagonist]: {
            includeRole: true,
            includeNicknames: true,
            includeShortDescription: false, // omitted if full description present
            includeFullDescription: true,
            includeRelations: true,
        },
        [CharacterRole.Main]: {
            includeRole: true,
            includeNicknames: true,
            includeShortDescription: false,
            includeFullDescription: true,
            includeRelations: true,
        },
        [CharacterRole.Antagonist]: {
            includeRole: true,
            includeNicknames: true,
            includeShortDescription: false,
            includeFullDescription: true,
            includeRelations: true,
        },
        [CharacterRole.Secondary]: {
            includeRole: true,
            includeNicknames: false,
            includeShortDescription: true,
            includeFullDescription: false,
            includeRelations: true,
        },
        [CharacterRole.Supporting]: {
            includeRole: true,
            includeNicknames: false,
            includeShortDescription: false,
            includeFullDescription: false,
            includeRelations: false,
        },
        [CharacterRole.Minor]: {
            includeRole: true,
            includeNicknames: false,
            includeShortDescription: false,
            includeFullDescription: false,
            includeRelations: false,
        },
        [CharacterRole.Cameo]: {
            includeRole: true,
            includeNicknames: false,
            includeShortDescription: false,
            includeFullDescription: false,
            includeRelations: false,
        },
    };

    export function formatCharactersMarkdown(state: AppState): string {
        const { mergedCharacters } = state;
        if (!mergedCharacters?.length) {
            return '';
        }

        const lines: string[] = [];
        lines.push('## Characters\n');

        // Sort characters by importance (protagonist first, cameo last)
        const sortedCharacters = [...mergedCharacters].sort((a, b) => {
            const importanceOrder = [
                CharacterRole.Protagonist,
                CharacterRole.Main,
                CharacterRole.Antagonist,
                CharacterRole.Secondary,
                CharacterRole.Supporting,
                CharacterRole.Minor,
                CharacterRole.Cameo,
            ];
            return importanceOrder.indexOf(a.role) - importanceOrder.indexOf(b.role);
        });

        for (const character of sortedCharacters) {
            const config = CHARACTER_ROLE_DETAIL[character.role];
            lines.push(`### ${character.name}`);

            if (config.includeRole) {
                lines.push(`**Role:** ${character.role}`);
            }

            if (config.includeNicknames && character.nicknames?.length) {
                lines.push(`**Nicknames:** ${character.nicknames.join(', ')}`);
            }

            const description = character.description || character.shortDescription;
            if (config.includeFullDescription && description) {
                lines.push(description);
            } else if (config.includeShortDescription && character.shortDescription) {
                lines.push(character.shortDescription);
            }

            if (config.includeRelations && character.relations?.length) {
                lines.push('**Relations:**');
                for (const relation of character.relations) {
                    lines.push(`- ${relation.name}: ${relation.relation}`);
                }
            }

            lines.push('');
        }

        return lines.join('\n');
    }

    export function formatLoreMarkdown(state: AppState): string {
        const { mergedLore } = state;
        if (!mergedLore?.length) {
            return '';
        }

        const lines: string[] = [];
        lines.push('## Lore\n');

        for (const entry of mergedLore) {
            lines.push(`### ${entry.title}`);
            if (entry.text) {
                lines.push(entry.text);
            }
            lines.push('');
        }

        return lines.join('\n');
    }

    export function formatUserSection(state: AppState): string {
        if (!state.userName) {
            return '';
        }

        return state.userDescription
            ? `## ${state.userName}:\n${state.userDescription}`
            : `## User name: ${state.userName}`;
    }

    export function formatLocationsMarkdown(state: AppState): string {
        const { mergedLocations } = state;
        if (!mergedLocations?.length) {
            return '';
        }

        const lines: string[] = [];
        lines.push('## Locations\n');

        for (const location of mergedLocations) {
            lines.push(`### ${location.name}`);

            const description = location.shortDescription || location.description;
            if (description) {
                lines.push(description);
            }

            lines.push(`**Scale:** ${location.scale}`);
            lines.push('');
        }

        return lines.join('\n');
    }

    export function substituteVars(state: AppState, text: string): string {
        const charName = state.currentWorld?.title || 'Assistant';
        const userName = state.userName || 'User';
        return text
            .replaceAll('{{system}}', state.chatSystemInstruction)
            .replaceAll('{{char}}', charName)
            .replaceAll('{{user}}', userName);
    }

    export function formatSystemPrompt(state: AppState, storyTokenBudget: number = 0): string {
        const { currentStory, currentWorld } = state;
        if (!currentStory) {
            return state.effectiveSystemInstruction;
        }

        const parts: string[] = [state.effectiveSystemInstruction];

        // Chat-only worlds: system instruction + user info, no story scaffolding
        if (currentWorld?.chatOnly) {
            parts.unshift('# **Roleplay Context**');
            const userSection = formatUserSection(state);
            if (userSection) {
                parts.push(userSection);
            }

            if (currentStory.scratchpad) {
                parts.push(`## Scratchpad`, currentStory.scratchpad);
            }
            parts.push('### **End of Roleplay Context**');
        } else {
            parts.push(`# Story Title: ${currentStory.title}`);

            const loreSection = formatLoreMarkdown(state);
            if (loreSection) {
                parts.push(loreSection);
            }

            const charactersSection = formatCharactersMarkdown(state);
            if (charactersSection) {
                parts.push(charactersSection);
            }

            const locationsSection = formatLocationsMarkdown(state);
            if (locationsSection) {
                parts.push(locationsSection);
            }

            if (currentStory.scratchpad) {
                parts.push(`## Scratchpad`, currentStory.scratchpad);
            }

            if (currentStory.text && storyTokenBudget > 0) {
                const storyText = formatStoryChunks(currentStory.text, currentStory.chapters ?? [], storyTokenBudget);
                if (storyText) {
                    parts.push(`## Story Text:`, storyText);
                }
            }
        }

        return parts.join('\n\n');
    }

    export function compilePrompt(
        state: AppState,
        newMessages: Iterable<ChatMessage> = [],
        excludedMessageIds: Iterable<string> = [],
    ): LLM.ChatCompletionRequest | null {
        const { currentStory, model, enableThinking, currentWorld } = state;

        if (!currentStory || !model) {
            return null;
        }

        const messages: ChatMessage[] = [];
        const excludedMessages = new Set(excludedMessageIds);
        for (const message of currentStory.chatMessages) {
            if (!excludedMessages.has(message.id)) {
                messages.push(message);
                excludedMessages.add(message.id);
            }
        }

        for (const message of newMessages) {
            if (!excludedMessages.has(message.id)) {
                messages.push(message);
                excludedMessages.add(message.id);
            }
        }

        const charName = currentWorld?.title || 'Assistant';
        const userName = state.userName || 'User';
        const applyVars = (msgs: ChatMessage[]) =>
            msgs.map(m => ({ ...m, content: substituteVars(state, m.content) }));

        // Chat-only world: format messages with name prefixes
        if (currentWorld?.chatOnly) {
            const formattedMessages: ChatMessage[] = messages.map(msg => {
                const prefix = msg.role === 'user' ? userName : charName;
                return { ...msg, content: `${prefix}: ${msg.content}` };
            });

            // Prepend system message
            formattedMessages.unshift({
                id: crypto.randomUUID(),
                role: 'system',
                content: formatSystemPrompt(state, 0),
            });

            return {
                model: model.id,
                messages: applyVars(formattedMessages),
                enable_thinking: false,
                max_tokens: model.max_length ? model.max_length : 2048,
                add_generation_prompt: true,
                banned_tokens: state.bannedTokens,
            };
        }

        // Estimate token budget for story text
        let storyTokenBudget = 0;
        if (model.max_context) {
            const nonStorySystem = formatSystemPrompt(state, 0);
            const chatText = messages.map(m => m.content).join('\n');
            const maxOutput = model.max_length ?? 2048;
            const otherTokens = approxTokens(nonStorySystem) + approxTokens(chatText) + maxOutput;
            storyTokenBudget = model.max_context - otherTokens;
        }

        messages.unshift({
            id: crypto.randomUUID(),
            role: 'system',
            content: formatSystemPrompt(state, storyTokenBudget),
        });

        return {
            model: model.id,
            messages: applyVars(messages),
            tools: Tools.getTools(),
            banned_tokens: state.bannedTokens,
            enable_thinking: enableThinking,
            max_tokens: model.max_length ? model.max_length : 2048,
        };
    }
}

export default Prompt;