import { createContext } from "preact";
import { useCallback, useEffect, useMemo, useReducer, useState, type Dispatch, type StateUpdater } from "preact/hooks";
import { MessageTools, type IMessage } from "../tools/messages";
import { useInputState } from "@common/hooks/useInputState";
import { type IConnection } from "../tools/connection";
import { loadObject, saveObject } from "../../../common/storage";
import { useInputCallback } from "@common/hooks/useInputCallback";
import { callUpdater, throttle } from "@common/utils";
import { Huggingface } from "../tools/huggingface";

interface IStory {
    lore: string;
    messages: IMessage[];
}

interface StoriesState {
    currentStory: string;
    stories: Record<string, IStory>;
    lore: string;
    messages: IMessage[];
}

interface StoryActionSetCurrent {
    currentStory: StateUpdater<string>;
}

interface StoryActionSetMessages {
    messages: StateUpdater<IMessage[]>;
}

interface StoryActionSetLore {
    lore: StateUpdater<string>;
}

interface StoryActionWithId {
    action: 'create' | 'delete';
    id: string;
    fromId?: string;
}

type StoryAction = StoryActionSetCurrent | StoryActionSetLore | StoryActionSetMessages | StoryActionWithId;

interface IContext {
    connection: IConnection;
    input: string;
    systemPrompt: string;
    userPrompt: string;
    summarizePrompt: string;
    summaryEnabled: boolean;
    bannedWords: string[];
    totalSpentKudos: number;
    stories: Record<string, IStory>;
    currentStory: string;
    // 
    triggerNext: boolean;
    continueLast: boolean;
}

interface IComputableContext {
    lore: string;
    messages: IMessage[];
}

interface IActions {
    setConnection: (connection: IConnection) => void;
    setInput: (url: string | Event) => void;
    setInstruct: (template: string | Event) => void;
    setLore: (lore: string | Event) => void;
    setSystemPrompt: (prompt: string | Event) => void;
    setUserPrompt: (prompt: string | Event) => void;
    setSummarizePrompt: (prompt: string | Event) => void;
    setBannedWords: (words: string[]) => void;
    setSummaryEnabled: (summaryEnabled: boolean) => void;
    setTotalSpentKudos: Dispatch<StateUpdater<number>>;

    setTriggerNext: (triggerNext: boolean) => void;
    setContinueLast: (continueLast: boolean) => void;

    setMessages: (messages: IMessage[]) => void;
    addMessage: (content: string, role: IMessage['role'], triggerNext?: boolean) => void;
    editMessage: (index: number, content: string, cost: number) => void;
    editSummary: (index: number, summary: string, cost: number) => void;
    deleteMessage: (index: number) => void;
    setCurrentSwipe: (index: number, swipe: number) => void;
    addSwipe: (index: number, content: string) => void;

    continueMessage: (continueLast?: boolean) => void;

    setCurrentStory: (id: string) => void;
    createStory: (id: string, fromId?: string) => void;
    deleteStory: (id: string) => void;
}

const SAVE_KEY = 'ai_game_save_state';
export const DEFAULT_STORY = 'default';

const INSTRUCT_MISTRAL = Huggingface.formatTemplate(`{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{% set loop_messages = messages[1:] %}{% else %}{% set loop_messages = messages %}{% endif %}{% for message in loop_messages %}{% if message['role'] == 'user' %}{% if loop.first and system_message is defined %}{{ ' [INST] ' + system_message + '\n\n' + message['content'] + ' [/INST]' }}{% else %}{{ ' [INST] ' + message['content'] + ' [/INST]' }}{% endif %}{% elif message['role'] == 'assistant' %}{{ ' ' + message['content'] + '</s>' }}{% endif %}{% endfor %}`);

const INSTRUCT_CHATML = Huggingface.formatTemplate(`{% for message in messages %}{{ '<|im_start|>' + message['role'] + '\n\n' + message['content'] + '<|im_end|>' + '\n' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n\n' }}{% endif %}`);

const INSTRUCT_LLAMA = Huggingface.formatTemplate(`{% for message in messages %}{{ '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}`);

const INSTRUCT_ALPACA = Huggingface.formatTemplate(`{% for message in messages %}{% if message['role'] == 'system' and message['content'] %}{{ message['content'] + '\n\n' }}{% elif message['role'] == 'user' %}{{ '### Instruction:\n\n' + message['content'] + '\n\n' }}{% elif message['role'] == 'assistant' %}{{ '### Response:\n\n' + message['content'] + '\n\n' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '### Response:\n\n' }}{% endif %}`);

