import { formatErrorMessage } from "@common/errors";
import { Type, type Static, type TObject } from '@common/typebox';
import { CharacterRole, LocationScale, type AppState, type Character, type Location } from "../contexts/state";
import LLM from "./llm";

const VALID_SCALES = Object.values(LocationScale);
const VALID_ROLES = Object.values(CharacterRole);
const LINES_LIMIT = 7;

export namespace Tools {
    interface Tool<T extends TObject = TObject> {
        description: string;
        parameters: T;
        handler(args: Static<T>, appState: AppState): unknown;
    }

    const tool = <T extends TObject = TObject>(t: Tool<T>): Tool<T> => t;

    const TOOLS = {
        'get_character': tool({
            handler: async (args, appState) => {
                if (!appState.currentStory) {
                    return 'Error: No story selected';
                }
                const character = appState.mergedCharacters.find(c => c.name === args.name.trim());
                if (!character) {
                    return `Error: Character "${args.name.trim()}" not found`;
                }
                let result = `# Character: ${character.name}\n\n`;
                result += `**Role:** ${character.role}\n\n`;
                if (character.nicknames && character.nicknames.length > 0) {
                    result += `**Nicknames:** ${character.nicknames.join(', ')}\n\n`;
                }
                result += `**Short Description:** ${character.shortDescription}\n\n`;
                if (character.description) {
                    result += `**Description:** ${character.description}\n\n`;
                }
                if (character.relations && character.relations.length > 0) {
                    result += `**Relations:**\n`;
                    for (const rel of character.relations) {
                        result += `- ${rel.name}: ${rel.relation}\n`;
                    }
                }
                return result.trim();
            },
            description: 'Get full information about a character by name',
            parameters: Type.Object({
                name: Type.String({ description: "The character's full name" }),
            }),
        }),
        'set_character': tool({
            handler: async (args, appState) => {
                if (!appState.currentStory || !appState.currentWorld) {
                    return 'Error: No story selected';
                }
                const existingCharacter = appState.mergedCharacters.find(c => c.name === args.name.trim());

                if (existingCharacter) {
                    // Edit existing character — find which container it lives in
                    const inStory = appState.currentStory.characters.some(c => c.id === existingCharacter.id);
                    const definedUpdates: Partial<Character> = {};
                    if (args.shortDescription !== undefined) {
                        definedUpdates.shortDescription = args.shortDescription;
                    }
                    if (args.description !== undefined) {
                        definedUpdates.description = args.description;
                    }
                    if (args.nicknames !== undefined) {
                        definedUpdates.nicknames = args.nicknames;
                    }
                    if (args.role !== undefined) {
                        definedUpdates.role = args.role;
                    }
                    if (args.relations !== undefined) {
                        return 'Error: set_character does not support updating relations. Use set_character_relation instead.';
                    }
                    appState.dispatch({
                        type: 'EDIT_CHARACTER',
                        worldId: appState.currentWorld.id,
                        storyId: inStory ? appState.currentStory.id : null,
                        characterId: existingCharacter.id,
                        updates: definedUpdates,
                    });
                    appState.dispatch({ type: 'SET_CURRENT_TAB', tab: 'characters' });
                    return `Character "${args.name.trim()}" updated successfully`;
                } else {
                    // Add new character — add to story level by default
                    if (!args.shortDescription) {
                        return 'Error: shortDescription is required when adding a new character';
                    }
                    if (args.role === undefined) {
                        return 'Error: role is required when adding a new character';
                    }
                    const existingCharacterNames = new Set(appState.mergedCharacters.map(c => c.name));
                    const invalidRelations: string[] = [];
                    for (const rel of args.relations || []) {
                        if (!existingCharacterNames.has(rel.name)) {
                            invalidRelations.push(`${rel.name} (${rel.relation})`);
                        }
                    }
                    appState.dispatch({
                        type: 'ADD_CHARACTER',
                        worldId: appState.currentWorld.id,
                        storyId: appState.currentStory.id,
                        character: {
                            id: crypto.randomUUID(),
                            name: args.name.trim(),
                            role: args.role,
                            nicknames: args.nicknames || [],
                            shortDescription: args.shortDescription.trim(),
                            description: args.description || '',
                            relations: args.relations || [],
                        },
                    });
                    appState.dispatch({ type: 'SET_CURRENT_TAB', tab: 'characters' });
                    let message = `Character "${args.name.trim()}" added successfully`;
                    if (invalidRelations.length > 0) {
                        message += `.\nNote: found invalid relations to non-existent characters: ${invalidRelations.join(', ')}`;
                    }
                    return message;
                }
            },
            description: 'Add or edit a character. If character exists, updates it; otherwise creates new one.',
            parameters: Type.Object({
                name: Type.String({ description: "The character's full name" }),
                shortDescription: Type.Optional(Type.String({ description: 'A brief description of the character (one line). Required on a new character.' })),
                role: Type.Optional(Type.Enum(VALID_ROLES, { description: `The character role in the story. Required when adding a new character.` })),
                nicknames: Type.Optional(Type.Array(Type.String(), { description: 'Optional list of nicknames' })),
                description: Type.Optional(Type.String({ description: 'Optional full character description' })),
                relations: Type.Optional(Type.Array(
                    Type.Object({
                        name: Type.String({ description: 'Related character name' }),
                        relation: Type.String({ description: 'Relationship type' }),
                    }),
                    { description: 'Optional list of relationships with other characters (allowed only when adding a new character)' }
                )),
            }),
        }),
        'set_character_relation': tool({
            handler: async (args, appState) => {
                if (!appState.currentStory || !appState.currentWorld) {
                    return 'Error: No story selected';
                }
                const character = appState.mergedCharacters.find(c => c.name === args.character_name.trim());
                if (!character) {
                    return `Error: Character "${args.character_name.trim()}" not found`;
                }
                const targetCharacter = appState.mergedCharacters.find(c => c.name === args.target_name.trim());
                if (!targetCharacter) {
                    return `Error: Target character "${args.target_name.trim()}" not found, please add it first.`;
                }
                const inStory = appState.currentStory.characters.some(c => c.id === character.id);
                const storyId = inStory ? appState.currentStory.id : null;
                const existingRelationIndex = character.relations.findIndex(r => r.name === args.target_name.trim());
                if (existingRelationIndex !== -1) {
                    appState.dispatch({
                        type: 'EDIT_CHARACTER_RELATION',
                        worldId: appState.currentWorld.id,
                        storyId,
                        characterId: character.id,
                        targetName: args.target_name.trim(),
                        updates: { relation: args.relation.trim() },
                    });
                    appState.dispatch({ type: 'SET_CURRENT_TAB', tab: 'characters' });
                    return `Relation updated: "${args.character_name.trim()}" -> "${args.target_name.trim()}" (${args.relation.trim()})`;
                } else {
                    appState.dispatch({
                        type: 'ADD_CHARACTER_RELATION',
                        worldId: appState.currentWorld.id,
                        storyId,
                        characterId: character.id,
                        relation: {
                            name: args.target_name.trim(),
                            relation: args.relation.trim(),
                        },
                    });
                    appState.dispatch({ type: 'SET_CURRENT_TAB', tab: 'characters' });
                    return `Relation added: "${args.character_name.trim()}" -> "${args.target_name.trim()}" (${args.relation.trim()})`;
                }
            },
            description: 'Add or edit a relationship between two characters. If relation exists, updates it; otherwise creates new one.',
            parameters: Type.Object({
                character_name: Type.String({ description: "The character's full name to add/edit the relation for" }),
                target_name: Type.String({ description: "The target character's full name" }),
                relation: Type.String({ description: 'The relationship type (e.g., friend, enemy, sibling)' }),
            }),
        }),
        'get_location': tool({
            handler: async (args, appState) => {
                if (!appState.currentStory) {
                    return 'Error: No story selected';
                }
                const location = appState.mergedLocations.find(l => l.name === args.name.trim());
                if (!location) {
                    return `Error: Location "${args.name.trim()}" not found`;
                }
                let result = `# Location: ${location.name}\n\n`;
                result += `**Short Description:** ${location.shortDescription}\n\n`;
                if (location.description) {
                    result += `**Description:** ${location.description}\n\n`;
                }
                result += `**Scale:** ${location.scale}\n`;
                return result.trim();
            },
            description: 'Get full information about a location by name',
            parameters: Type.Object({
                name: Type.String({ description: "The location's full name" }),
            }),
        }),
        'set_location': tool({
            handler: async (args, appState) => {
                if (!appState.currentStory || !appState.currentWorld) {
                    return 'Error: No story selected';
                }
                const existingLocation = appState.mergedLocations.find(l => l.name === args.name.trim());
                if (existingLocation) {
                    const inStory = appState.currentStory.locations.some(l => l.id === existingLocation.id);
                    const definedUpdates: Partial<Location> = {};
                    if (args.shortDescription) {
                        definedUpdates.shortDescription = args.shortDescription;
                    }
                    if (args.description) {
                        definedUpdates.description = args.description;
                    }
                    if (args.scale != null) {
                        definedUpdates.scale = args.scale;
                    }
                    appState.dispatch({
                        type: 'EDIT_LOCATION',
                        worldId: appState.currentWorld.id,
                        storyId: inStory ? appState.currentStory.id : null,
                        locationId: existingLocation.id,
                        updates: definedUpdates,
                    });
                    appState.dispatch({ type: 'SET_CURRENT_TAB', tab: 'locations' });
                    return `Location "${args.name.trim()}" updated successfully`;
                } else {
                    if (!args.shortDescription) {
                        return 'Error: shortDescription is required when adding a new location';
                    }
                    if (args.scale == null) {
                        return 'Error: scale is required when adding a new location';
                    }
                    appState.dispatch({
                        type: 'ADD_LOCATION',
                        worldId: appState.currentWorld.id,
                        storyId: appState.currentStory.id,
                        location: {
                            id: crypto.randomUUID(),
                            name: args.name.trim(),
                            shortDescription: args.shortDescription.trim(),
                            description: args.description || '',
                            scale: args.scale,
                        },
                    });
                    appState.dispatch({ type: 'SET_CURRENT_TAB', tab: 'locations' });
                    return `Location "${args.name.trim()}" added successfully`;
                }
            },
            description: 'Add or edit a location. If location exists, updates it; otherwise creates new one.',
            parameters: Type.Object({
                name: Type.String({ description: "The location's full name" }),
                shortDescription: Type.Optional(Type.String({ description: 'A brief description of the location (one line). Required when adding a new location.' })),
                description: Type.Optional(Type.String({ description: 'Optional full location description' })),
                scale: Type.Optional(Type.Enum(VALID_SCALES, {
                    description: `Location scale (enum): ${VALID_SCALES.join(', ')}`,
                })),
            }),
        }),
        'set_lore_entry': tool({
            handler: async (args, appState) => {
                if (!appState.currentStory || !appState.currentWorld) {
                    return 'Error: No story selected';
                }
                const existingEntry = appState.mergedLore.find(e => e.title.toLowerCase() === args.title.trim().toLowerCase());

                if (existingEntry) {
                    const inStory = appState.currentStory.lore.some(e => e.id === existingEntry.id);
                    appState.dispatch({
                        type: 'EDIT_LORE_ENTRY',
                        worldId: appState.currentWorld.id,
                        storyId: inStory ? appState.currentStory.id : null,
                        entryId: existingEntry.id,
                        updates: { text: args.text },
                    });
                    appState.dispatch({ type: 'SET_CURRENT_TAB', tab: 'lore' });
                    return `Lore entry "${existingEntry.title}" updated successfully`;
                } else {
                    appState.dispatch({
                        type: 'ADD_LORE_ENTRY',
                        worldId: appState.currentWorld.id,
                        storyId: appState.currentStory.id,
                        entry: {
                            id: crypto.randomUUID(),
                            title: args.title.trim(),
                            text: args.text,
                        },
                    });
                    appState.dispatch({ type: 'SET_CURRENT_TAB', tab: 'lore' });
                    return `Lore entry "${args.title.trim()}" added successfully`;
                }
            },
            description: 'Add or edit a lore entry. If entry with matching title exists, updates it; otherwise creates new one.',
            parameters: Type.Object({
                title: Type.String({ description: "The lore entry's title" }),
                text: Type.String({ description: 'The lore entry content.' }),
            }),
        }),
        'edit_text': tool({
            handler: async (args, appState) => {
                if (!appState.currentStory || !appState.currentWorld) {
                    return 'Error: No story selected';
                }
                const target = args.target ?? 'story';
                const isScratchpad = target === 'scratchpad';
                const currentText = isScratchpad ? (appState.currentStory.scratchpad ?? '') : appState.currentStory.text;
                const tab = isScratchpad ? 'scratchpad' : 'story';

                const dispatchEdit = (text: string) => appState.dispatch(
                    isScratchpad
                        ? { type: 'EDIT_SCRATCHPAD', worldId: appState.currentWorld!.id, id: appState.currentStory!.id, text }
                        : { type: 'EDIT_STORY', worldId: appState.currentWorld!.id, id: appState.currentStory!.id, text, highlightText: args.new_text }
                );

                // Append mode: when old_text is not provided, append new_text
                if (args.old_text == null) {
                    if (!isScratchpad) {
                        const isAppendToStory = (tc: LLM.ToolCall) => {
                            const jsonArgs = JSON.stringify(tc.function.arguments);
                            return tc.function.name === 'edit_text'
                                && (
                                    !jsonArgs.match(/\\?"old_text\\?"/)
                                    || jsonArgs.match(/\\?"old_text\\?":\\?"\\?"/)
                                )
                                && (
                                    !jsonArgs.match(/\\?"target\\?"/)
                                    || jsonArgs.match(/\\?"target\\?":\\?"story\\?"/)
                                );
                        }


                        // Check that there was a user message since the last edit_text append
                        const messages = appState.currentStory.chatMessages;
                        let hasUserMessageSinceLastAppend = true;
                        for (let i = messages.length - 1; i >= 0; i--) {
                            const m = messages[i];
                            if (m.role === 'user') break;
                            if (m.role === 'assistant' && m.tool_calls?.some(isAppendToStory)) {
                                hasUserMessageSinceLastAppend = false;
                                break;
                            }
                        }
                        if (!hasUserMessageSinceLastAppend) {
                            return `Error: You cannot add new text to the story until user approves your last edit. Stop right there.`;
                        }
                    }

                    const numLines = args.new_text.split('\n').filter(l => l.trim()).length;

                    let cropped = false;
                    if (!isScratchpad && numLines > LINES_LIMIT) {
                        const lines = args.new_text.split('\n');
                        let kept = 0;
                        const croppedLines: string[] = [];
                        for (const line of lines) {
                            if (line.trim()) {
                                if (kept >= LINES_LIMIT) break;
                                kept++;
                            }
                            croppedLines.push(line);
                        }
                        args.new_text = croppedLines.join('\n');
                        cropped = true;
                    }

                    dispatchEdit(currentText + '\n' + args.new_text);
                    appState.dispatch({ type: 'SET_CURRENT_TAB', tab });
                    let message = cropped
                        ? `Added text:\n${args.new_text.split('\n').filter(l => l.trim()).map(l => '> ' + l).join('\n')}\n\nNote: The rest was cropped due to ${LINES_LIMIT} lines limit!`
                        : `Text appended to ${target} successfully.`;

                    message += `\nNote: you can't continue nor add more text until user's approval. Stop any attempts and wait.`

                    return message;
                }

                // Replace mode
                const occurrences = currentText.split(args.old_text).length - 1;
                if (occurrences === 0) {
                    return `Error: old_text not found in ${target}`;
                }
                if (occurrences > 1 && !args.replace_all) {
                    return `Error: old_text appears multiple times in ${target}`;
                }

                dispatchEdit(currentText.replaceAll(args.old_text, args.new_text));
                appState.dispatch({ type: 'SET_CURRENT_TAB', tab });
                return `${target.charAt(0).toUpperCase() + target.slice(1)} edited successfully`;
            },
            description: `Replace or append text in the story or scratchpad. When old_text is omitted, appends new_text to the end. Case-sensitive. When appending to the main story, you can add no more than ${LINES_LIMIT} non-empty lines at once.`,
            parameters: Type.Object({
                new_text: Type.String({ description: 'The new text to replace old_text with, or to append if old_text is omitted' }),
                old_text: Type.Optional(Type.String({ description: 'The text to find and replace. If omitted, new_text will be appended' })),
                replace_all: Type.Optional(Type.Boolean({ description: 'If true, replace all occurrences of old_text' })),
                target: Type.Optional(Type.Enum(['story', 'scratchpad'] as const, { description: 'Which text to edit: "story" (default) or "scratchpad"' })),
            }),
        }),
        'grep': tool({
            handler: async (args, appState) => {
                if (!appState.currentStory) {
                    return 'Error: No story selected';
                }

                const sources: { name: string; content: string }[] = [
                    { name: 'story', content: appState.currentStory.text },
                    ...appState.mergedLore.flatMap(e => [
                        { name: `lore:${e.title}`, content: e.title },
                        { name: `lore:${e.title}`, content: e.text },
                    ]),
                    ...appState.mergedCharacters.flatMap(c => [
                        { name: `character:${c.name}`, content: c.name },
                        { name: `character:${c.name}`, content: c.shortDescription },
                        { name: `character:${c.name}`, content: c.description || '' },
                        ...c.relations.map(rel => (
                            { name: `character:${c.name}`, content: `relation: ${rel.name} (${rel.relation})` }
                        )),
                    ]),
                    ...appState.mergedLocations.flatMap(l => [
                        { name: `location:${l.name}`, content: l.name },
                        { name: `location:${l.name}`, content: l.shortDescription },
                        { name: `location:${l.name}`, content: l.description || '' },
                    ]),
                ];

                const allMatches: { source: string; line: number; content: string }[] = [];
                const pattern = new RegExp(args.pattern, args.case_sensitive ? 'g' : 'gi');

                for (const src of sources) {
                    const lines = src.content.split('\n');
                    for (let i = 0; i < lines.length; i++) {
                        if (pattern.test(lines[i])) {
                            allMatches.push({
                                source: src.name,
                                line: i + 1,
                                content: lines[i].trim(),
                            });
                        }
                    }
                }

                if (allMatches.length === 0) {
                    return `No matches found for pattern: ${args.pattern}`;
                }

                const limit = args.limit ?? 20;
                const truncated = allMatches.length > limit;
                const displayed = truncated ? allMatches.slice(0, limit) : allMatches;

                let result = `Found ${allMatches.length} match(es) for pattern "${args.pattern}":\n`;
                result += displayed.map(m => `[${m.source}] Line ${m.line}: ${m.content}`).join('\n');
                if (truncated) {
                    result += `\n... and ${allMatches.length - limit} more match(es)`;
                }
                return result;
            },
            description: 'Search for a pattern in the story text, lore, characters, and locations (includes world-level data)',
            parameters: Type.Object({
                pattern: Type.String({ description: 'The JS regex pattern to search for' }),
                case_sensitive: Type.Optional(Type.Boolean({ description: 'If true, search is case-sensitive (default: false)' })),
                limit: Type.Optional(Type.Integer({ description: 'Maximum number of matches to return (default: 20)' })),
            }),
        }),
        'generate_image': tool({
            handler: async (args, appState) => {
                if (!appState.connection) {
                    return 'Error: No connection configured';
                }
                if (!appState.imageModel) {
                    return 'Error: No image model configured';
                }
                const { width, height, negative_prompt, sampler_name } = appState.imageGenerationSettings;
                const response = await LLM.generateImage(appState.connection, {
                    model: appState.imageModel.id,
                    prompt: args.prompt,
                    output_format: 'jpeg',
                    image_settings: {
                        width,
                        height,
                        negative_prompt: negative_prompt || undefined,
                        sampler_name: sampler_name || undefined,
                    },
                });
                if ('error' in response) {
                    return `Error: ${response.error}`;
                }
                const b64 = response.data[0]?.b64_json;
                if (!b64) {
                    return 'Error: No image data returned';
                }
                return `![${args.prompt}](data:image/jpeg;base64,${b64})`;
            },
            description: 'Generate an image from a text prompt. Format prompt as tags: masterpiece, best quality, ...',
            parameters: Type.Object({
                prompt: Type.String({ description: 'The image generation prompt' }),
            }),
        }),
    };

