/**
 * SaaS Store - SaaS Platform Connection State Management
 *
 * Manages SaaS login state, account info, connection mode,
 * and available models. Persists auth state to localStorage
 * via saas-client helpers.
 *
 * Connection modes:
 * - 'tauri':   Local Kernel via Tauri (default)
 * - 'gateway': External Gateway via WebSocket
 * - 'saas':    SaaS backend relay
 */

import { create } from 'zustand';
import {
  saasClient,
  SaaSApiError,
  loadSaaSSession,
  loadSaaSSessionSync,
  saveSaaSSession,
  clearSaaSSession,
  saveConnectionMode,
  loadConnectionMode,
  type SaaSAccountInfo,
  type SaaSModelInfo,
  type SaaSLoginResponse,
  type TotpSetupResponse,
  type SyncConfigRequest,
  type AgentTemplateAvailable,
  type AgentTemplateFull,
  type BillingPlan,
  type SubscriptionInfo,
  type PaymentResult,
  type PaymentStatus,
} from '../lib/saas-client';
import { createLogger } from '../lib/logger';
import {
  initTelemetryCollector,
  stopTelemetryCollector,
} from '../lib/telemetry-collector';
import {
  startPromptOTASync,
  stopPromptOTASync,
} from '../lib/llm-service';

const log = createLogger('SaaSStore');

// === Device ID ===

/** Generate or load a persistent device ID for this browser instance */
function getOrCreateDeviceId(): string {
  const KEY = 'zclaw-device-id';
  const existing = localStorage.getItem(KEY);
  if (existing) return existing;
  const newId = crypto.randomUUID();
  localStorage.setItem(KEY, newId);
  return newId;
}

const DEVICE_ID = getOrCreateDeviceId();

// === Types ===

export type ConnectionMode = 'tauri' | 'gateway' | 'saas';

export interface SaaSStateSlice {
  isLoggedIn: boolean;
  account: SaaSAccountInfo | null;
  saasUrl: string;
  authToken: string | null;
  connectionMode: ConnectionMode;
  availableModels: SaaSModelInfo[];
  isLoading: boolean;
  error: string | null;
  totpRequired: boolean;
  totpSetupData: TotpSetupResponse | null;
  /** Whether SaaS backend is currently reachable */
  saasReachable: boolean;
  /** Agent templates available for onboarding */
  availableTemplates: AgentTemplateAvailable[];
  /** Currently assigned template (null if not yet assigned or assignment removed) */
  assignedTemplate: AgentTemplateFull | null;
  /** Consecutive heartbeat/health-check failures */
  _consecutiveFailures: number;
  _heartbeatTimer?: ReturnType<typeof setInterval>;
  _healthCheckTimer?: ReturnType<typeof setInterval>;

  // === Billing State ===
  plans: BillingPlan[];
  subscription: SubscriptionInfo | null;
  billingLoading: boolean;
  billingError: string | null;
}

export interface SaaSActionsSlice {
  login: (saasUrl: string, username: string, password: string) => Promise<void>;
  loginWithTotp: (saasUrl: string, username: string, password: string, totpCode: string) => Promise<void>;
  register: (saasUrl: string, username: string, email: string, password: string, displayName?: string) => Promise<void>;
  logout: () => Promise<void>;
  setConnectionMode: (mode: ConnectionMode) => void;
  fetchAvailableModels: () => Promise<void>;
  syncConfigFromSaaS: () => Promise<void>;
  pushConfigToSaaS: () => Promise<void>;
  registerCurrentDevice: () => Promise<void>;
  fetchAvailableTemplates: () => Promise<void>;
  /** Assign a template to the current account */
  assignTemplate: (templateId: string) => Promise<void>;
  /** Fetch the currently assigned template (auto-called after login) */
  fetchAssignedTemplate: () => Promise<void>;
  /** Unassign the current template */
  unassignTemplate: () => Promise<void>;
  clearError: () => void;
  restoreSession: () => void;
  // === Billing Actions ===
  fetchPlans: () => Promise<void>;
  fetchSubscription: () => Promise<void>;
  fetchBillingOverview: () => Promise<void>;
  createPayment: (planId: string, method: 'alipay' | 'wechat') => Promise<PaymentResult | null>;
  pollPaymentStatus: (paymentId: string) => Promise<PaymentStatus | null>;
  clearBillingError: () => void;
  setupTotp: () => Promise<TotpSetupResponse>;
  verifyTotp: (code: string) => Promise<void>;
  disableTotp: (password: string) => Promise<void>;
  cancelTotpSetup: () => void;
}

