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

import { useEffect, useRef, useState } from "react";
import {
  type KisRuntimeCredentials,
  useKisRuntimeStore,
} from "@/features/dashboard/store/use-kis-runtime-store";
import type {
  DashboardRealtimeTradeTick,
  DashboardStockOrderBookResponse,
} from "@/features/dashboard/types/dashboard.types";
import {
  buildKisRealtimeMessage,
  parseKisRealtimeOrderbook,
  parseKisRealtimeTickBatch,
} from "@/features/dashboard/utils/kis-realtime.utils";
import {
  DOMESTIC_KIS_SESSION_OVERRIDE_STORAGE_KEY,
  resolveDomesticKisSession,
  shouldUseAfterHoursSinglePriceTr,
  shouldUseExpectedExecutionTr,
  type DomesticKisSession,
} from "@/lib/kis/domestic-market-session";

const TRADE_TR_ID = "H0STCNT0";
const TRADE_TR_ID_EXPECTED = "H0STANC0";
const TRADE_TR_ID_OVERTIME = "H0STOUP0";
const ORDERBOOK_TR_ID = "H0STASP0";
const ORDERBOOK_TR_ID_OVERTIME = "H0STOAA0";
const MAX_TRADE_TICKS = 10;

function resolveTradeTrId(
  env: KisRuntimeCredentials["tradingEnv"],
  session: DomesticKisSession,
) {
  if (env === "mock") return TRADE_TR_ID;
  if (shouldUseAfterHoursSinglePriceTr(session)) return TRADE_TR_ID_OVERTIME;
  if (shouldUseExpectedExecutionTr(session)) return TRADE_TR_ID_EXPECTED;
  return TRADE_TR_ID;
}

function resolveOrderBookTrId(
  env: KisRuntimeCredentials["tradingEnv"],
  session: DomesticKisSession,
) {
  if (env === "mock") return ORDERBOOK_TR_ID;
  if (shouldUseAfterHoursSinglePriceTr(session)) return ORDERBOOK_TR_ID_OVERTIME;
  return ORDERBOOK_TR_ID;
}

/**
 * @description KIS 실시간 체결/호가를 단일 WebSocket으로 구독합니다.
 * @see features/dashboard/components/DashboardContainer.tsx useKisTradeWebSocket 호출
 * @see lib/kis/domestic-market-session.ts 장 세션 계산 및 테스트용 override
 */
