auto-trade/features/settings/store/use-kis-runtime-store.ts

"use client";

import { fetchKisWebSocketApproval } from "@/features/settings/apis/kis-auth.api";
import type { KisTradingEnv } from "@/features/trade/types/trade.types";
import { createJSONStorage, persist } from "zustand/middleware";
import { create } from "zustand";

/**
 * @file features/settings/store/use-kis-runtime-store.ts
 * @description Stores KIS input, verification, and websocket connection state.
 * @see features/trade/hooks/useKisTradeWebSocket.ts
 */
export interface KisRuntimeCredentials {
  appKey: string;
  appSecret: string;
  tradingEnv: KisTradingEnv;
  accountNo: string;
}

interface KisWsConnection {
  approvalKey: string;
  wsUrl: string;
}

interface KisRuntimeStoreState {
  kisTradingEnvInput: KisTradingEnv;
  kisAppKeyInput: string;
  kisAppSecretInput: string;
  kisAccountNoInput: string;

  verifiedCredentials: KisRuntimeCredentials | null;
  isKisVerified: boolean;
  isKisProfileVerified: boolean;
  verifiedAccountNo: string | null;
  tradingEnv: KisTradingEnv;

  wsApprovalKey: string | null;
  wsUrl: string | null;

  _hasHydrated: boolean;
}

interface KisRuntimeStoreActions {
  setKisTradingEnvInput: (tradingEnv: KisTradingEnv) => void;
  setKisAppKeyInput: (appKey: string) => void;
  setKisAppSecretInput: (appSecret: string) => void;
  setKisAccountNoInput: (accountNo: string) => void;
  setVerifiedKisSession: (
    credentials: KisRuntimeCredentials,
    tradingEnv: KisTradingEnv,
  ) => void;
  setVerifiedKisProfile: (profile: {
    accountNo: string;
  }) => void;
  invalidateKisProfileVerification: () => void;
  invalidateKisVerification: () => void;
  clearKisRuntimeSession: (tradingEnv: KisTradingEnv) => void;
  getOrFetchWsConnection: () => Promise<KisWsConnection | null>;
  setHasHydrated: (state: boolean) => void;
}

const INITIAL_STATE: KisRuntimeStoreState = {
  kisTradingEnvInput: "real",
  kisAppKeyInput: "",
  kisAppSecretInput: "",
  kisAccountNoInput: "",
  verifiedCredentials: null,
  isKisVerified: false,
  isKisProfileVerified: false,
  verifiedAccountNo: null,
  tradingEnv: "real",
  wsApprovalKey: null,
  wsUrl: null,
  _hasHydrated: false,
};

const RESET_PROFILE_STATE = {
  isKisProfileVerified: false,
  verifiedAccountNo: null,
} as const;

const RESET_VERIFICATION_STATE = {
  verifiedCredentials: null,
  isKisVerified: false,
  ...RESET_PROFILE_STATE,
  wsApprovalKey: null,
  wsUrl: null,
};

let wsConnectionPromise: Promise<KisWsConnection | null> | null = null;

/**
 * @description Runtime store for KIS session.
 * @see features/settings/components/KisAuthForm.tsx
 */
export const useKisRuntimeStore = create<
  KisRuntimeStoreState & KisRuntimeStoreActions
>()(
  persist(
    (set, get) => ({
      ...INITIAL_STATE,

      setKisTradingEnvInput: (tradingEnv) =>
        set({
          kisTradingEnvInput: tradingEnv,
          ...RESET_VERIFICATION_STATE,
        }),

      setKisAppKeyInput: (appKey) =>
        set({
          kisAppKeyInput: appKey,
          ...RESET_VERIFICATION_STATE,
        }),

      setKisAppSecretInput: (appSecret) =>
        set({
          kisAppSecretInput: appSecret,
          ...RESET_VERIFICATION_STATE,
        }),

      setKisAccountNoInput: (accountNo) =>
        set((state) => ({
          kisAccountNoInput: accountNo,
          ...RESET_PROFILE_STATE,
          verifiedCredentials: state.verifiedCredentials
            ? {
                ...state.verifiedCredentials,
                accountNo: "",
              }
            : null,
        })),

      setVerifiedKisSession: (credentials, tradingEnv) =>
        set({
          verifiedCredentials: credentials,
          isKisVerified: true,
          tradingEnv,
          wsApprovalKey: null,
          wsUrl: null,
        }),

      setVerifiedKisProfile: ({ accountNo }) =>
        set((state) => ({
          isKisProfileVerified: true,
          verifiedAccountNo: accountNo,
          verifiedCredentials: state.verifiedCredentials
            ? {
                ...state.verifiedCredentials,
                accountNo,
              }
            : state.verifiedCredentials,
          wsApprovalKey: null,
          wsUrl: null,
        })),

      invalidateKisProfileVerification: () =>
        set((state) => ({
          ...RESET_PROFILE_STATE,
          verifiedCredentials: state.verifiedCredentials
            ? {
                ...state.verifiedCredentials,
                accountNo: "",
              }
            : state.verifiedCredentials,
          wsApprovalKey: null,
          wsUrl: null,
        })),

      invalidateKisVerification: () =>
        set({
          ...RESET_VERIFICATION_STATE,
        }),

      clearKisRuntimeSession: (tradingEnv) =>
        set({
          kisTradingEnvInput: tradingEnv,
          kisAppKeyInput: "",
          kisAppSecretInput: "",
          kisAccountNoInput: "",
          ...RESET_VERIFICATION_STATE,
          tradingEnv,
        }),

      getOrFetchWsConnection: async () => {
        const { wsApprovalKey, wsUrl, verifiedCredentials } = get();

        if (wsApprovalKey && wsUrl) {
          return { approvalKey: wsApprovalKey, wsUrl };
        }

        if (!verifiedCredentials) {
          return null;
        }

        if (wsConnectionPromise) {
          return wsConnectionPromise;
        }

        wsConnectionPromise = (async () => {
          try {
            const data = await fetchKisWebSocketApproval(verifiedCredentials);
            if (!data.approvalKey || !data.wsUrl) {
              return null;
            }

            const nextConnection = {
              approvalKey: data.approvalKey,
              wsUrl: data.wsUrl,
            } satisfies KisWsConnection;

            set({
              wsApprovalKey: nextConnection.approvalKey,
              wsUrl: nextConnection.wsUrl,
            });

            return nextConnection;
          } catch (error) {
            console.error(error);
            return null;
          } finally {
            wsConnectionPromise = null;
          }
        })();

        return wsConnectionPromise;
      },
      setHasHydrated: (state) => {
        set({
          _hasHydrated: state,
        });
      },
    }),
    {
      name: "autotrade-kis-runtime-store",
      storage: createJSONStorage(() => localStorage),
      onRehydrateStorage: () => (state) => {
        state?.setHasHydrated(true);
      },
      partialize: (state) => ({
        kisTradingEnvInput: state.kisTradingEnvInput,
        kisAppKeyInput: state.kisAppKeyInput,
        kisAppSecretInput: state.kisAppSecretInput,
        kisAccountNoInput: state.kisAccountNoInput,
        verifiedCredentials: state.verifiedCredentials,
        isKisVerified: state.isKisVerified,
        isKisProfileVerified: state.isKisProfileVerified,
        verifiedAccountNo: state.verifiedAccountNo,
        tradingEnv: state.tradingEnv,
        // wsApprovalKey/wsUrl are kept in memory only (expiration-sensitive).
      }),
    },
  ),
);