const INSTRUCT_METHARME = Huggingface.formatTemplate(`{% for message in messages %}{% if message['role'] == 'system' and message['content'] %}{{ '<|system|>' + message['content'] }}{% elif message['role'] == 'user' %}{{ '<|user|>' + message['content'] }}{% elif message['role'] == 'assistant' %}{{'<|model|>' + message['content'] }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|model|>' }}{% endif %}`);

const INSTRUCT_GEMMA = Huggingface.formatTemplate(`{% for message in messages %}{% if (message['role'] == 'assistant') %}{% set role = 'model' %}{% else %}{% set role = message['role'] %}{% endif %}{{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}{% endfor %}{% if add_generation_prompt %}{{ '<start_of_turn>model\n' }}{% endif %}`);

export const INSTRUCTS = {
    'Mistral': INSTRUCT_MISTRAL,
    'ChatML': INSTRUCT_CHATML,
    'LLama': INSTRUCT_LLAMA,
    'Alpaca': INSTRUCT_ALPACA,
    'Metharme': INSTRUCT_METHARME,
    'Gemma': INSTRUCT_GEMMA,
}

const DEFAULT_CONTEXT: IContext = {
    connection: {
        type: 'kobold',
        url: 'http://localhost:5001',
        instruct: INSTRUCT_MISTRAL,
    },
    input: '',
    systemPrompt: 'You are a creative writer. Write a story based on the world description below. Story should be adult and mature; and could include swearing, violence and unfairness. Portray characters realistically and stay in the lore.',
    stories: {},
    currentStory: DEFAULT_STORY,
    userPrompt: Huggingface.formatTemplate(`{% if isStart -%}
Write a novel using information above as a reference.
{%- else -%}
Continue the story forward.
{%- endif %}

{% if prompt -%}
This is the description of what should happen next in your answer: {{ prompt | trim }}
{% endif %}
Remember that this story should be infinite and go forever.
Make sure to follow the world description and rules exactly. Avoid cliffhangers and pauses, be creative.`),
    summarizePrompt: 'Summarize following text in one paragraph:\n\n{{ message }}\n\nAnswer with shortened text only.',
    summaryEnabled: true,
    bannedWords: [],
    totalSpentKudos: 0,
    triggerNext: false,
    continueLast: false,
};

const EMPTY_STORY: IStory = {
    lore: '',
    messages: [],
};

const saveContext = throttle(async (context: IContext & IComputableContext) => {
    const contextToSave: Partial<IContext & IComputableContext> = { ...context };
    delete contextToSave.triggerNext;
    delete contextToSave.continueLast;
    delete contextToSave.lore;
    delete contextToSave.messages;

    return saveObject(SAVE_KEY, contextToSave);
}, 1000, true);

export type IStateContext = IContext & IActions & IComputableContext;

export const StateContext = createContext<IStateContext>({} as IStateContext);

const loadedContext = await loadObject(SAVE_KEY, DEFAULT_CONTEXT);

const storyReducer = (state: StoriesState, action: StoryAction): StoriesState => {
    let { currentStory, stories } = state;

    if ('id' in action) {
        switch (action.action) {
            case 'create':
                stories[action.id] = {
                    ...EMPTY_STORY,
                    lore: (action.fromId && stories[action.fromId]?.lore) ?? EMPTY_STORY.lore,
                };
                break;
            case 'delete':
                if (action.id !== DEFAULT_STORY) {
                    delete stories[action.id];
                    if (action.id === currentStory) {
                        currentStory = DEFAULT_STORY;
                    }
                }
                break;
        }
    } else if ('currentStory' in action) {
        currentStory = callUpdater(action.currentStory, currentStory);
    } else if ('messages' in action) {
        stories[currentStory].messages = callUpdater(action.messages, stories[currentStory]?.messages ?? []);
    } else if ('lore' in action) {
        stories[currentStory].lore = callUpdater(action.lore, stories[currentStory]?.lore ?? '');
    }

    return {
        currentStory,
        stories,
        lore: stories[currentStory].lore ?? '',
        messages: stories[currentStory].messages ?? [],
    };
};

