import { createContext } from "preact";
import { useContext, useMemo } from "preact/hooks";
import { useStoredReducer } from "@common/hooks/useStored";

import LLM from "../utils/llm";
import Chapters from "../utils/chapters";

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

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

export type Tab = "story" | "lore" | "characters" | "locations" | "chapters" | "scratchpad" | "prompt" | "menu" | "system" | "chat";

export enum CharacterRole {
    Protagonist = 'protagonist',
    Antagonist = 'antagonist',
    Main = 'main',
    Secondary = 'secondary',
    Supporting = 'supporting',
    Minor = 'minor',
    Cameo = 'cameo',
}

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

export enum LocationScale {
    Room = 'room',
    House = 'house',
    Street = 'street',
    Village = 'village',
    City = 'city',
    Region = 'region',
    Country = 'country',
    Continent = 'continent',
    World = 'world',
    Universe = 'universe',
}

export interface Location {
    id: string;
    name: string;
    shortDescription: string;
    description: string;
    scale: LocationScale;
}

export interface LoreEntry {
    id: string;
    title: string;
    text: string;
}

export interface Story {
    id: string;
    title: string;
    text: string;
    scratchpad: string;
    lore: LoreEntry[];
    characters: Character[];
    locations: Location[];
    chatMessages: ChatMessage[];
    chapters: Chapters.Chapter[];
    lastEditedText?: string;
}

export interface World {
    id: string;
    title: string;
    chatOnly: boolean;
    lore: LoreEntry[];
    characters: Character[];
    locations: Location[];
    stories: Story[];
    systemInstructionOverride?: string;
}

// ─── Type Guards ──────────────────────────────────────────────────────────────

export function isWorld(obj: unknown): obj is World {
    if (typeof obj !== 'object' || obj === null) return false;
    const w = obj as Record<string, unknown>;
    return (
        typeof w.id === 'string' &&
        typeof w.title === 'string' &&
        typeof w.chatOnly === 'boolean' &&
        Array.isArray(w.lore) &&
        Array.isArray(w.characters) &&
        Array.isArray(w.locations) &&
        Array.isArray(w.stories)
    );
}

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

interface IState {
    worlds: World[];
    currentWorldId: string | null;
    currentStoryId: string | null;
    currentTab: Tab;
    chatOpen: boolean;
    connection: LLM.Connection | null;
    model: LLM.ModelInfo | null;
    enableThinking: boolean;
    bannedTokens: string[];
    systemInstruction: string;
}

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