export function useKisTradeWebSocket(
  symbol: string | undefined,
  credentials: KisRuntimeCredentials | null,
  isVerified: boolean,
  onTick?: (tick: DashboardRealtimeTradeTick) => void,
  options?: {
    orderBookSymbol?: string;
    onOrderBookMessage?: (data: DashboardStockOrderBookResponse) => void;
  },
) {
  const [latestTick, setLatestTick] =
    useState<DashboardRealtimeTradeTick | null>(null);
  const [recentTradeTicks, setRecentTradeTicks] = useState<
    DashboardRealtimeTradeTick[]
  >([]);
  const [isConnected, setIsConnected] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [lastTickAt, setLastTickAt] = useState<number | null>(null);
  const [marketSession, setMarketSession] = useState<DomesticKisSession>(() =>
    resolveSessionInClient(),
  );

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

  const obSymbol = options?.orderBookSymbol;
  const onOrderBookMsg = options?.onOrderBookMessage;
  const realtimeTrId = credentials
    ? resolveTradeTrId(credentials.tradingEnv, marketSession)
    : TRADE_TR_ID;

  // KST 장 세션을 주기적으로 재평가합니다.
  useEffect(() => {
    const timerId = window.setInterval(() => {
      const nextSession = resolveSessionInClient();
      setMarketSession((prev) => (prev === nextSession ? prev : nextSession));
    }, 30_000);

    return () => window.clearInterval(timerId);
  }, []);

  // 연결은 되었는데 체결이 오래 안 들어오는 경우 안내합니다.
  useEffect(() => {
    if (!isConnected || lastTickAt) return;

    const timer = window.setTimeout(() => {
      setError(
        "실시간 연결은 되었지만 체결 데이터가 없습니다. 장 시간(한국시간) 여부를 확인해 주세요.",
      );
    }, 8000);

    return () => window.clearTimeout(timer);
  }, [isConnected, lastTickAt]);

  useEffect(() => {
    setLatestTick(null);
    setRecentTradeTicks([]);
    setError(null);
    setLastTickAt(null);
    seenTickRef.current.clear();

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

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

    const tradeTrId = resolveTradeTrId(credentials.tradingEnv, marketSession);
    const orderBookTrId = obSymbol
      ? resolveOrderBookTrId(credentials.tradingEnv, marketSession)
      : null;

    const connect = async () => {
      try {
        setError(null);
        setIsConnected(false);

        const approvalKey = await useKisRuntimeStore
          .getState()
          .getOrFetchApprovalKey();

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

        if (disposed) return;
        approvalKeyRef.current = approvalKey;

        const wsBase =
          process.env.NEXT_PUBLIC_KIS_WS_URL ||
          "ws://ops.koreainvestment.com:21000";

        socket = new WebSocket(`${wsBase}/tryitout/${tradeTrId}`);
        socketRef.current = socket;

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

          socket?.send(
            JSON.stringify(
              buildKisRealtimeMessage(
                approvalKeyRef.current,
                symbol,
                tradeTrId,
                "1",
              ),
            ),
          );

          if (obSymbol && orderBookTrId) {
            socket?.send(
              JSON.stringify(
                buildKisRealtimeMessage(
                  approvalKeyRef.current,
                  obSymbol,
                  orderBookTrId,
                  "1",
                ),
              ),
            );
          }

          setIsConnected(true);
        };

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

          if (obSymbol && onOrderBookMsg) {
            const orderBook = parseKisRealtimeOrderbook(event.data, obSymbol);
            if (orderBook) {
              orderBook.tradingEnv = credentials.tradingEnv;
              onOrderBookMsg(orderBook);
              return;
            }
          }

          const ticks = parseKisRealtimeTickBatch(event.data, symbol);
          if (ticks.length === 0) return;

          const meaningfulTicks = ticks.filter((tick) => tick.tradeVolume > 0);
          const dedupedTicks = meaningfulTicks.filter((tick) => {
            const key = `${tick.tickTime}-${tick.price}-${tick.tradeVolume}`;
            if (seenTickRef.current.has(key)) return false;
            seenTickRef.current.add(key);
            return true;
          });

          const latest = ticks[ticks.length - 1];
          setLatestTick(latest);

          if (dedupedTicks.length > 0) {
            setRecentTradeTicks((prev) =>
              [...dedupedTicks.reverse(), ...prev].slice(0, MAX_TRADE_TICKS),
            );
          }

          setError(null);
          setLastTickAt(Date.now());
          onTick?.(latest);
        };

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

        socket.onclose = () => {
          if (!disposed) setIsConnected(false);
        };
      } catch (err) {
        if (disposed) return;

        setError(
          err instanceof Error
            ? err.message
            : "실시간 웹소켓 초기화 중 오류가 발생했습니다.",
        );
        setIsConnected(false);
      }
    };

    void connect();
    const seenRef = seenTickRef.current;

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

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

        if (obSymbol && orderBookTrId) {
          socket.send(
            JSON.stringify(
              buildKisRealtimeMessage(key, obSymbol, orderBookTrId, "2"),
            ),
          );
        }
      }

      socket?.close();
      if (socketRef.current === socket) socketRef.current = null;
      approvalKeyRef.current = null;
      seenRef.clear();
    };
  }, [
    symbol,
    isVerified,
    credentials,
    marketSession,
    onTick,
    obSymbol,
    onOrderBookMsg,
  ]);

  return {
    latestTick,
    recentTradeTicks,
    isConnected,
    error,
    lastTickAt,
    realtimeTrId,
  };
}

function resolveSessionInClient() {
  if (typeof window === "undefined") {
    return resolveDomesticKisSession();
  }

  try {
    const override = window.localStorage.getItem(
      DOMESTIC_KIS_SESSION_OVERRIDE_STORAGE_KEY,
    );
    return resolveDomesticKisSession(override);
  } catch {
    return resolveDomesticKisSession();
  }
}