apix/lib/store.ts

import { create } from 'zustand';
import { persist } from 'zustand/middleware';
import { db, ImageItem as DBImageItem } from './db';


// ImageItem definition is now in db.ts but we keep a compatible interface if needed, 
// or import it. Let's reuse DBImageItem for consistency in the store.
export type ImageItem = DBImageItem;


export interface HistoryItem {
    id: string;
    url: string;
    category: string;
    originalName: string;
}

export interface VideoItem {
    id: string;
    url: string; // Blob URL or remote URL
    prompt: string;
    thumbnail?: string; // Optional thumbnail from source image
    createdAt?: number;
}

export type ViewType = 'gallery' | 'settings' | 'library' | 'history';
export type ReferenceCategory = 'subject' | 'scene' | 'style' | 'video'; // Added 'video'

interface AppState {
    // Navigation
    currentView: ViewType;
    setCurrentView: (view: ViewType) => void;

    // Reference Selection
    selectionMode: ReferenceCategory | null;
    setSelectionMode: (mode: ReferenceCategory | null) => void;

    // Multiple references per category (arrays)
    references: {
        subject?: Array<{ id: string; thumbnail?: string }>;
        scene?: Array<{ id: string; thumbnail?: string }>;
        style?: Array<{ id: string; thumbnail?: string }>;
        video?: Array<{ id: string; thumbnail?: string }>;
    };
    // Add a reference to a category (appends to array)
    addReference: (category: ReferenceCategory, ref: { id: string; thumbnail?: string }) => void;
    // Remove a specific reference by ID from a category
    removeReference: (category: ReferenceCategory, refId: string) => void;
    // Clear all references for a category
    clearReferences: (category: ReferenceCategory) => void;
    // Legacy setter for backwards compatibility (replaces all refs in a category)
    setReference: (category: ReferenceCategory, ref: { id: string; thumbnail?: string } | undefined) => void;

    prompt: string;
    setPrompt: (p: string) => void;

    gallery: ImageItem[];
    loadGallery: () => Promise<void>;
    addToGallery: (image: ImageItem) => Promise<void>;
    removeFromGallery: (id: number) => Promise<void>;
    clearGallery: () => Promise<void>;

    isGenerating: boolean;
    setIsGenerating: (isGenerating: boolean) => void;

    showCookieExpired: boolean;
    setShowCookieExpired: (show: boolean) => void;


    // Videos
    videos: VideoItem[];
    addVideo: (video: VideoItem) => void;
    removeVideo: (id: string) => void;

    history: HistoryItem[];
    setHistory: (items: HistoryItem[]) => void;
    removeFromHistory: (id: string) => void;

    settings: {
        aspectRatio: string;
        preciseMode: boolean;
        imageCount: number;
        theme: 'light' | 'dark';
        // Provider selection
        provider: 'whisk' | 'grok' | 'meta';
        // Whisk (Google)
        whiskCookies: string;
        // Grok (xAI)
        grokApiKey: string;
        grokCookies: string;
        // Meta AI
        metaCookies: string;
    };
    setSettings: (s: Partial<AppState['settings']>) => void;
}

export const useStore = create<AppState>()(
    persist(
        (set) => ({
            currentView: 'gallery',
            setCurrentView: (view) => set({ currentView: view }),

            selectionMode: null,
            setSelectionMode: (mode) => set({ selectionMode: mode }),

            references: {},
            // Add a reference to a category array
            addReference: (category, ref) => set((state) => ({
                references: {
                    ...state.references,
                    [category]: [...(state.references[category] || []), ref]
                }
            })),
            // Remove a specific reference by ID
            removeReference: (category, refId) => set((state) => ({
                references: {
                    ...state.references,
                    [category]: (state.references[category] || []).filter(r => r.id !== refId)
                }
            })),
            // Clear all references for a category
            clearReferences: (category) => set((state) => ({
                references: {
                    ...state.references,
                    [category]: []
                }
            })),
            // Legacy setter (replaces entire category with single ref or clears if undefined)
            setReference: (category, ref) => set((state) => ({
                references: { ...state.references, [category]: ref ? [ref] : undefined }
            })),

            prompt: '',
            setPrompt: (p) => set({ prompt: p }),

            gallery: [],
            loadGallery: async () => {
                const items = await db.gallery.toArray();
                // Sort by createdAt desc if needed
                items.sort((a, b) => b.createdAt - a.createdAt);
                set({ gallery: items });
            },
            addToGallery: async (img) => {
                const id = await db.gallery.add(img);
                const newImg = { ...img, id };
                set((state) => ({ gallery: [newImg, ...state.gallery] }));
            },
            removeFromGallery: async (id) => {
                if (!id) return;
                await db.gallery.delete(id);
                set((state) => ({
                    gallery: state.gallery.filter((item) => item.id !== id)
                }));
            },
            clearGallery: async () => {
                await db.gallery.clear();
                set({ gallery: [] });
            },


            isGenerating: false,
            setIsGenerating: (isGenerating) => set({ isGenerating }),

            showCookieExpired: false,
            setShowCookieExpired: (show) => set({ showCookieExpired: show }),

            // Videos
            videos: [],
            addVideo: (video) => set((state) => ({ videos: [video, ...state.videos] })),
            removeVideo: (id) => set((state) => ({ videos: state.videos.filter(v => v.id !== id) })),

            history: [],
            setHistory: (items) => set({ history: items }),
            removeFromHistory: (id) => set((state) => ({
                history: state.history.filter(item => item.id !== id)
            })),

            settings: {
                aspectRatio: '1:1',
                preciseMode: false,
                imageCount: 4,
                theme: 'dark',
                provider: 'whisk',
                whiskCookies: '',
                grokApiKey: '',
                grokCookies: '',
                metaCookies: ''
            },
            setSettings: (s) => set((state) => ({ settings: { ...state.settings, ...s } }))
        }),
        {
            name: 'kv-pix-storage',
            partialize: (state) => ({
                settings: state.settings,
                // gallery: state.gallery, // Don't persist gallery to localStorage
                history: state.history,
                videos: state.videos // Persist videos
            }),
        }
    )
);