tsgames/src/games/storywriter/contexts/state.tsx

import { createContext } from "preact";
import { useContext, useMemo, useReducer } from "preact/hooks";

import LLM from "../utils/llm";
import { useStoredReducer } from "@common/hooks/useStoredState";

// ─── Types ────────────────────────────────────────────────────────────────────

export type ChatMessage = LLM.ChatMessage & {
    id: string;
}

export type Tab = "story" | "lore" | "characters" | "locations"

export interface Character {
    id: string;
    name: string;
    nicknames: string[];
    shortDescription: string;
    description: string;
    // Meaningful relationships with other characters
    relations: ({
        name: string; // character full name
        relation: string; // daughter/friend/neighbour/etc
    })[];
}

export interface Story {
    id: string;
    title: string;
    text: string;
    lore: string;
    characters: Character[];
    currentTab: Tab;
    chatMessages: ChatMessage[];
}

// ─── State ───────────────────────────────────────────────────────────────────

interface IState {
    stories: Story[];
    currentStoryId: string | null;
    connection: LLM.Connection | null;
    model: LLM.ModelInfo | null;
    enableThinking: boolean;
    bannedTokens: string[];
}

// ─── Actions ─────────────────────────────────────────────────────────────────

type Action =
    | { type: 'CREATE_STORY'; title: string }
    | { type: 'RENAME_STORY'; id: string; title: string }
    | { type: 'EDIT_STORY'; id: string; text: string }
    | { type: 'EDIT_LORE'; id: string; lore: string }
    | { type: 'SET_CURRENT_TAB'; id: string; tab: Tab }
    | { type: 'DELETE_STORY'; id: string }
    | { type: 'SELECT_STORY'; id: string }
    | { type: 'ADD_CHAT_MESSAGE'; storyId: string; message: ChatMessage }
    | { type: 'CLEAR_CHAT'; storyId: string }
    | { type: 'SET_CONNECTION'; connection: LLM.Connection | null }
    | { type: 'SET_MODEL'; model: LLM.ModelInfo | null }
    | { type: 'SET_ENABLE_THINKING'; enable: boolean }
    | { type: 'SET_BANNED_TOKENS'; tokens: string[] }
    | { type: 'ADD_CHARACTER'; storyId: string; character: Character }
    | { type: 'EDIT_CHARACTER'; storyId: string; characterId: string; updates: Partial<Omit<Character, 'relations'>> }
    | { type: 'DELETE_CHARACTER'; storyId: string; characterId: string }
    | { type: 'ADD_CHARACTER_RELATION'; storyId: string; characterId: string; relation: Character['relations'][number] }
    | { type: 'EDIT_CHARACTER_RELATION'; storyId: string; characterId: string; targetName: string; updates: Partial<Character['relations'][number]> }
    | { type: 'DELETE_CHARACTER_RELATION'; storyId: string; characterId: string; targetName: string };

// ─── Initial State ───────────────────────────────────────────────────────────

const DEFAULT_STATE: IState = {
    stories: [],
    currentStoryId: null,
    connection: null,
    model: null,
    enableThinking: false,
    bannedTokens: [],
};

// ─── Reducer ─────────────────────────────────────────────────────────────────