    export type ToolName = keyof typeof TOOLS;
    export const isToolName = (name: string): name is ToolName => name in TOOLS;

    export function getTools(only?: ToolName[]): LLM.Tool[] {
        return Object.entries(TOOLS)
            .filter(([key]) => !only || only.includes(key as ToolName))
            .map(([key, tool]) => ({
                type: 'function',
                function: {
                    name: key,
                    description: tool.description,
                    parameters: tool.parameters,
                },
            }));
    }

    export const CHAT_ONLY_TOOLS = ['generate_image'] as const;

    function parseArg(arg: unknown): unknown {
        if (typeof arg !== 'string') return arg;

        try {
            const obj = JSON.parse(arg);
            if (Array.isArray(obj)) {
                return obj.map(parseArg);
            }
            if (typeof obj === 'object') {
                return Object.fromEntries(
                    Object.entries(obj)
                        .map(([key, value]) => [key, parseArg(value)])
                )
            }
            return obj;
        } catch { }
        return arg;
    }

    export async function executeTool(appState: AppState, toolCall: LLM.ToolCall): Promise<string> {
        const { function: fn } = toolCall;
        const args = parseArg(fn.arguments);

        if (!isToolName(fn.name)) {
            return `Unknown tool: ${fn.name}`;
        }

        const tool = TOOLS[fn.name];
        if (!tool) {
            return `Unknown tool: ${fn.name}`;
        }

        const errors = Type.Check(tool.parameters, args);
        if (errors.length > 0) {
            return errors.map(e => e.message).join('\n');
        }

        try {
            const result = await tool.handler(args as any, appState);
            if (typeof result === 'string') {
                return result;
            }
            return JSON.stringify(result);
        } catch (err) {
            return formatErrorMessage(err);
        }
    }
}
