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 Story {
    id: string;
    title: string;
    text: string;
    lore: string;
    currentTab: Tab;
    chatMessages: ChatMessage[];
}

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

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

// ─── 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 };

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

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

// ─── 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: '',
                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,
            };
        }
    }
}

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

export interface AppState {
    stories: Story[];
    currentStory: Story | null;
    connection: LLM.Connection | null;
    model: LLM.ModelInfo | null;
    enableThinking: boolean;
    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,
        dispatch,
    }), [state]);

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