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

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

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

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

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

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

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

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

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

type Action =
    | { type: 'CREATE_STORY'; title: string }
    | { type: 'RENAME_STORY'; id: string; title: string }
    | { type: 'EDIT_STORY'; id: string; text: string; highlightText?: string }
    | { type: 'EDIT_SCRATCHPAD'; id: string; text: string }
    | { type: 'ADD_LORE_ENTRY'; storyId: string; entry: LoreEntry }
    | { type: 'EDIT_LORE_ENTRY'; storyId: string; entryId: string; updates: Partial<LoreEntry> }
    | { type: 'DELETE_LORE_ENTRY'; storyId: string; entryId: string }
    | { type: 'REORDER_LORE_ENTRIES'; storyId: string; entryIds: string[] }
    | { type: 'SET_SYSTEM_INSTRUCTION'; systemInstruction: string }
    | { type: 'SET_CURRENT_TAB'; tab: Tab }
    | { type: 'SET_CHAT_OPEN'; open: boolean }
    | { type: 'DELETE_STORY'; id: string }
    | { type: 'SELECT_STORY'; id: string }
    | { type: 'DUPLICATE_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, 'id' | 'name' | '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 }
    | { type: 'ADD_LOCATION'; storyId: string; location: Location }
    | { type: 'EDIT_LOCATION'; storyId: string; locationId: string; updates: Partial<Location> }
    | { type: 'DELETE_LOCATION'; storyId: string; locationId: string }
    | { type: 'STORE_CHAPTER_SUMMARY'; storyId: string; header: string; hash: Chapters.Hash; summary: string }
    | { type: 'CLEAN_CHAPTER_SUMMARIES'; storyId: string; validHashes: Record<string, Chapters.Hash[]> };

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

const DEFAULT_STATE: IState = {
    stories: [],
    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_STORY': {
            const story: Story = {
                id: crypto.randomUUID(),
                title: action.title,
                text: '',
                scratchpad: '',
                lore: [],
                characters: [],
                locations: [],
                chatMessages: [],
                chapters: [],
            };
            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, lastEditedText: action.highlightText }
                        : s
                ),
            };
        }
        case 'ADD_LORE_ENTRY': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.storyId
                        ? { ...s, lore: [...s.lore, action.entry] }
                        : s
                ),
            };
        }
        case 'EDIT_LORE_ENTRY': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.storyId
                        ? {
                            ...s,
                            lore: s.lore.map(e =>
                                e.id === action.entryId
                                    ? { ...e, ...action.updates }
                                    : e
                            ),
                        }
                        : s
                ),
            };
        }
        case 'DELETE_LORE_ENTRY': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.storyId
                        ? { ...s, lore: s.lore.filter(e => e.id !== action.entryId) }
                        : s
                ),
            };
        }
        case 'REORDER_LORE_ENTRIES': {
            return {
                ...state,
                stories: state.stories.map(s => {
                    if (s.id !== action.storyId) return s;
                    const entryMap = new Map(s.lore.map(e => [e.id, e]));
                    const reordered = action.entryIds
                        .map(id => entryMap.get(id))
                        .filter((e): e is LoreEntry => e !== undefined);
                    // Add any entries that weren't in the new order (safety)
                    for (const entry of s.lore) {
                        if (!action.entryIds.includes(entry.id)) {
                            reordered.push(entry);
                        }
                    }
                    return { ...s, lore: reordered };
                }),
            };
        }
        case 'SET_SYSTEM_INSTRUCTION': {
            return {
                ...state,
                systemInstruction: action.systemInstruction,
            };
        }
        case 'SET_CURRENT_TAB': {
            return { ...state, currentTab: action.tab };
        }
        case 'SET_CHAT_OPEN': {
            return { ...state, chatOpen: action.open };
        }
        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 'DUPLICATE_STORY': {
            const original = state.stories.find(s => s.id === action.id);
            if (!original) return state;
            const newStory: Story = {
                id: crypto.randomUUID(),
                title: `${original.title} (Copy)`,
                text: '',
                scratchpad: '',
                lore: [...original.lore],
                characters: original.characters,
                locations: original.locations,
                chatMessages: [],
                chapters: [],
            };
            return {
                ...state,
                stories: [...state.stories, newStory],
                currentStoryId: newStory.id,
            };
        }
        case 'SELECT_STORY': {
            return { ...state, currentStoryId: action.id, currentTab: 'story' };
        }
        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 => {
                    if (s.id !== action.storyId) return s;
                    const deletedChar = s.characters.find(c => c.id === action.characterId);
                    if (!deletedChar) return s;
                    return {
                        ...s,
                        characters: s.characters
                            .filter(c => c.id !== action.characterId)
                            .map(c => ({
                                ...c,
                                relations: c.relations.filter(r => r.name !== deletedChar.name),
                            })),
                    };
                }),
            };
        }
        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
                ),
            };
        }
        case 'ADD_LOCATION': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.storyId
                        ? { ...s, locations: [...s.locations, action.location] }
                        : s
                ),
            };
        }
        case 'EDIT_LOCATION': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.storyId
                        ? {
                            ...s,
                            locations: s.locations.map(l =>
                                l.id === action.locationId
                                    ? { ...l, ...action.updates }
                                    : l
                            ),
                        }
                        : s
                ),
            };
        }
        case 'DELETE_LOCATION': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.storyId
                        ? { ...s, locations: s.locations.filter(l => l.id !== action.locationId) }
                        : s
                ),
            };
        }
        case 'CLEAN_CHAPTER_SUMMARIES': {
            return {
                ...state,
                stories: state.stories.map(s => {
                    if (s.id !== action.storyId) return 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 'EDIT_SCRATCHPAD': {
            return {
                ...state,
                stories: state.stories.map(s =>
                    s.id === action.id ? { ...s, scratchpad: action.text } : s
                ),
            };
        }
        case 'STORE_CHAPTER_SUMMARY': {
            return {
                ...state,
                stories: state.stories.map(s => {
                    if (s.id !== action.storyId) return 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],
                    };
                }),
            };
        }
    }
}

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

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

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