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

import { useEffect, useRef, useState } from "react";
import {
  type KisRuntimeCredentials,
  useKisRuntimeStore,
} from "@/features/settings/store/use-kis-runtime-store";
import type {
  DashboardRealtimeTradeTick,
  DashboardStockOrderBookResponse,
} from "@/features/trade/types/trade.types";
import {
  buildKisRealtimeMessage,
  parseKisRealtimeOrderbook,
  parseKisRealtimeTickBatch,
} from "@/features/trade/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 TRADE_TR_ID_OVERTIME_EXPECTED = "H0STOAC0";
const TRADE_TR_ID_TOTAL = "H0UNCNT0";
const TRADE_TR_ID_TOTAL_EXPECTED = "H0UNANC0";
const ORDERBOOK_TR_ID = "H0STASP0";
const ORDERBOOK_TR_ID_OVERTIME = "H0STOAA0";
const MAX_TRADE_TICKS = 10;
const WS_DEBUG_STORAGE_KEY = "KIS_WS_DEBUG";

/**
 * @description 장 구간/시장별 누락을 줄이기 위해 TR ID를 우선순위 배열로 반환합니다.
 * @see features/trade/hooks/useKisTradeWebSocket.ts useEffect WebSocket 구독 등록
 * @see temp-kis-domestic-functions-ws.py ccnl_krx/ccnl_total/exp_ccnl_krx/exp_ccnl_total/overtime_ccnl_krx
 */
function resolveTradeTrIds(
  env: KisRuntimeCredentials["tradingEnv"],
  session: DomesticKisSession,
) {
  if (env === "mock") return [TRADE_TR_ID];

  if (shouldUseAfterHoursSinglePriceTr(session)) {
    // 시간외 단일가(16:00~18:00): 전용 TR + 통합 TR 백업
    return uniqueTrIds([
      TRADE_TR_ID_OVERTIME,
      TRADE_TR_ID_OVERTIME_EXPECTED,
      TRADE_TR_ID_TOTAL,
      TRADE_TR_ID_TOTAL_EXPECTED,
    ]);
  }

  if (shouldUseExpectedExecutionTr(session)) {
    // 동시호가 구간(장전/장마감): 예상체결 TR을 우선, 일반체결/통합체결을 백업
    return uniqueTrIds([
      TRADE_TR_ID_EXPECTED,
      TRADE_TR_ID_TOTAL_EXPECTED,
      TRADE_TR_ID,
      TRADE_TR_ID_TOTAL,
    ]);
  }

  if (session === "afterCloseFixedPrice") {
    // 시간외 종가(15:40~16:00): 브로커별 라우팅 차이를 대비해 일반/시간외/통합 TR을 함께 구독
    return uniqueTrIds([
      TRADE_TR_ID,
      TRADE_TR_ID_TOTAL,
      TRADE_TR_ID_OVERTIME,
      TRADE_TR_ID_OVERTIME_EXPECTED,
      TRADE_TR_ID_TOTAL_EXPECTED,
    ]);
  }

  return uniqueTrIds([TRADE_TR_ID, TRADE_TR_ID_TOTAL]);
}

/**
 * @description 장 구간별 호가 TR ID 후보를 우선순위 배열로 반환합니다.
 * @see features/trade/hooks/useKisTradeWebSocket.ts useEffect WebSocket 구독 등록
 * @see temp-kis-domestic-functions-ws.py asking_price_krx/asking_price_total/overtime_asking_price_krx
 */
function resolveOrderBookTrIds(
  env: KisRuntimeCredentials["tradingEnv"],
  session: DomesticKisSession,
) {
  if (env === "mock") return [ORDERBOOK_TR_ID];

  if (shouldUseAfterHoursSinglePriceTr(session)) {
    // 시간외 단일가(16:00~18:00)는 KRX 전용 호가 TR만 구독합니다.
    // 통합 TR(H0UNASP0)을 같이 구독하면 종목별로 포맷/잔량이 섞여 보일 수 있습니다.
    return uniqueTrIds([ORDERBOOK_TR_ID_OVERTIME]);
  }

  if (session === "afterCloseFixedPrice") {
    return uniqueTrIds([ORDERBOOK_TR_ID_OVERTIME]);
  }

  // UI 흐름: 호가창 UI -> useKisTradeWebSocket onmessage -> onOrderBookMessage
  // -> TradeContainer setRealtimeOrderBook -> useOrderBook 병합 -> OrderBook 렌더
  // 장중에는 KRX 전용(H0STASP0)만 구독해 값이 번갈아 덮이는 현상을 방지합니다.
  return uniqueTrIds([ORDERBOOK_TR_ID]);
}

/**
 * @description 콘솔 디버그 플래그를 확인합니다.
 * @see features/trade/hooks/useKisTradeWebSocket.ts socket.onmessage
 */
function isWsDebugEnabled() {
  if (typeof window === "undefined") return false;

  try {
    return window.localStorage.getItem(WS_DEBUG_STORAGE_KEY) === "1";
  } catch {
    return false;
  }
}

/**
 * @description 실시간 웹소켓 제어(JSON) 메시지를 파싱합니다.
 * @see features/trade/hooks/useKisTradeWebSocket.ts socket.onmessage
 */
function parseWsControlMessage(raw: string) {
  if (!raw.startsWith("{")) return null;

  try {
    return JSON.parse(raw) as {
      header?: { tr_id?: string };
      body?: { rt_cd?: string; msg1?: string };
    };
  } catch {
    return null;
  }
}

