import { create } from 'zustand';
import { persist } from 'zustand/middleware';

export interface Message {
    id: string;
    role: 'user' | 'assistant';
    content: string;
    date?: string;
    createTime?: string;
}

export interface Session {
    id: string;
    topic: string;
    appId?: string;
    messages: Message[];
    createTime?: string;
    updateTime?: string;
}

interface ChatState {
    sessions: Session[];
    currentSessionId: string | null;
    selectedAppId: string | null;
    sidebarOpen: boolean;
    loading: boolean;

    // Actions
    addSession: (session: Session) => void;
    updateSession: (id: string, updater: (session: Session) => void) => void;
    deleteSession: (id: string) => void;
    clearSessions: () => void;
    setCurrentSessionId: (id: string | null) => void;
    setSelectedAppId: (id: string | null) => void;
    setSidebarOpen: (open: boolean) => void;
    setLoading: (loading: boolean) => void;
    addMessage: (sessionId: string, message: Message) => void;
    getCurrentSession: () => Session | null;
    updateCurrentSession: (updater: (session: Session) => void) => void;
}

const initialState = {
    sessions: [],
    currentSessionId: null,
    selectedAppId: null,
    sidebarOpen: true,
    loading: false,
};

export const useChatStore = create<ChatState>()(
    persist(
        (set, get) => ({
            ...initialState,

            addSession: (session) => {
                set((state) => ({
                    sessions: [session, ...state.sessions],
                    currentSessionId: session.id,
                }));
            },

            updateSession: (id, updater) => {
                set((state) => ({
                    sessions: state.sessions.map((session) => {
                        if (session.id === id) {
                            const updated = { ...session };
                            updater(updated);
                            return updated;
                        }
                        return session;
                    }),
                }));
            },

            deleteSession: (id) => {
                set((state) => ({
                    sessions: state.sessions.filter((session) => session.id !== id),
                    currentSessionId: state.currentSessionId === id ? null : state.currentSessionId,
                }));
            },

            clearSessions: () => {
                set({ sessions: [], currentSessionId: null });
            },

            setCurrentSessionId: (id) => {
                set({ currentSessionId: id });
            },

            setSelectedAppId: (id) => {
                set({ selectedAppId: id });
            },

            setSidebarOpen: (open) => {
                set({ sidebarOpen: open });
            },

            setLoading: (loading) => {
                set({ loading });
            },

            addMessage: (sessionId, message) => {
                set((state) => ({
                    sessions: state.sessions.map((session) => {
                        if (session.id === sessionId) {
                            return {
                                ...session,
                                messages: [...session.messages, message],
                                updateTime: new Date().toISOString(),
                            };
                        }
                        return session;
                    }),
                }));
            },

            getCurrentSession: () => {
                const { sessions, currentSessionId } = get();
                return sessions.find((s) => s.id === currentSessionId) || null;
            },

            updateCurrentSession: (updater) => {
                const { sessions, currentSessionId } = get();
                if (!currentSessionId) return;
                
                set({
                    sessions: sessions.map((session) => {
                        if (session.id === currentSessionId) {
                            const updated = { ...session };
                            updater(updated);
                            return updated;
                        }
                        return session;
                    }),
                });
            },
        }),
        {
            name: 'universal-chat-storage',
            partialize: (state) => ({
                sessions: state.sessions,
                currentSessionId: state.currentSessionId,
                selectedAppId: state.selectedAppId,
            }),
        }
    )
);