type Action =
    // World actions
    | { type: 'CREATE_WORLD'; title: string; chatOnly?: boolean }
    | { type: 'RENAME_WORLD'; worldId: string; title: string }
    | { type: 'DELETE_WORLD'; worldId: string }
    | { type: 'SELECT_WORLD'; worldId: string }
    // Story actions
    | { type: 'CREATE_STORY'; worldId: string }
    | { type: 'RENAME_STORY'; worldId: string; id: string; title: string }
    | { type: 'EDIT_STORY'; worldId: string; id: string; text: string; highlightText?: string }
    | { type: 'EDIT_SCRATCHPAD'; worldId: string; id: string; text: string }
    | { type: 'DELETE_STORY'; worldId: string; id: string }
    | { type: 'SELECT_STORY'; worldId: string; id: string }
    | { type: 'DUPLICATE_STORY'; worldId: string; id: string }
    // Story lore
    | { type: 'ADD_LORE_ENTRY'; worldId: string; storyId: string | null; entry: LoreEntry }
    | { type: 'EDIT_LORE_ENTRY'; worldId: string; storyId: string | null; entryId: string; updates: Partial<LoreEntry> }
    | { type: 'DELETE_LORE_ENTRY'; worldId: string; storyId: string | null; entryId: string }
    | { type: 'REORDER_LORE_ENTRIES'; worldId: string; storyId: string | null; entryIds: string[] }
    // Settings
    | { type: 'SET_SYSTEM_INSTRUCTION'; systemInstruction: string }
    | { type: 'SET_WORLD_SYSTEM_INSTRUCTION_OVERRIDE'; worldId: string; systemInstructionOverride: string | undefined }
    | { type: 'SET_CURRENT_TAB'; tab: Tab }
    | { type: 'SET_CHAT_OPEN'; open: boolean }
    // Chat
    | { type: 'ADD_CHAT_MESSAGE'; worldId: string; storyId: string; message: ChatMessage }
    | { type: 'CLEAR_CHAT'; worldId: string; storyId: string }
    | { type: 'DELETE_CHAT_MESSAGES_FROM'; worldId: string; storyId: string; messageId: string }
    | { type: 'EDIT_CHAT_MESSAGE'; worldId: string; storyId: string; messageId: string; content: string }
    // Connection
    | { 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[] }
    // Characters
    | { type: 'ADD_CHARACTER'; worldId: string; storyId: string | null; character: Character }
    | { type: 'EDIT_CHARACTER'; worldId: string; storyId: string | null; characterId: string; updates: Partial<Omit<Character, 'id' | 'name' | 'relations'>> }
    | { type: 'DELETE_CHARACTER'; worldId: string; storyId: string | null; characterId: string }
    | { type: 'ADD_CHARACTER_RELATION'; worldId: string; storyId: string | null; characterId: string; relation: Character['relations'][number] }
    | { type: 'EDIT_CHARACTER_RELATION'; worldId: string; storyId: string | null; characterId: string; targetName: string; updates: Partial<Character['relations'][number]> }
    | { type: 'DELETE_CHARACTER_RELATION'; worldId: string; storyId: string | null; characterId: string; targetName: string }
    // Locations
    | { type: 'ADD_LOCATION'; worldId: string; storyId: string | null; location: Location }
    | { type: 'EDIT_LOCATION'; worldId: string; storyId: string | null; locationId: string; updates: Partial<Location> }
    | { type: 'DELETE_LOCATION'; worldId: string; storyId: string | null; locationId: string }
    // Chapters
    | { type: 'STORE_CHAPTER_SUMMARY'; worldId: string; storyId: string; header: string; hash: Chapters.Hash; summary: string }
    | { type: 'CLEAN_CHAPTER_SUMMARIES'; worldId: string; storyId: string; validHashes: Record<string, Chapters.Hash[]> }
    // Import/Export
    | { type: 'IMPORT_WORLD'; world: World };

// ─── Helpers ─────────────────────────────────────────────────────────────────

function updateWorld(state: IState, worldId: string, updater: (w: World) => World): IState {
    return {
        ...state,
        worlds: state.worlds.map(w => w.id === worldId ? updater(w) : w),
    };
}

function updateStory(state: IState, worldId: string, storyId: string, updater: (s: Story) => Story): IState {
    return updateWorld(state, worldId, w => ({
        ...w,
        stories: w.stories.map(s => s.id === storyId ? updater(s) : s),
    }));
}


function updateLoreContainer(
    state: IState,
    worldId: string,
    storyId: string | null,
    updater: (c: { lore: LoreEntry[]; characters: Character[]; locations: Location[] }) => Partial<{ lore: LoreEntry[]; characters: Character[]; locations: Location[] }>
): IState {
    if (storyId) {
        return updateStory(state, worldId, storyId, s => ({ ...s, ...updater(s) }));
    }
    return updateWorld(state, worldId, w => ({ ...w, ...updater(w) }));
}

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

const DEFAULT_STATE: IState = {
    worlds: [],
    currentWorldId: null,
    currentStoryId: null,
    currentTab: 'menu',
    chatOpen: false,
    connection: null,
    model: null,
    enableThinking: false,
    bannedTokens: [],
    systemInstruction: `You are a creative writing assistant. Help the user develop their story by writing engaging content, maintaining consistency with the established characters, settings, and plot. Follow the user's instructions while staying true to the story's tone and style.

Write using markdown to highlight special parts.
Supported markdown subset:
 - *italic*
 - **bold**
 - "quotes"
 - \`monospace\`
 - # Header 1
 - ## Header 2
 - ### Header 3
 - > blockquote
 - Tables
 - Ordered lists
 - Unordered lists (only with \`- \` markers)
 - Only top-level lists (no nesting)

Show the chapters with \`# Chapter\` headers.
Add important details not yet ready to be included in the story to the scratchpad: character motivations, hidden plot points, etc.
You **must** use \`edit_text\` tool to write to the story.
Keep the reports after editing concise to save token budget, just say "Done" with minimal to no thinking.
The most actual state of the story is provided below, use it as ground truth.`,
};

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