function reducer(state: IState, action: Action): IState {
    switch (action.type) {
        case 'CREATE_STORY': {
            const story: Story = {
                id: crypto.randomUUID(),
                title: action.title,
                text: '',
                lore: '',
                characters: [],
                currentTab: 'story',
                chatMessages: [],
            };
            return {
                ...state,
                stories: [...state.stories, story],
                currentStoryId: story.id,
            };
        }
        case 'RENAME_STORY': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.id ? { ...s, title: action.title } : s
                ),
            };
        }
        case 'EDIT_STORY': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.id ? { ...s, text: action.text } : s
                ),
            };
        }
        case 'EDIT_LORE': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.id ? { ...s, lore: action.lore } : s
                ),
            };
        }
        case 'SET_CURRENT_TAB': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.id ? { ...s, currentTab: action.tab } : s
                ),
            };
        }
        case 'DELETE_STORY': {
            const remaining = state.stories.filter(s => s.id !== action.id);
            const deletingCurrent = state.currentStoryId === action.id;
            return {
                ...state,
                stories: remaining,
                currentStoryId: deletingCurrent ? null : state.currentStoryId,
            };
        }
        case 'SELECT_STORY': {
            return {
                ...state,
                currentStoryId: action.id,
            };
        }
        case 'ADD_CHAT_MESSAGE': {
            return {
                ...state,
                stories: state.stories.map(s => {
                    if (s.id !== action.storyId) return s;
                    const existingIndex = s.chatMessages.findIndex(m => m.id === action.message.id);
                    if (existingIndex !== -1) {
                        // Overwrite existing message with same id
                        const updatedMessages = [...s.chatMessages];
                        updatedMessages[existingIndex] = action.message;
                        return { ...s, chatMessages: updatedMessages };
                    }
                    // Append new message
                    return { ...s, chatMessages: [...s.chatMessages, action.message] };
                }),
            };
        }
        case 'CLEAR_CHAT': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.storyId ? { ...s, chatMessages: [] } : s
                ),
            };
        }
        case 'SET_CONNECTION': {
            return {
                ...state,
                connection: action.connection,
            };
        }
        case 'SET_MODEL': {
            return {
                ...state,
                model: action.model,
            };
        }
        case 'SET_ENABLE_THINKING': {
            return {
                ...state,
                enableThinking: action.enable,
            };
        }
        case 'SET_BANNED_TOKENS': {
            return {
                ...state,
                bannedTokens: action.tokens,
            };
        }
        case 'ADD_CHARACTER': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.storyId
                        ? { ...s, characters: [...s.characters, action.character] }
                        : s
                ),
            };
        }
        case 'EDIT_CHARACTER': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.storyId
                        ? {
                              ...s,
                              characters: s.characters.map(c =>
                                  c.id === action.characterId
                                      ? { ...c, ...action.updates }
                                      : c
                              ),
                          }
                        : s
                ),
            };
        }
        case 'DELETE_CHARACTER': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.storyId
                        ? { ...s, characters: s.characters.filter(c => c.id !== action.characterId) }
                        : s
                ),
            };
        }
        case 'ADD_CHARACTER_RELATION': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.storyId
                        ? {
                              ...s,
                              characters: s.characters.map(c =>
                                  c.id === action.characterId
                                      ? { ...c, relations: [...c.relations, action.relation] }
                                      : c
                              ),
                          }
                        : s
                ),
            };
        }
        case 'EDIT_CHARACTER_RELATION': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.storyId
                        ? {
                              ...s,
                              characters: s.characters.map(c =>
                                  c.id === action.characterId
                                      ? {
                                            ...c,
                                            relations: c.relations.map(r =>
                                                r.name === action.targetName
                                                    ? { ...r, ...action.updates }
                                                    : r
                                            ),
                                        }
                                      : c
                              ),
                          }
                        : s
                ),
            };
        }
        case 'DELETE_CHARACTER_RELATION': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.storyId
                        ? {
                              ...s,
                              characters: s.characters.map(c =>
                                  c.id === action.characterId
                                      ? { ...c, relations: c.relations.filter(r => r.name !== action.targetName) }
                                      : c
                              ),
                          }
                        : s
                ),
            };
        }
    }
}

// ─── Context ─────────────────────────────────────────────────────────────────

export interface AppState {
    stories: Story[];
    currentStory: Story | null;
    connection: LLM.Connection | null;
    model: LLM.ModelInfo | null;
    enableThinking: boolean;
    bannedTokens: string[];
    dispatch: (action: Action) => void;
}

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

export const useAppState = () => useContext(StateContext);

// ─── Provider ────────────────────────────────────────────────────────────────

export const StateContextProvider = ({ children }: { children?: any }) => {
    const [state, dispatch] = useStoredReducer('storywriter.state', reducer, DEFAULT_STATE);

    const value = useMemo<AppState>(() => ({
        stories: state.stories,
        currentStory: state.stories.find(s => s.id === state.currentStoryId) ?? null,
        connection: state.connection,
        model: state.model,
        enableThinking: state.enableThinking,
        bannedTokens: state.bannedTokens ?? [],
        dispatch,
    }), [state]);

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