export type SaaSStore = SaaSStateSlice & SaaSActionsSlice;

// === Constants ===

const DEFAULT_SAAS_URL = import.meta.env.VITE_SAAS_URL || 'http://127.0.0.1:8080';

// === Helpers ===

/** Determine the initial connection mode from persisted state */
function resolveInitialMode(sessionMeta: { saasUrl: string; account: SaaSAccountInfo | null } | null): ConnectionMode {
  const persistedMode = loadConnectionMode();
  if (persistedMode === 'tauri' || persistedMode === 'gateway' || persistedMode === 'saas') {
    return persistedMode;
  }
  return sessionMeta ? 'saas' : 'tauri';
}

// === Store Implementation ===

export const useSaaSStore = create<SaaSStore>((set, get) => {
  // Restore session metadata synchronously (URL + account only).
  // Token is loaded from secure storage asynchronously by restoreSession().
  const sessionMeta = loadSaaSSessionSync();
  const initialMode = resolveInitialMode(sessionMeta);

  // If session URL exists, configure the singleton client base URL
  if (sessionMeta) {
    saasClient.setBaseUrl(sessionMeta.saasUrl);
  }

  return {
    // === Initial State ===
    // isLoggedIn starts false — will be set to true by restoreSession()
    isLoggedIn: false,
    account: sessionMeta?.account ?? null,
    saasUrl: sessionMeta?.saasUrl ?? DEFAULT_SAAS_URL,
    authToken: null, // In-memory only — loaded from secure storage by restoreSession()
    connectionMode: initialMode,
    availableModels: [],
    isLoading: false,
    error: null,
    totpRequired: false,
    totpSetupData: null,
    saasReachable: true,
    availableTemplates: [],
    assignedTemplate: null,
    _consecutiveFailures: 0,

    // === Billing State ===
    plans: [],
    subscription: null,
    billingLoading: false,
    billingError: null,

    // === Actions ===

    login: async (saasUrl: string, username: string, password: string) => {
      set({ isLoading: true, error: null });

      const trimmedUrl = saasUrl.trim();
      const trimmedUsername = username.trim();
      const normalizedUrl = trimmedUrl.replace(/\/+$/, '');
      const requestUrl = normalizedUrl || window.location.origin;

      try {
        // 空 trimmedUrl 表示走 Vite proxy（开发模式），允许通过
        if (!trimmedUsername) {
          throw new Error('请输入用户名');
        }
        if (!password) {
          throw new Error('请输入密码');
        }

        // Configure singleton client and attempt login
        saasClient.setBaseUrl(normalizedUrl);
        const loginData: SaaSLoginResponse = await saasClient.login(trimmedUsername, password);

        // Persist session: token + refresh token → secure storage, metadata → localStorage
        const sessionData = {
          token: loginData.token,
          refreshToken: loginData.refresh_token,
          account: loginData.account,
          saasUrl: normalizedUrl,
        };
        try {
          await saveSaaSSession(sessionData);
        } catch (e) {
          log.warn('Failed to persist SaaS session after login', { error: e });
        }
        saveConnectionMode('saas');

        set({
          isLoggedIn: true,
          account: loginData.account,
          saasUrl: normalizedUrl,
          authToken: null, // Not stored in Zustand state — saasClient holds in memory
          connectionMode: 'saas',
          isLoading: false,
          error: null,
        });

        // Register device and start heartbeat in background
        get().registerCurrentDevice().catch((err: unknown) => {
          log.warn('Failed to register device:', err);
        });

        // Fetch available templates in background (non-blocking)
        get().fetchAvailableTemplates().catch((err: unknown) => {
          log.warn('Failed to fetch templates after login:', err);
        });

        // Fetch assigned template in background (non-blocking)
        get().fetchAssignedTemplate().catch((err: unknown) => {
          log.warn('Failed to fetch assigned template after login:', err);
        });

        // Fetch billing data in background (non-blocking)
        get().fetchBillingOverview().catch((err: unknown) => {
          log.warn('Failed to fetch billing after login:', err);
        });

        // Fetch available models in background (non-blocking)
        get().fetchAvailableModels().catch((err: unknown) => {
          log.warn('Failed to fetch models after login:', err);
        });

        // Auto-pull SaaS config in background (non-blocking)
        get().syncConfigFromSaaS().then(() => {
          // After pull, push any locally modified configs back to SaaS
          get().pushConfigToSaaS().catch((err: unknown) => {
            log.warn('Failed to push config to SaaS:', err);
          });
        }).catch((err: unknown) => {
          log.warn('Failed to sync config after login:', err);
        });

        // Initialize telemetry collector
        initTelemetryCollector(DEVICE_ID);

        // Start Prompt OTA sync (background, non-blocking)
        startPromptOTASync(DEVICE_ID);
      } catch (err: unknown) {
        // Check for TOTP required signal
        if (err instanceof SaaSApiError && err.code === 'TOTP_ERROR' && err.status === 400) {
          set({ isLoading: false, totpRequired: true, error: null });
          return;
        }

        const message = err instanceof SaaSApiError
          ? err.message
          : err instanceof Error
            ? err.message
            : String(err);

        const isTimeout = message.includes('signal timed out')
          || message.includes('Timeout')
          || message.includes('timed out')
          || message.includes('AbortError');

        const isConnectionRefused = message.includes('Failed to fetch')
          || message.includes('NetworkError')
          || message.includes('ECONNREFUSED')
          || message.includes('connection refused');

        const userMessage = isTimeout
          ? `连接 SaaS 服务器超时，请确认后端服务正在运行: ${requestUrl}`
          : isConnectionRefused
            ? `无法连接到 SaaS 服务器，请确认后端服务已启动: ${requestUrl}`
            : message;

        set({ isLoading: false, error: userMessage });
        throw new Error(userMessage);
      }
    },

    loginWithTotp: async (saasUrl: string, username: string, password: string, totpCode: string) => {
      set({ isLoading: true, error: null, totpRequired: false });

      try {
        const normalizedUrl = saasUrl.trim().replace(/\/+$/, '');
        saasClient.setBaseUrl(normalizedUrl);
        const loginData = await saasClient.login(username.trim(), password, totpCode);

        const sessionData = {
          token: loginData.token,
          refreshToken: loginData.refresh_token,
          account: loginData.account,
          saasUrl: normalizedUrl,
        };
        try {
          await saveSaaSSession(sessionData);
        } catch (e) {
          log.warn('Failed to persist SaaS session after TOTP login', { error: e });
        }
        saveConnectionMode('saas');

        set({
          isLoggedIn: true,
          account: loginData.account,
          saasUrl: normalizedUrl,
          authToken: null,
          connectionMode: 'saas',
          isLoading: false,
          error: null,
          totpRequired: false,
        });

        get().registerCurrentDevice().catch((err: unknown) => {
          log.warn('Failed to register device:', err);
        });
        get().fetchAvailableModels().catch((err: unknown) => {
          log.warn('Failed to fetch models:', err);
        });

        // Initialize telemetry collector
        initTelemetryCollector(DEVICE_ID);

        // Start Prompt OTA sync (background, non-blocking)
        startPromptOTASync(DEVICE_ID);
      } catch (err: unknown) {
        const message = err instanceof SaaSApiError ? err.message
          : err instanceof Error ? err.message : String(err);
        set({ isLoading: false, error: message });
        throw new Error(message);
      }
    },

    register: async (saasUrl: string, username: string, email: string, password: string, displayName?: string) => {
      set({ isLoading: true, error: null });

      try {
        const trimmedUrl = saasUrl.trim();
        // 空 trimmedUrl 表示走 Vite proxy（开发模式），允许通过
        if (!username.trim()) {
          throw new Error('请输入用户名');
        }
        if (!email.trim()) {
          throw new Error('请输入邮箱');
        }
        if (!password) {
          throw new Error('请输入密码');
        }

        const normalizedUrl = trimmedUrl.replace(/\/+$/, '');

        saasClient.setBaseUrl(normalizedUrl);
        const registerData: SaaSLoginResponse = await saasClient.register({
          username: username.trim(),
          email: email.trim(),
          password,
          display_name: displayName,
        });

        const sessionData = {
          token: registerData.token,
          refreshToken: registerData.refresh_token,
          account: registerData.account,
          saasUrl: normalizedUrl,
        };
        try {
          await saveSaaSSession(sessionData);
        } catch (e) {
          log.warn('Failed to persist SaaS session after register', { error: e });
        }
        saveConnectionMode('saas');

        set({
          isLoggedIn: true,
          account: registerData.account,
          saasUrl: normalizedUrl,
          authToken: null,
          connectionMode: 'saas',
          isLoading: false,
          error: null,
        });

        get().registerCurrentDevice().catch((err: unknown) => {
          log.warn('Failed to register device after register:', err);
        });

        get().fetchAvailableModels().catch((err: unknown) => {
          log.warn('Failed to fetch models after register:', err);
        });

        // Initialize telemetry collector
        initTelemetryCollector(DEVICE_ID);

        // Start Prompt OTA sync
        startPromptOTASync(DEVICE_ID);
      } catch (err: unknown) {
        const message = err instanceof SaaSApiError
          ? err.message
          : err instanceof Error
            ? err.message
            : String(err);

        set({ isLoading: false, error: message });
        throw new Error(message);
      }
    },

    logout: async () => {
      saasClient.setToken(null);
      await clearSaaSSession();
      saveConnectionMode('tauri');
      stopTelemetryCollector();
      stopPromptOTASync();

      // Clear currentModel so next connection uses fresh model resolution
      try {
        const { useConversationStore } = require('./chat/conversationStore');
        useConversationStore.getState().setCurrentModel('');
      } catch { /* non-critical */ }

      set({
        isLoggedIn: false,
        account: null,
        authToken: null,
        connectionMode: 'tauri',
        availableModels: [],
        availableTemplates: [],
        assignedTemplate: null,
        plans: [],
        subscription: null,
        billingLoading: false,
        billingError: null,
        error: null,
        totpRequired: false,
        totpSetupData: null,
      });
    },

    setConnectionMode: (mode: ConnectionMode) => {
      const { isLoggedIn } = get();

      // Cannot switch to SaaS mode if not logged in
      if (mode === 'saas' && !isLoggedIn) {
        return;
      }

      saveConnectionMode(mode);
      set({ connectionMode: mode });
    },

    fetchAvailableModels: async () => {
      const { isLoggedIn, saasUrl } = get();

      if (!isLoggedIn) {
        set({ availableModels: [] });
        return;
      }

      try {
        saasClient.setBaseUrl(saasUrl);
        const models = await saasClient.listModels();
        set({ availableModels: models });

        // Sync currentModel: if the stored model is not in the available list,
        // switch to the first available model to prevent 404 errors
        if (models.length > 0) {
          try {
            const { useConversationStore } = await import('./chat/conversationStore');
            const current = useConversationStore.getState().currentModel;
            const modelIds = models.map(m => m.alias || m.id);
            if (current && !modelIds.includes(current)) {
              const firstModel = models[0];
              const fallbackId = firstModel.alias || firstModel.id;
              useConversationStore.getState().setCurrentModel(fallbackId);
              log.info(`Synced currentModel: ${current} not available, switched to ${fallbackId}`);
            }
          } catch (syncErr) {
            log.warn('Failed to sync currentModel after fetching models:', syncErr);
          }
        }
      } catch (err: unknown) {
        log.warn('Failed to fetch available models:', err);
        // Do not set error state - model fetch failure is non-critical
        set({ availableModels: [] });
      }
    },

    /**
     * Push locally modified configs to SaaS (push direction of bidirectional sync).
     * Collects all "dirty" config keys, computes diff, and syncs via merge.
     */
    pushConfigToSaaS: async () => {
      const { isLoggedIn, saasUrl } = get();
      if (!isLoggedIn) return;

      try {
        saasClient.setBaseUrl(saasUrl);

        // Collect all dirty config keys
        const dirtyKeys: string[] = [];
        const dirtyValues: Record<string, unknown> = {};
        let i = 0;
        while (true) {
          const key = localStorage.key(i);
          if (!key) break;
          i++;
          if (key.startsWith('zclaw-config-dirty.') && localStorage.getItem(key) === '1') {
            const configKey = key.replace('zclaw-config-dirty.', '');
            const storageKey = `zclaw-config.${configKey}`;
            const value = localStorage.getItem(storageKey);
            if (value !== null) {
              dirtyKeys.push(configKey);
              dirtyValues[configKey] = value;
            }
          }
        }

        if (dirtyKeys.length === 0) return;

        // Generate a client fingerprint
        const fingerprint = DEVICE_ID;
        const syncRequest = {
          client_fingerprint: fingerprint,
          action: 'merge' as const,
          config_keys: dirtyKeys,
          client_values: dirtyValues,
        };

        // Compute diff first (dry run)
        const diff = await saasClient.computeConfigDiff(syncRequest as SyncConfigRequest);

        if (diff.conflicts > 0) {
          log.warn(`Config sync has ${diff.conflicts} conflicts, using merge strategy`);
        }

        // Perform actual sync
        const result = await saasClient.syncConfig(syncRequest);
        log.info(`Config push result: ${result.updated} updated, ${result.created} created, ${result.skipped} skipped`);

        // Clear dirty flags for successfully synced keys
        for (const key of dirtyKeys) {
          localStorage.removeItem(`zclaw-config-dirty.${key}`);
        }
      } catch (err: unknown) {
        log.warn('Failed to push config to SaaS:', err);
      }
    },

    /** Pull SaaS config and apply to local storage (startup auto-sync) */
    syncConfigFromSaaS: async () => {
      const { isLoggedIn, saasUrl } = get();

      if (!isLoggedIn) return;

      try {
        saasClient.setBaseUrl(saasUrl);

        // Read last sync timestamp from localStorage
        const lastSyncKey = 'zclaw-config-last-sync';
        const lastSync = localStorage.getItem(lastSyncKey) || undefined;

        const result = await saasClient.pullConfig(lastSync);

        if (result.configs.length === 0) {
          log.info('No config updates from SaaS');
          return;
        }

        // Apply SaaS config values to localStorage
        // Each config is stored as zclaw-config.{category}.{key}
        for (const config of result.configs) {
          if (config.value === null) continue;

          const storageKey = `zclaw-config.${config.category}.${config.key}`;
          const existing = localStorage.getItem(storageKey);

          // Diff check: skip if local was modified since last pull
          const dirtyKey = `zclaw-config-dirty.${config.category}.${config.key}`;
          const lastPulledValue = localStorage.getItem(`zclaw-config-pulled.${config.category}.${config.key}`);

          if (dirtyKey && localStorage.getItem(dirtyKey) === '1') {
            // Local was modified since last pull → keep local, skip overwrite
            log.warn(`Config conflict, keeping local: ${config.key}`);
            continue;
          }

          // If existing value differs from what we last pulled AND differs from SaaS, local was modified
          if (existing !== null && lastPulledValue !== null && existing !== lastPulledValue && existing !== config.value) {
            log.warn(`Config conflict (local modified), keeping local: ${config.key}`);
            continue;
          }

          // Only update if the value has actually changed
          if (existing !== config.value) {
            localStorage.setItem(storageKey, config.value);
            // Record the pulled value for future diff checks
            localStorage.setItem(`zclaw-config-pulled.${config.category}.${config.key}`, config.value);
            log.info(`Config synced: ${config.key}`);
          }
        }

        // Update last sync timestamp
        localStorage.setItem(lastSyncKey, result.pulled_at);
        log.info(`Synced ${result.configs.length} config items from SaaS`);

        // Propagate Kernel-relevant configs to Rust backend
        const kernelCategories = ['agent', 'llm'];
        const kernelConfigs = result.configs.filter(
          (c) => kernelCategories.includes(c.category) && c.value !== null
        );
        if (kernelConfigs.length > 0 && typeof window !== 'undefined' && '__TAURI_INTERNALS__' in window) {
          try {
            const { invoke } = await import('@tauri-apps/api/core');
            await invoke('kernel_apply_saas_config', {
              configs: kernelConfigs.map((c) => ({
                category: c.category,
                key: c.key,
                value: c.value,
              })),
            });
            log.info(`Propagated ${kernelConfigs.length} Kernel configs to Rust backend`);
          } catch (invokeErr: unknown) {
            log.warn('Failed to propagate configs to Kernel (non-fatal):', invokeErr);
          }
        }
      } catch (err: unknown) {
        log.warn('Failed to sync config from SaaS:', err);
      }
    },

    registerCurrentDevice: async () => {
      const { isLoggedIn, saasUrl } = get();

      if (!isLoggedIn) {
        return;
      }

      try {
        saasClient.setBaseUrl(saasUrl);
        await saasClient.registerDevice({
          device_id: DEVICE_ID,
          device_name: `${navigator.userAgent.split(' ').slice(0, 3).join(' ')}`,
          platform: navigator.platform,
          app_version: (typeof __APP_VERSION__ !== 'undefined' ? __APP_VERSION__ : 'unknown'),
        });
        log.info('Device registered successfully');

        // Start periodic heartbeat (every 5 minutes) with failure tracking
        if (typeof window !== 'undefined' && !get()._heartbeatTimer) {
          const DEGRADE_AFTER_FAILURES = 3; // Degrade after 3 consecutive failures (~15 min)
          const timer = window.setInterval(async () => {
            const state = get();
            if (!state.isLoggedIn) {
              window.clearInterval(timer);
              return;
            }
            try {
              await saasClient.deviceHeartbeat(DEVICE_ID);
              // Reset failure count on success
              if (state._consecutiveFailures > 0) {
                log.info(`Heartbeat recovered after ${state._consecutiveFailures} failures`);
              }
              set({ _consecutiveFailures: 0, saasReachable: true } as unknown as Partial<SaaSStore>);
            } catch (err) {
              const failures = state._consecutiveFailures + 1;
              log.warn(`Heartbeat failed (${failures}/${DEGRADE_AFTER_FAILURES}): ${err}`);
              set({ _consecutiveFailures: failures } as unknown as Partial<SaaSStore>);

              // Auto-degrade to local mode after threshold
              if (failures >= DEGRADE_AFTER_FAILURES && state.connectionMode === 'saas') {
                log.warn(`SaaS unreachable after ${failures} attempts — degrading to local mode`);
                set({
                  saasReachable: false,
                  connectionMode: 'tauri',
                } as unknown as Partial<SaaSStore>);
                saveConnectionMode('tauri');
              }
            }
          }, 5 * 60 * 1000);
          set({ _heartbeatTimer: timer, _consecutiveFailures: 0 } as unknown as Partial<SaaSStore>);
        }
      } catch (err: unknown) {
        log.warn('Failed to register device:', err);
      }
    },

    fetchAvailableTemplates: async () => {
      try {
        const templates = await saasClient.fetchAvailableTemplates();
        set({ availableTemplates: templates });
      } catch {
        // Graceful degradation - don't block login
        set({ availableTemplates: [] });
      }
    },

    assignTemplate: async (templateId: string) => {
      try {
        const template = await saasClient.assignTemplate(templateId);
        set({ assignedTemplate: template });
      } catch (err) {
        log.warn('Failed to assign template:', err);
        // Don't throw — let wizard continue with fallback flow
      }
    },

    fetchAssignedTemplate: async () => {
      try {
        const template = await saasClient.getAssignedTemplate();
        set({ assignedTemplate: template });
      } catch {
        // Not critical — null is fine
        set({ assignedTemplate: null });
      }
    },

    unassignTemplate: async () => {
      try {
        await saasClient.unassignTemplate();
        set({ assignedTemplate: null });
      } catch (err) {
        log.warn('Failed to unassign template:', err);
      }
    },

    clearError: () => {
      set({ error: null });
    },

    restoreSession: async () => {
      const restored = await loadSaaSSession();
      if (!restored) return;

      saasClient.setBaseUrl(restored.saasUrl);

      // Strategy: access token → refresh token → cookie auth → fail
      let account: SaaSAccountInfo | null = null;
      let newToken: string | null = restored.token;

      if (restored.token) {
        // Token from secure storage — use as Bearer
        saasClient.setToken(restored.token);
        try {
          account = await saasClient.me();
        } catch {
          // Access token expired — clear and try refresh
          saasClient.setToken(null);
          newToken = null;
        }
      }

      if (!account && restored.refreshToken) {
        // Try refresh token from secure storage
        saasClient.setRefreshToken(restored.refreshToken);
        try {
          const refreshed = await saasClient.refreshMutex();
          newToken = refreshed;
          saasClient.setToken(refreshed);
          account = await saasClient.me();
          // Persist the new tokens back to secure storage
          try {
            const { saveSaaSSession: save } = await import('../lib/saas-session');
            await save({
              token: refreshed,
              refreshToken: saasClient.getRefreshToken(),
              account,
              saasUrl: restored.saasUrl,
            });
          } catch (e) {
            log.warn('Failed to persist refreshed session', { error: e });
          }
        } catch {
          // Refresh token also expired or invalid
          saasClient.setRefreshToken(null);
          saasClient.setToken(null);
          newToken = null;
        }
      }

      if (!account) {
        // Try cookie-based auth (works for same-origin, e.g. admin panel)
        account = await saasClient.restoreFromCookie();
      }

      if (!account) {
        // All methods failed — user needs to re-login
        set({
          isLoggedIn: false,
          account: null,
          saasUrl: restored.saasUrl,
          authToken: null,
        });
        return;
      }

      set({
        isLoggedIn: true,
        account,
        saasUrl: restored.saasUrl,
        authToken: newToken,
        connectionMode: loadConnectionMode() === 'saas' ? 'saas' : 'tauri',
      });
      get().fetchAvailableModels().catch(() => {});
      get().fetchAvailableTemplates().catch(() => {});
      get().fetchAssignedTemplate().catch(() => {});
      get().fetchBillingOverview().catch(() => {});
      get().syncConfigFromSaaS().then(() => {
        get().pushConfigToSaaS().catch(() => {});
      }).catch(() => {});
      initTelemetryCollector(DEVICE_ID);
      startPromptOTASync(DEVICE_ID);
    },

    setupTotp: async () => {
      set({ isLoading: true, error: null });
      try {
        const setupData = await saasClient.setupTotp();
        set({ totpSetupData: setupData, isLoading: false });
        return setupData;
      } catch (err: unknown) {
        const message = err instanceof SaaSApiError ? err.message
          : err instanceof Error ? err.message : String(err);
        set({ isLoading: false, error: message });
        throw new Error(message);
      }
    },

    verifyTotp: async (code: string) => {
      set({ isLoading: true, error: null });
      try {
        await saasClient.verifyTotp(code);
        const account = await saasClient.me();
        const { saasUrl } = get();
        saveSaaSSession({ token: null, refreshToken: null, account, saasUrl }).catch((e) =>
          log.warn('Failed to persist SaaS session after verifyTotp', { error: e })
        ); // Token in saasClient memory only
        set({ totpSetupData: null, isLoading: false, account });
      } catch (err: unknown) {
        const message = err instanceof SaaSApiError ? err.message
          : err instanceof Error ? err.message : String(err);
        set({ isLoading: false, error: message });
        throw new Error(message);
      }
    },

    disableTotp: async (password: string) => {
      set({ isLoading: true, error: null });
      try {
        await saasClient.disableTotp(password);
        const account = await saasClient.me();
        const { saasUrl } = get();
        saveSaaSSession({ token: null, refreshToken: null, account, saasUrl }).catch((e) =>
          log.warn('Failed to persist SaaS session after disableTotp', { error: e })
        );
        set({ isLoading: false, account });
      } catch (err: unknown) {
        const message = err instanceof SaaSApiError ? err.message
          : err instanceof Error ? err.message : String(err);
        set({ isLoading: false, error: message });
        throw new Error(message);
      }
    },

    cancelTotpSetup: () => {
      set({ totpSetupData: null });
    },

    // === Billing Actions ===

    fetchPlans: async () => {
      try {
        const plans = await saasClient.listPlans();
        set({ plans });
      } catch (err: unknown) {
        log.warn('Failed to fetch plans:', err);
      }
    },

    fetchSubscription: async () => {
      try {
        const sub = await saasClient.getSubscription();
        set({ subscription: sub });
      } catch (err: unknown) {
        log.warn('Failed to fetch subscription:', err);
        set({ subscription: null });
      }
    },

    fetchBillingOverview: async () => {
      set({ billingLoading: true, billingError: null });
      try {
        const [plans, sub] = await Promise.all([
          saasClient.listPlans(),
          saasClient.getSubscription(),
        ]);
        set({ plans, subscription: sub, billingLoading: false });
      } catch (err: unknown) {
        const msg = err instanceof Error ? err.message : String(err);
        set({ billingLoading: false, billingError: msg });
      }
    },

    createPayment: async (planId: string, method: 'alipay' | 'wechat') => {
      set({ billingLoading: true, billingError: null });
      try {
        const result = await saasClient.createPayment({ plan_id: planId, payment_method: method });
        set({ billingLoading: false });
        return result;
      } catch (err: unknown) {
        const msg = err instanceof Error ? err.message : String(err);
        set({ billingLoading: false, billingError: msg });
        return null;
      }
    },

    pollPaymentStatus: async (paymentId: string) => {
      try {
        const status = await saasClient.getPaymentStatus(paymentId);
        if (status.status === 'succeeded') {
          // Refresh subscription after successful payment
          await get().fetchSubscription();
        }
        return status;
      } catch (err: unknown) {
        log.warn('Failed to poll payment status:', err);
        return null;
      }
    },

    clearBillingError: () => {
      set({ billingError: null });
    },
  };
});