noteflow/client/src/contexts/workspace-context.tsx

// Workspace context for managing current user/workspace identity

import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
import { IdentityDefaults } from '@/api';
import { getAPI } from '@/api/interface';
import { extractErrorMessage } from '@/api';
import type { GetCurrentUserResponse, Workspace } from '@/api/types';
import { useConnectionState } from '@/contexts/connection-state';
import { WorkspaceContext, type WorkspaceContextValue } from '@/contexts/workspace-state';
import { workspaceStorageKey } from '@/contexts/storage';
import { readStorageRaw, writeStorageRaw } from '@/lib/storage/utils';

const fallbackUser: GetCurrentUserResponse = {
  user_id: IdentityDefaults.DEFAULT_USER_ID,
  workspace_id: IdentityDefaults.DEFAULT_WORKSPACE_ID,
  display_name: IdentityDefaults.DEFAULT_USER_NAME,
  is_authenticated: false,
};
const fallbackWorkspace: Workspace = {
  id: IdentityDefaults.DEFAULT_WORKSPACE_ID,
  name: IdentityDefaults.DEFAULT_WORKSPACE_NAME,
  role: 'owner',
  is_default: true,
};

function readStoredWorkspaceId(): string | null {
  const value = readStorageRaw(workspaceStorageKey(), '');
  return value || null;
}

function persistWorkspaceId(workspaceId: string): void {
  writeStorageRaw(workspaceStorageKey(), workspaceId, 'workspace_persist');
}

function resolveWorkspace(workspaces: Workspace[], preferredId: string | null): Workspace | null {
  if (!workspaces.length) {
    return null;
  }
  if (preferredId) {
    const byId = workspaces.find((workspace) => workspace.id === preferredId);
    if (byId) {
      return byId;
    }
  }
  const defaultWorkspace = workspaces.find((workspace) => workspace.is_default);
  return defaultWorkspace ?? workspaces[0] ?? null;
}

export function WorkspaceProvider({ children }: { children: React.ReactNode }) {
  const { isConnected } = useConnectionState();
  const [currentWorkspace, setCurrentWorkspace] = useState<Workspace | null>(null);
  const [workspaces, setWorkspaces] = useState<Workspace[]>([]);
  const [currentUser, setCurrentUser] = useState<GetCurrentUserResponse | null>(null);
  const [isLoading, setIsLoading] = useState(true);
  const [error, setError] = useState<string | null>(null);

  const loadContext = useCallback(async () => {
    setIsLoading(true);
    setError(null);
    try {
      const api = getAPI();
      const [user, workspaceResponse] = await Promise.all([
        api.getCurrentUser(),
        api.listWorkspaces(),
      ]);

      const availableWorkspaces =
        workspaceResponse.workspaces.length > 0
          ? workspaceResponse.workspaces
          : [fallbackWorkspace];

      setCurrentUser(user ?? fallbackUser);
      setWorkspaces(availableWorkspaces);

      const storedId = readStoredWorkspaceId();
      const selected = resolveWorkspace(availableWorkspaces, storedId);
      setCurrentWorkspace(selected);
      if (selected) {
        persistWorkspaceId(selected.id);
      }
    } catch (err) {
      setError(extractErrorMessage(err, 'Failed to load workspace context'));
      setCurrentUser(fallbackUser);
      setWorkspaces([fallbackWorkspace]);
      setCurrentWorkspace(fallbackWorkspace);
      persistWorkspaceId(fallbackWorkspace.id);
    } finally {
      setIsLoading(false);
    }
  }, []);

  const hasLoaded = useRef(false);
  useEffect(() => {
    if (hasLoaded.current) {
      return;
    }
    if (!isConnected) {
      return;
    }
    hasLoaded.current = true;
    void loadContext();
  }, [loadContext, isConnected]);

  // Use ref for workspaces to avoid stale closure in switchWorkspace
  const workspacesRef = useRef(workspaces);
  useEffect(() => {
    workspacesRef.current = workspaces;
  }, [workspaces]);

  const switchWorkspace = useCallback(async (workspaceId: string) => {
    if (!workspaceId) {
      return;
    }
    setIsLoading(true);
    setError(null);
    try {
      const api = getAPI();
      const response = await api.switchWorkspace(workspaceId);
      // Use ref to get current workspaces without stale closure
      const selected =
        response.workspace ??
        workspacesRef.current.find((workspace) => workspace.id === workspaceId);
      if (!response.success || !selected) {
        throw new Error('Workspace not found');
      }
      setCurrentWorkspace(selected);
      persistWorkspaceId(selected.id);
    } catch (err) {
      setError(extractErrorMessage(err, 'Failed to switch workspace'));
      throw err;
    } finally {
      setIsLoading(false);
    }
  }, []);

  const value = useMemo<WorkspaceContextValue>(
    () => ({
      currentWorkspace,
      workspaces,
      currentUser,
      switchWorkspace,
      isLoading,
      error,
    }),
    [currentWorkspace, workspaces, currentUser, switchWorkspace, isLoading, error]
  );

  return <WorkspaceContext.Provider value={value}>{children}</WorkspaceContext.Provider>;
}