Genarrative/src/services/customWorldAgentUiState.ts

import type { CustomWorldAgentUiState } from '../types';

export const CUSTOM_WORLD_AGENT_SESSION_QUERY_KEY = 'customWorldSessionId';
export const CUSTOM_WORLD_AGENT_OPERATION_QUERY_KEY = 'customWorldOperationId';
export const CUSTOM_WORLD_AGENT_UI_STATE_STORAGE_KEY =
  'genarrative.custom-world-agent-ui.v1';

type CustomWorldAgentUiEnvironment = {
  location?: {
    pathname: string;
    search: string;
  } | null;
  history?: {
    replaceState: (
      data: unknown,
      unused: string,
      url?: string | URL | null,
    ) => void;
  } | null;
  sessionStorage?: Pick<Storage, 'getItem' | 'setItem' | 'removeItem'> | null;
};

function resolveEnvironment(
  env?: CustomWorldAgentUiEnvironment,
): Required<CustomWorldAgentUiEnvironment> {
  if (env) {
    return {
      location: env.location ?? null,
      history: env.history ?? null,
      sessionStorage: env.sessionStorage ?? null,
    };
  }

  if (typeof window === 'undefined') {
    return {
      location: null,
      history: null,
      sessionStorage: null,
    };
  }

  return {
    location: window.location,
    history: window.history,
    sessionStorage: window.sessionStorage,
  };
}

function normalizeValue(value: unknown) {
  return typeof value === 'string' && value.trim() ? value.trim() : null;
}

export function readCustomWorldAgentUiState(
  env?: CustomWorldAgentUiEnvironment,
): CustomWorldAgentUiState {
  const resolved = resolveEnvironment(env);
  const params = new URLSearchParams(resolved.location?.search ?? '');
  const stateFromQuery: CustomWorldAgentUiState = {
    activeSessionId: normalizeValue(
      params.get(CUSTOM_WORLD_AGENT_SESSION_QUERY_KEY),
    ),
    activeOperationId: normalizeValue(
      params.get(CUSTOM_WORLD_AGENT_OPERATION_QUERY_KEY),
    ),
  };

  if (stateFromQuery.activeSessionId || stateFromQuery.activeOperationId) {
    return stateFromQuery;
  }

  const storedValue = resolved.sessionStorage?.getItem(
    CUSTOM_WORLD_AGENT_UI_STATE_STORAGE_KEY,
  );
  if (!storedValue) {
    return {};
  }

  try {
    const parsed = JSON.parse(storedValue) as CustomWorldAgentUiState;
    return {
      activeSessionId: normalizeValue(parsed.activeSessionId),
      activeOperationId: normalizeValue(parsed.activeOperationId),
    };
  } catch {
    resolved.sessionStorage?.removeItem(CUSTOM_WORLD_AGENT_UI_STATE_STORAGE_KEY);
    return {};
  }
}

export function writeCustomWorldAgentUiState(
  state: CustomWorldAgentUiState,
  env?: CustomWorldAgentUiEnvironment,
) {
  const resolved = resolveEnvironment(env);
  const activeSessionId = normalizeValue(state.activeSessionId);
  const activeOperationId = normalizeValue(state.activeOperationId);
  const nextState: CustomWorldAgentUiState = {
    activeSessionId,
    activeOperationId,
  };

  if (resolved.location && resolved.history?.replaceState) {
    const params = new URLSearchParams(resolved.location.search);
    if (activeSessionId) {
      params.set(CUSTOM_WORLD_AGENT_SESSION_QUERY_KEY, activeSessionId);
    } else {
      params.delete(CUSTOM_WORLD_AGENT_SESSION_QUERY_KEY);
    }

    if (activeOperationId) {
      params.set(CUSTOM_WORLD_AGENT_OPERATION_QUERY_KEY, activeOperationId);
    } else {
      params.delete(CUSTOM_WORLD_AGENT_OPERATION_QUERY_KEY);
    }

    const search = params.toString();
    const nextUrl = search
      ? `${resolved.location.pathname}?${search}`
      : resolved.location.pathname;
    resolved.history.replaceState(null, '', nextUrl);
  }

  if (resolved.sessionStorage) {
    if (activeSessionId || activeOperationId) {
      resolved.sessionStorage.setItem(
        CUSTOM_WORLD_AGENT_UI_STATE_STORAGE_KEY,
        JSON.stringify(nextState),
      );
    } else {
      resolved.sessionStorage.removeItem(CUSTOM_WORLD_AGENT_UI_STATE_STORAGE_KEY);
    }
  }
}

export function clearCustomWorldAgentUiState(
  env?: CustomWorldAgentUiEnvironment,
) {
  writeCustomWorldAgentUiState({}, env);
}