function reducer(state: IState, action: Action): IState {
    switch (action.type) {
        case 'CREATE_WORLD': {
            const world: World = {
                id: crypto.randomUUID(),
                title: action.title,
                chatOnly: action.chatOnly ?? false,
                lore: [],
                characters: [],
                locations: [],
                stories: [],
            };
            return {
                ...state,
                worlds: [...state.worlds, world],
                currentWorldId: world.id,
                currentStoryId: null,
                currentTab: 'menu',
            };
        }
        case 'RENAME_WORLD': {
            return updateWorld(state, action.worldId, w => ({ ...w, title: action.title }));
        }
        case 'DELETE_WORLD': {
            const remaining = state.worlds.filter(w => w.id !== action.worldId);
            const deletingCurrent = state.currentWorldId === action.worldId;
            return {
                ...state,
                worlds: remaining,
                currentWorldId: deletingCurrent ? null : state.currentWorldId,
                currentStoryId: deletingCurrent ? null : state.currentStoryId,
            };
        }
        case 'SELECT_WORLD': {
            return {
                ...state,
                currentWorldId: action.worldId,
                currentStoryId: null,
                currentTab: 'menu',
            };
        }
        case 'CREATE_STORY': {
            const story: Story = {
                id: crypto.randomUUID(),
                title: '',
                text: '',
                scratchpad: '',
                lore: [],
                characters: [],
                locations: [],
                chatMessages: [],
                chapters: [],
            };
            const world = state.worlds.find(w => w.id === action.worldId);
            story.title = world?.chatOnly ? 'New chat' : 'New Story';
            return {
                ...updateWorld(state, action.worldId, w => ({ ...w, stories: [...w.stories, story] })),
                currentWorldId: action.worldId,
                currentStoryId: story.id,
                currentTab: 'menu',
            };
        }
        case 'RENAME_STORY': {
            return updateStory(state, action.worldId, action.id, s => ({ ...s, title: action.title }));
        }
        case 'EDIT_STORY': {
            return updateStory(state, action.worldId, action.id, s => ({
                ...s,
                text: action.text,
                lastEditedText: action.highlightText,
            }));
        }
        case 'EDIT_SCRATCHPAD': {
            return updateStory(state, action.worldId, action.id, s => ({ ...s, scratchpad: action.text }));
        }
        case 'DELETE_STORY': {
            const deletingCurrent = state.currentStoryId === action.id;
            return {
                ...updateWorld(state, action.worldId, w => ({
                    ...w,
                    stories: w.stories.filter(s => s.id !== action.id),
                })),
                currentStoryId: deletingCurrent ? null : state.currentStoryId,
            };
        }
        case 'SELECT_STORY': {
            const world = state.worlds.find(w => w.id === action.worldId);
            return {
                ...state,
                currentWorldId: action.worldId,
                currentStoryId: action.id,
                currentTab: world?.chatOnly ? 'chat' : 'story',
            };
        }
        case 'DUPLICATE_STORY': {
            const world = state.worlds.find(w => w.id === action.worldId);
            const original = world?.stories.find(s => s.id === action.id);
            if (!original) return state;
            const firstMessage = original.chatMessages[0];
            const chatMessages = world?.chatOnly && firstMessage && firstMessage.role === 'assistant' ? [firstMessage] : [];
            const newStory: Story = {
                id: crypto.randomUUID(),
                title: `${original.title} (Copy)`,
                text: '',
                scratchpad: '',
                lore: [...original.lore],
                characters: [...original.characters],
                locations: [...original.locations],
                chatMessages,
                chapters: [],
            };
            return {
                ...updateWorld(state, action.worldId, w => ({ ...w, stories: [...w.stories, newStory] })),
                currentStoryId: newStory.id,
                currentTab: 'menu',
            };
        }
        case 'ADD_LORE_ENTRY': {
            return updateLoreContainer(state, action.worldId, action.storyId, c => ({
                lore: [...c.lore, action.entry],
            }));
        }
        case 'EDIT_LORE_ENTRY': {
            return updateLoreContainer(state, action.worldId, action.storyId, c => ({
                lore: c.lore.map(e => e.id === action.entryId ? { ...e, ...action.updates } : e),
            }));
        }
        case 'DELETE_LORE_ENTRY': {
            return updateLoreContainer(state, action.worldId, action.storyId, c => ({
                lore: c.lore.filter(e => e.id !== action.entryId),
            }));
        }
        case 'REORDER_LORE_ENTRIES': {
            return updateLoreContainer(state, action.worldId, action.storyId, c => {
                const entryMap = new Map(c.lore.map(e => [e.id, e]));
                const reordered = action.entryIds
                    .map(id => entryMap.get(id))
                    .filter((e): e is LoreEntry => e !== undefined);
                for (const entry of c.lore) {
                    if (!action.entryIds.includes(entry.id)) reordered.push(entry);
                }
                return { lore: reordered };
            });
        }
        case 'SET_SYSTEM_INSTRUCTION': {
            return { ...state, systemInstruction: action.systemInstruction };
        }
        case 'SET_WORLD_SYSTEM_INSTRUCTION_OVERRIDE': {
            return updateWorld(state, action.worldId, w => ({ ...w, systemInstructionOverride: action.systemInstructionOverride }));
        }
        case 'SET_CURRENT_TAB': {
            return { ...state, currentTab: action.tab };
        }
        case 'SET_CHAT_OPEN': {
            return { ...state, chatOpen: action.open };
        }
        case 'ADD_CHAT_MESSAGE': {
            return updateStory(state, action.worldId, action.storyId, s => {
                const existingIndex = s.chatMessages.findIndex(m => m.id === action.message.id);
                if (existingIndex !== -1) {
                    const updatedMessages = [...s.chatMessages];
                    updatedMessages[existingIndex] = action.message;
                    return { ...s, chatMessages: updatedMessages };
                }
                return { ...s, chatMessages: [...s.chatMessages, action.message] };
            });
        }
        case 'CLEAR_CHAT': {
            return updateStory(state, action.worldId, action.storyId, s => ({ ...s, chatMessages: [] }));
        }
        case 'DELETE_CHAT_MESSAGES_FROM': {
            return updateStory(state, action.worldId, action.storyId, s => {
                const messageIndex = s.chatMessages.findIndex(m => m.id === action.messageId);
                if (messageIndex === -1) return s;
                return { ...s, chatMessages: s.chatMessages.slice(0, messageIndex) };
            });
        }
        case 'EDIT_CHAT_MESSAGE': {
            return updateStory(state, action.worldId, action.storyId, s => ({
                ...s,
                chatMessages: s.chatMessages.map(m =>
                    m.id === action.messageId ? { ...m, content: action.content } : m
                ),
            }));
        }
        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 updateLoreContainer(state, action.worldId, action.storyId, c => ({
                characters: [...c.characters, action.character],
            }));
        }
        case 'EDIT_CHARACTER': {
            return updateLoreContainer(state, action.worldId, action.storyId, c => ({
                characters: c.characters.map(ch =>
                    ch.id === action.characterId ? { ...ch, ...action.updates } : ch
                ),
            }));
        }
        case 'DELETE_CHARACTER': {
            return updateLoreContainer(state, action.worldId, action.storyId, c => {
                const deleted = c.characters.find(ch => ch.id === action.characterId);
                if (!deleted) return {};
                return {
                    characters: c.characters
                        .filter(ch => ch.id !== action.characterId)
                        .map(ch => ({
                            ...ch,
                            relations: ch.relations.filter(r => r.name !== deleted.name),
                        })),
                };
            });
        }
        case 'ADD_CHARACTER_RELATION': {
            return updateLoreContainer(state, action.worldId, action.storyId, c => ({
                characters: c.characters.map(ch =>
                    ch.id === action.characterId
                        ? { ...ch, relations: [...ch.relations, action.relation] }
                        : ch
                ),
            }));
        }
        case 'EDIT_CHARACTER_RELATION': {
            return updateLoreContainer(state, action.worldId, action.storyId, c => ({
                characters: c.characters.map(ch =>
                    ch.id === action.characterId
                        ? {
                            ...ch,
                            relations: ch.relations.map(r =>
                                r.name === action.targetName ? { ...r, ...action.updates } : r
                            ),
                        }
                        : ch
                ),
            }));
        }
        case 'DELETE_CHARACTER_RELATION': {
            return updateLoreContainer(state, action.worldId, action.storyId, c => ({
                characters: c.characters.map(ch =>
                    ch.id === action.characterId
                        ? { ...ch, relations: ch.relations.filter(r => r.name !== action.targetName) }
                        : ch
                ),
            }));
        }
        case 'ADD_LOCATION': {
            return updateLoreContainer(state, action.worldId, action.storyId, c => ({
                locations: [...c.locations, action.location],
            }));
        }
        case 'EDIT_LOCATION': {
            return updateLoreContainer(state, action.worldId, action.storyId, c => ({
                locations: c.locations.map(l =>
                    l.id === action.locationId ? { ...l, ...action.updates } : l
                ),
            }));
        }
        case 'DELETE_LOCATION': {
            return updateLoreContainer(state, action.worldId, action.storyId, c => ({
                locations: c.locations.filter(l => l.id !== action.locationId),
            }));
        }
        case 'STORE_CHAPTER_SUMMARY': {
            return updateStory(state, action.worldId, action.storyId, s => {
                const chapters = s.chapters ?? [];
                const existing = chapters.find(c => c.header === action.header);
                const updated = existing
                    ? Chapters.storeSummary(existing, action.hash, action.summary)
                    : Chapters.storeSummary(Chapters.emptyChapter(action.header), action.hash, action.summary);
                return {
                    ...s,
                    chapters: existing
                        ? chapters.map(c => c.header === action.header ? updated : c)
                        : [...chapters, updated],
                };
            });
        }
        case 'CLEAN_CHAPTER_SUMMARIES': {
            return updateStory(state, action.worldId, action.storyId, s => {
                const chapters = (s.chapters ?? [])
                    .filter(c => action.validHashes[c.header] !== undefined)
                    .map(c => {
                        const valid = new Set(action.validHashes[c.header]);
                        const summaryCache = Object.fromEntries(
                            Object.entries(c.summaryCache).filter(([hash]) => valid.has(hash))
                        );
                        return { ...c, summaryCache };
                    });
                return { ...s, chapters };
            });
        }
        case 'IMPORT_WORLD': {
            const exists = state.worlds.some(w => w.id === action.world.id);
            const world = exists ? { ...action.world, id: crypto.randomUUID() } : action.world;
            return {
                ...state,
                worlds: [...state.worlds, world],
                currentWorldId: world.id,
                currentStoryId: null,
                currentTab: 'menu',
            };
        }
    }
}

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