export const StateContextProvider = ({ children }: { children?: any }) => {
    const [connection, setConnection] = useState<IConnection>(loadedContext.connection);
    const [input, setInput] = useInputState(loadedContext.input);
    const [systemPrompt, setSystemPrompt] = useInputState(loadedContext.systemPrompt);
    const [userPrompt, setUserPrompt] = useInputState(loadedContext.userPrompt);
    const [summarizePrompt, setSummarizePrompt] = useInputState(loadedContext.summarizePrompt);
    const [bannedWords, setBannedWords] = useState<string[]>(loadedContext.bannedWords);
    const [summaryEnabled, setSummaryEnabled] = useState(loadedContext.summaryEnabled);
    const [totalSpentKudos, setTotalSpentKudos] = useState(loadedContext.totalSpentKudos);

    const [storiesState, storyDispatch] = useReducer(storyReducer, {
        stories: loadedContext.stories,
        currentStory: loadedContext.currentStory,
        lore: loadedContext.stories[loadedContext.currentStory]?.lore ?? '',
        messages: loadedContext.stories[loadedContext.currentStory]?.messages ?? [],
    });

    const [triggerNext, setTriggerNext] = useState(false);
    const [continueLast, setContinueLast] = useState(false);
    const [instruct, setInstruct] = useInputState(connection.instruct);

    useEffect(() => setConnection({ ...connection, instruct }), [instruct]);

    const setLore = useInputCallback((lore) => {
        storyDispatch({ lore });
    }, []);

    const setMessages = useCallback((messages: StateUpdater<IMessage[]>) => {
        storyDispatch({ messages });
    }, []);

    const setCurrentStory = useCallback((currentStory: StateUpdater<string>) => {
        storyDispatch({ currentStory });
    }, []);

    const actions: IActions = useMemo(() => ({
        setConnection,
        setInput,
        setInstruct,
        setSystemPrompt,
        setUserPrompt,
        setSummarizePrompt,
        setLore,
        setSummaryEnabled,

        setTriggerNext,
        setContinueLast,
        setTotalSpentKudos,
        setCurrentStory,

        setBannedWords: (words) => setBannedWords(words.slice()),

        setMessages: (newMessages) => setMessages(newMessages.slice()),
        addMessage: (content, role, triggerNext = false) => {
            setMessages(messages => [
                ...messages,
                MessageTools.create(content, role),
            ]);
            setTriggerNext(triggerNext);
        },
        editMessage: (index, content, cost) => {
            setMessages(messages => MessageTools.updateSwipe(messages, index, { content }, cost));
        },
        editSummary: (index, summary, cost) => {
            setMessages(messages => MessageTools.updateSwipe(messages, index, { summary }, cost));
        },
        deleteMessage: (index) => setMessages(messages =>
            messages.filter((_, i) => i !== index)
        ),
        setCurrentSwipe: (index, currentSwipe) => {
            let shouldTrigger = false;

            setMessages(messages =>
                messages.map(
                    (message, i) => {
                        if (i === index) {
                            const swipes = message.swipes.slice();
                            const latestSwipe = swipes.at(-1);
                            if (currentSwipe >= swipes.length) {
                                if (latestSwipe && latestSwipe.content.length > 0) {
                                    currentSwipe = swipes.length;
                                    swipes.push({ content: '', cost: 0 });
                                } else {
                                    currentSwipe = swipes.length - 1;
                                }
                                shouldTrigger = true;
                            } else while (currentSwipe < 0) {
                                currentSwipe += swipes.length;
                            }
                            return {
                                ...message, swipes, currentSwipe
                            };
                        } else {
                            return message;
                        }
                    }
                )
            );

            setTriggerNext(shouldTrigger);
        },
        addSwipe: (index, content) => setMessages(messages =>
            messages.map(
                (message, i) => {
                    if (i === index) {
                        const swipes = [...message.swipes, { content, cost: 0 }];

                        return {
                            ...message,
                            swipes,
                            currentSwipe: swipes.length - 1,
                        };
                    } else {
                        return message;
                    }
                }
            )
        ),
        continueMessage: (c = false) => {
            setTriggerNext(true);
            setContinueLast(c);
        },
        createStory: (id: string, fromId?: string) => {
            storyDispatch({ id, action: 'create', fromId });
        },
        deleteStory: (id: string) => {
            storyDispatch({ id, action: 'delete' });
        }
    }), []);

    const rawContext: IContext & IComputableContext = {
        connection,
        input,
        systemPrompt,
        userPrompt,
        summarizePrompt,
        summaryEnabled,
        bannedWords,
        totalSpentKudos,
        ...storiesState,
        //
        triggerNext,
        continueLast,
    };

    const context = useMemo(() => rawContext, Object.values(rawContext));

    useEffect(() => {
        saveContext(context);
    }, [context]);

    const value = useMemo(() => ({ ...context, ...actions }), [context, actions])

    return (
        <StateContext.Provider value={value}>
            {children}
        </StateContext.Provider>
    );
}