src/store/useStore.ts

import { create } from 'zustand';
import { User, Workflow, AIResponse } from '@shared/types';
import i18n from '@/i18n';
import { getDatabase } from '@/lib/db';

interface AppState {
  user: any | null;
  token: string | null;
  language: string;
  setLanguage: (lang: string) => void;
  setUser: (user: any | null) => void;
  setToken: (token: string | null) => void;
  logout: () => void;
  
  currentSessionId: string | null;
  setCurrentSessionId: (id: string | null) => void;
  chatHistory: any[];
  setChatHistory: (history: any[] | ((prev: any[]) => any[])) => void;
  addChatMessage: (message: any) => void;
  updateLastChatMessage: (content: string) => void;
  clearChat: () => void;
  
  workflows: any[];
  activeWorkflow: any | null;
  setActiveWorkflow: (workflow: any | null) => void;
  updateWorkflow: (workflow: any) => void;
  
  // RxDB Syncing
  initDatabase: () => Promise<void>;
  
  // Payment Modal State
  isPaymentModalOpen: boolean;
  selectedPlanId: string | null;
  openPaymentModal: (planId?: string) => void;
  closePaymentModal: () => void;
}

export const useStore = create<AppState>((set, get) => ({
  user: JSON.parse(localStorage.getItem('user') || 'null'),
  token: localStorage.getItem('token'),
  language: 'zh',
  setLanguage: (lang) => {
    i18n.changeLanguage(lang);
    set({ language: lang });
  },
  setUser: (user) => {
    if (user) localStorage.setItem('user', JSON.stringify(user));
    else localStorage.removeItem('user');
    set({ user });
  },
  setToken: (token) => {
    if (token) localStorage.setItem('token', token);
    else localStorage.removeItem('token');
    set({ token });
  },
  logout: () => {
    localStorage.removeItem('user');
    localStorage.removeItem('token');
    set({ user: null, token: null, currentSessionId: null, chatHistory: [] });
  },
  
  currentSessionId: null,
  setCurrentSessionId: (id) => set({ currentSessionId: id }),
  chatHistory: [],
  setChatHistory: (history) => set((state) => {
    const newHistory = typeof history === 'function' ? history(state.chatHistory) : history;
    return { chatHistory: newHistory };
  }),
  addChatMessage: async (message) => {
    set((state) => ({ 
      chatHistory: [...state.chatHistory, message] 
    }));
    // Save to RxDB
    const db = await getDatabase();
    await db.chats.insert({
      id: message.id || Math.random().toString(36).substring(7),
      sessionId: get().currentSessionId || 'default',
      role: message.role,
      content: message.content,
      timestamp: Date.now(),
      status: 'sent'
    });
  },
  updateLastChatMessage: async (content) => {
    set((state) => {
      const lastMsg = state.chatHistory[state.chatHistory.length - 1];
      const updatedHistory = state.chatHistory.map((msg, i) => 
        i === state.chatHistory.length - 1 ? { ...msg, content: msg.content + content } : msg
      );
      return { chatHistory: updatedHistory };
    });
    
    // Update RxDB
    const lastMsg = get().chatHistory[get().chatHistory.length - 1];
    if (lastMsg) {
      const db = await getDatabase();
      const doc = await db.chats.findOne(lastMsg.id).exec();
      if (doc) {
        await doc.patch({ content: lastMsg.content });
      }
    }
  },
  clearChat: () => set({ chatHistory: [], currentSessionId: null }),
  
  workflows: [],
  activeWorkflow: null,
  setActiveWorkflow: (workflow) => set({ activeWorkflow: workflow }),
  updateWorkflow: async (workflow) => {
    set((state) => ({
      workflows: state.workflows.map((w) => w.id === workflow.id ? workflow : w)
    }));
    // Save to RxDB
    const db = await getDatabase();
    await db.workflows.upsert({
      ...workflow,
      updatedAt: Date.now()
    });
  },

  initDatabase: async () => {
    const db = await getDatabase();
    // Load chats for current session if exists
    const sessionId = get().currentSessionId || 'default';
    const chats = await db.chats.find({
      selector: { sessionId },
      sort: [{ timestamp: 'asc' }]
    }).exec();
    set({ chatHistory: chats.map(d => d.toJSON()) });
    
    // Load workflows
    const workflows = await db.workflows.find().exec();
    set({ workflows: workflows.map(d => d.toJSON()) });
  },

  // Payment Modal
  isPaymentModalOpen: false,
  selectedPlanId: null,
  openPaymentModal: (planId) => set({ isPaymentModalOpen: true, selectedPlanId: planId || null }),
  closePaymentModal: () => set({ isPaymentModalOpen: false, selectedPlanId: null }),
}));