export interface AppState {
    worlds: World[];
    currentWorld: World | null;
    currentStory: Story | null;
    /** Combined lore/characters/locations: world-level merged with story-level */
    mergedLore: LoreEntry[];
    mergedCharacters: Character[];
    mergedLocations: Location[];
    currentTab: Tab;
    chatOpen: boolean;
    connection: LLM.Connection | null;
    model: LLM.ModelInfo | null;
    enableThinking: boolean;
    bannedTokens: string[];
    systemInstruction: string;
    /** Effective system instruction: world override if set, otherwise global */
    effectiveSystemInstruction: 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>(() => {
        const currentWorld = state.worlds.find(w => w.id === state.currentWorldId) ?? null;
        const currentStory = currentWorld?.stories.find(s => s.id === state.currentStoryId) ?? null;

        // Merge world-level + story-level, story takes priority (appended after)
        const mergedLore = [
            ...(currentWorld?.lore ?? []),
            ...(currentStory?.lore ?? []),
        ];
        const mergedCharacters = [
            ...(currentWorld?.characters ?? []),
            ...(currentStory?.characters ?? []),
        ];
        const mergedLocations = [
            ...(currentWorld?.locations ?? []),
            ...(currentStory?.locations ?? []),
        ];

        return {
            worlds: state.worlds,
            currentWorld,
            currentStory,
            mergedLore,
            mergedCharacters,
            mergedLocations,
            currentTab: state.currentTab,
            chatOpen: state.chatOpen,
            connection: state.connection,
            model: state.model,
            enableThinking: state.enableThinking,
            bannedTokens: state.bannedTokens ?? [],
            systemInstruction: state.systemInstruction ?? '',
            effectiveSystemInstruction: currentWorld?.systemInstructionOverride ?? state.systemInstruction ?? '',
            dispatch,
        };
    }, [state]);

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