/**
 * @description 실시간 원문에서 파이프 구분 TR ID를 빠르게 추출합니다.
 * @see features/trade/hooks/useKisTradeWebSocket.ts socket.onmessage
 */
function peekPipeTrId(raw: string) {
  const parts = raw.split("|");
  return parts.length > 1 ? parts[1] : "";
}

function uniqueTrIds(ids: string[]) {
  return [...new Set(ids)];
}

/**
 * @description Subscribes trade ticks and orderbook over one websocket.
 * @see features/trade/components/TradeContainer.tsx
 * @see lib/kis/domestic-market-session.ts
 */
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 realtimeTrIds = credentials
    ? resolveTradeTrIds(credentials.tradingEnv, marketSession)
    : [TRADE_TR_ID];
  const realtimeTrId = credentials
    ? realtimeTrIds[0] ?? TRADE_TR_ID
    : TRADE_TR_ID;

  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(
        "실시간 연결은 되었지만 체결 데이터가 없습니다. 장 시간(한국시간)과 TR ID를 확인해 주세요.",
      );
    }, 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 debugEnabled = isWsDebugEnabled();

    const tradeTrIds = resolveTradeTrIds(credentials.tradingEnv, marketSession);
    const orderBookTrIds =
      obSymbol && onOrderBookMsg
        ? resolveOrderBookTrIds(credentials.tradingEnv, marketSession)
        : [];

    const subscribe = (
      key: string,
      targetSymbol: string,
      trId: string,
      trType: "1" | "2",
    ) => {
      socket?.send(
        JSON.stringify(buildKisRealtimeMessage(key, targetSymbol, trId, trType)),
      );
    };

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

        const wsConnection = await useKisRuntimeStore
          .getState()
          .getOrFetchWsConnection();

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

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

        // 공식 샘플과 동일하게 /tryitout 엔드포인트로 연결하고, TR은 payload로 구독합니다.
        socket = new WebSocket(`${wsConnection.wsUrl}/tryitout`);
        socketRef.current = socket;

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

          for (const trId of tradeTrIds) {
            subscribe(approvalKeyRef.current, symbol, trId, "1");
          }

          if (obSymbol) {
            for (const trId of orderBookTrIds) {
              subscribe(approvalKeyRef.current, obSymbol, trId, "1");
            }
          }

          if (debugEnabled) {
            console.info("[KisRealtime] Subscribed", {
              symbol,
              marketSession,
              tradeTrIds,
              orderBookSymbol: obSymbol ?? null,
              orderBookTrIds,
            });
          }

          setIsConnected(true);
        };

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

          const control = parseWsControlMessage(event.data);
          if (control) {
            const trId = control.header?.tr_id ?? "";
            if (trId === "PINGPONG") {
              // 서버 Keepalive에 응답하지 않으면 연결이 끊길 수 있습니다.
              socket?.send(event.data);
              return;
            }

            if (debugEnabled) {
              console.info("[KisRealtime] Control", {
                trId,
                rt_cd: control.body?.rt_cd,
                message: control.body?.msg1,
              });
            }
            return;
          }

          if (obSymbol && onOrderBookMsg) {
            const orderBook = parseKisRealtimeOrderbook(event.data, obSymbol);
            if (orderBook) {
              orderBook.tradingEnv = credentials.tradingEnv;
              if (debugEnabled) {
                console.debug("[KisRealtime] OrderBook", {
                  trId: peekPipeTrId(event.data),
                  symbol: orderBook.symbol,
                  businessHour: orderBook.businessHour,
                  hourClassCode: orderBook.hourClassCode,
                });
              }
              onOrderBookMsg(orderBook);
              return;
            }
          }

          const ticks = parseKisRealtimeTickBatch(event.data, symbol);
          if (ticks.length === 0) {
            if (debugEnabled && event.data.includes("|")) {
              console.debug("[KisRealtime] Unparsed payload", {
                trId: peekPipeTrId(event.data),
                preview: event.data.slice(0, 220),
              });
            }
            return;
          }

          const meaningfulTicks = ticks.filter((tick) => tick.tradeVolume > 0);
          if (meaningfulTicks.length === 0) {
            if (debugEnabled) {
              console.debug("[KisRealtime] Ignored zero-volume ticks", {
                trId: peekPipeTrId(event.data),
                parsedCount: ticks.length,
              });
            }
            return;
          }

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

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

          if (debugEnabled) {
            console.debug("[KisRealtime] Tick", {
              trId: peekPipeTrId(event.data),
              symbol: latest.symbol,
              tickTime: latest.tickTime,
              price: latest.price,
              tradeVolume: latest.tradeVolume,
              executionClassCode: latest.executionClassCode,
              buyExecutionCount: latest.buyExecutionCount,
              sellExecutionCount: latest.sellExecutionCount,
              netBuyExecutionCount: latest.netBuyExecutionCount,
              parsedCount: ticks.length,
            });
          }

          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) {
            if (debugEnabled) {
              console.warn("[KisRealtime] WebSocket error", {
                symbol,
                marketSession,
                tradeTrIds,
              });
            }
            setIsConnected(false);
          }
        };

        socket.onclose = () => {
          if (!disposed) {
            if (debugEnabled) {
              console.warn("[KisRealtime] WebSocket closed", {
                symbol,
                marketSession,
                tradeTrIds,
              });
            }
            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) {
        for (const trId of tradeTrIds) {
          subscribe(key, symbol, trId, "2");
        }

        if (obSymbol) {
          for (const trId of orderBookTrIds) {
            subscribe(key, obSymbol, trId, "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();
  }
}