auto-trade/features/dashboard/hooks/useKisOrderbookWebSocket.ts

import { useEffect, useRef, useState } from "react";
import {
  type KisRuntimeCredentials,
  useKisRuntimeStore,
} from "@/features/dashboard/store/use-kis-runtime-store";
import type { DashboardStockOrderBookResponse } from "@/features/dashboard/types/dashboard.types";
import {
  buildKisRealtimeMessage,
  parseKisRealtimeOrderbook,
} from "@/features/dashboard/utils/kis-realtime.utils";

const KRX_ORDERBOOK_TR_ID = "H0STASP0";
const KRX_OVERTIME_ORDERBOOK_TR_ID = "H0STOAA0";
const DEFAULT_ORDERBOOK_TR_ID = "H0UNASP0";

/**
 * @description 한국 시간 기준으로 시간외 단일가 구간(16:00~18:00)인지 확인합니다.
 * @see resolveOrderBookTrId 시간대별 TR ID 선택
 */
function isKrxOvertimeInKst(now = new Date()) {
  const formatter = new Intl.DateTimeFormat("en-US", {
    timeZone: "Asia/Seoul",
    weekday: "short",
    hour: "2-digit",
    minute: "2-digit",
    hour12: false,
  });
  const parts = formatter.formatToParts(now);
  const partMap = new Map(parts.map((part) => [part.type, part.value]));
  const weekday = partMap.get("weekday");

  if (weekday === "Sat" || weekday === "Sun") {
    return false;
  }

  const hour = Number(partMap.get("hour") ?? "0");
  const minute = Number(partMap.get("minute") ?? "0");
  const totalMinutes = hour * 60 + minute;

  return totalMinutes >= 16 * 60 && totalMinutes < 18 * 60;
}

/**
 * @description 시장/시간대에 맞는 국내주식 호가 TR ID를 선택합니다.
 * @see .tmp/open-trading-api/examples_user/domestic_stock/domestic_stock_functions_ws.py
 */
function resolveOrderBookTrId(market: "KOSPI" | "KOSDAQ" | undefined) {
  if (market === "KOSPI" || market === "KOSDAQ") {
    return isKrxOvertimeInKst()
      ? KRX_OVERTIME_ORDERBOOK_TR_ID
      : KRX_ORDERBOOK_TR_ID;
  }

  return DEFAULT_ORDERBOOK_TR_ID;
}

/**
 * @description KIS 실시간 호가 웹소켓 훅
 * @see parseKisRealtimeOrderbook 웹소켓 payload 파싱
 */
export function useKisOrderbookWebSocket(
  symbol: string | undefined,
  market: "KOSPI" | "KOSDAQ" | undefined,
  credentials: KisRuntimeCredentials | null,
  isVerified: boolean,
) {
  const [realtimeOrderBook, setRealtimeOrderBook] =
    useState<DashboardStockOrderBookResponse | null>(null);
  const [isConnected, setIsConnected] = useState(false);
  const [error, setError] = useState<string | null>(null);

  const socketRef = useRef<WebSocket | null>(null);
  const approvalKeyRef = useRef<string | null>(null);

  const trId = resolveOrderBookTrId(market);

  useEffect(() => {
    setRealtimeOrderBook(null);
    setError(null);

    if (!symbol || !isVerified || !credentials) {
      socketRef.current?.close();
      socketRef.current = null;
      approvalKeyRef.current = null;
      return;
    }

    let disposed = false;
    let socket: WebSocket | null = null;

    const connect = async () => {
      try {
        console.log("[OrderBook WS] 연결 시작", { symbol, trId, market });
        const approvalKey = await useKisRuntimeStore
          .getState()
          .getOrFetchApprovalKey();

        if (!approvalKey) {
          throw new Error("웹소켓 승인키 발급에 실패했습니다.");
        }
        if (disposed) return;

        approvalKeyRef.current = approvalKey;

        const wsBaseUrl =
          process.env.NEXT_PUBLIC_KIS_WS_URL ||
          "ws://ops.koreainvestment.com:21000";
        const wsUrl = `${wsBaseUrl}/tryitout/${trId}`;
        socket = new WebSocket(wsUrl);
        socketRef.current = socket;

        socket.onopen = () => {
          if (disposed || !approvalKeyRef.current) return;

          const subscribeMessage = buildKisRealtimeMessage(
            approvalKeyRef.current,
            symbol,
            trId,
            "1",
          );
          socket?.send(JSON.stringify(subscribeMessage));
          setIsConnected(true);
        };

        socket.onmessage = (event) => {
          if (disposed || typeof event.data !== "string") return;

          const orderBook = parseKisRealtimeOrderbook(event.data, symbol);
          if (!orderBook) {
            return;
          }

          orderBook.tradingEnv = credentials.tradingEnv;
          setRealtimeOrderBook(orderBook);
        };

        socket.onerror = () => {
          if (disposed) return;
          setIsConnected(false);
        };

        socket.onclose = () => {
          if (disposed) return;
          setIsConnected(false);
        };
      } catch (err) {
        if (disposed) return;
        setError(
          err instanceof Error
            ? err.message
            : "실시간 호가 웹소켓 초기화 중 오류가 발생했습니다.",
        );
        setIsConnected(false);
      }
    };

    void connect();

    return () => {
      disposed = true;
      setIsConnected(false);

      const approvalKey = approvalKeyRef.current;
      if (socket?.readyState === WebSocket.OPEN && approvalKey) {
        const unsubscribeMessage = buildKisRealtimeMessage(
          approvalKey,
          symbol,
          trId,
          "2",
        );
        socket.send(JSON.stringify(unsubscribeMessage));
      }

      socket?.close();
      if (socketRef.current === socket) {
        socketRef.current = null;
      }
      approvalKeyRef.current = null;
    };
  }, [isVerified, symbol, market, credentials, trId]);

  return {
    realtimeOrderBook,
    isConnected,
    error,
  };
}