import LLM from "./llm";
import Chapters from "./chapters";
import { type AppState, CharacterRole, type ChatMessage } from "../contexts/state";
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 { currentStory } = state;
        if (!currentStory || !currentStory.characters?.length) {
            return '';
        }

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

        // Sort characters by importance (protagonist first, cameo last)
        const sortedCharacters = [...currentStory.characters].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 { currentStory } = state;
        if (!currentStory?.lore?.length) {
            return '';
        }

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

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

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

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

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

        for (const location of currentStory.locations) {
            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 formatSystemPrompt(state: AppState, storyTokenBudget: number = 0): string {
        const { currentStory } = state;
        if (!currentStory) {
            return state.systemInstruction;
        }

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

        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: ChatMessage[] = []): LLM.ChatCompletionRequest | null {
        const { currentStory, model, enableThinking } = state;

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

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

        const messages: ChatMessage[] = [
            { 
                id: crypto.randomUUID(),
                role: 'system',
                content: formatSystemPrompt(state, storyTokenBudget),
             },
            ...currentStory.chatMessages,
        ];

        const presentMessages = new Set(messages.map(m => m.id));
        for (const newMessage of newMessages) {
            if (!presentMessages.has(newMessage.id)) {
                messages.push(newMessage);
            }
        }

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

export default Prompt;
