"""
HTS 도메인 특화 규칙 기반 쿼리 확장기 (QueryExpander)

- LLM 없음 → 10ms 이내 실행
- HTS 핵심 용어 동의어 사전으로 쿼리 확장
- BM25 검색에서 사용 (Dense 임베딩에는 사용 안 함)

사용 예:
    expander = QueryExpander()
    expander.expand("주문 취소")
    # → "주문 취소 매수 매도 거래 체결 미체결 정정 주문변경"
"""

from typing import List

# ---------------------------------------------------------------------------
# HTS 도메인 동의어 사전
# 키: 원본 단어, 값: 함께 추가할 관련어 리스트 (원본 포함)
# ---------------------------------------------------------------------------
SYNONYM_MAP = {
    # ── 주문 계열 ──────────────────────────────────────────────────────────
    "취소":     ["취소", "정정", "주문변경", "미체결", "주문취소", "정정취소"],
    "정정":     ["정정", "취소", "수정", "변경", "주문변경", "주문수정"],
    "주문":     ["주문", "매수", "매도", "거래", "체결", "미체결"],
    "매수":     ["매수", "사다", "구매", "주문", "체결"],
    "매도":     ["매도", "팔다", "판매", "주문", "체결"],
    "사고":     ["매수", "사다", "구매", "주문"],
    "팔고":     ["매도", "팔다", "판매", "주문"],
    "예약":     ["예약", "예약주문", "자동주문", "시간지정", "걸어놓다"],
    "일괄":     ["일괄", "묶음", "동시", "한꺼번에", "대량주문"],
    "자동":     ["자동", "예약", "예약주문", "자동주문", "걸어놓다"],
    "걸어":     ["예약", "예약주문", "자동주문", "걸어놓다"],

    # ── 잔고 / 손익 계열 ───────────────────────────────────────────────────
    "잔고":     ["잔고", "보유", "수량", "평가금액", "자산", "보유주식"],
    "손익":     ["손익", "수익", "수익률", "평가손익", "실현손익"],
    "수익":     ["수익", "손익", "수익률", "평가손익", "이익", "수익현황"],
    "얼마":     ["잔고", "평가금액", "수익률", "손익", "잔고조회", "수익", "보유"],
    "올랐":     ["손익", "수익률", "평가손익", "수익", "잔고"],
    "내렸":     ["손익", "평가손익", "손실", "수익률"],
    "보유":     ["보유", "잔고", "내주식", "포트폴리오", "보유주식"],
    "주식":     ["주식", "잔고", "보유주식", "주식잔고", "국내주식잔고"],
    "확인":     ["확인", "조회", "잔고확인", "현황", "조회하다"],
    "내주식":   ["잔고", "보유주식", "국내잔고", "수익", "평가금액"],

    # ── 시세 계열 ──────────────────────────────────────────────────────────
    "시세":     ["시세", "현재가", "가격", "주가", "호가", "실시간"],
    "현재가":   ["현재가", "시세", "실시간", "주가", "체결가"],
    "가격":     ["현재가", "시세", "주가", "호가", "체결가"],
    "호가":     ["호가", "체결", "호가창", "매수호가", "매도호가"],
    "체결":     ["체결", "호가", "체결내역", "체결가", "거래"],
    # "지금"은 자연어 시간 강조 표현 — 현재가/시세로 확장 시 잔고 쿼리에서 역효과
    # 예: "지금 내 주식 얼마야" → 호가로 오매칭. BM25 자체 IDF로 처리.
    # "지금":     ["현재가", "실시간", "시세", "주가"],  # 비활성화

    # ── 차트 계열 ──────────────────────────────────────────────────────────
    "차트":     ["차트", "일봉", "주봉", "캔들", "그래프", "시세차트"],
    "일봉":     ["일봉", "차트", "일별시세", "캔들차트", "주봉"],
    "그래프":   ["차트", "그래프", "일봉", "캔들"],

    # ── 조건검색 계열 ──────────────────────────────────────────────────────
    "조건":     ["조건", "조건검색", "스크리닝", "필터", "검색식"],
    "스크리닝": ["스크리닝", "조건검색", "종목검색", "필터", "조건"],
    "검색식":   ["검색식", "조건검색", "조건식", "스크리닝"],
    "실시간":   ["실시간", "조건검색", "현재가", "시세"],

    # ── 자산 / 계좌 계열 ───────────────────────────────────────────────────
    "이체":     ["이체", "계좌이체", "자금이동", "송금", "계좌간이체"],
    "계좌":     ["계좌", "잔고", "자산", "통합조회", "계좌잔고"],
    "자산":     ["자산", "계좌", "잔고", "포트폴리오", "통합조회"],
    "통합":     ["통합", "통합조회", "계좌잔고", "자산", "전체"],
    "송금":     ["이체", "계좌이체", "자금이동", "송금"],

    # ── 투자정보 계열 ──────────────────────────────────────────────────────
    "공시":     ["공시", "전자공시", "공시검색", "사업보고서", "기업공시"],
    "뉴스":     ["뉴스", "종목뉴스", "시장정보", "이슈", "뉴스검색"],
    "정보":     ["투자정보", "뉴스", "공시", "시장정보"],

    # ── 설정 계열 ──────────────────────────────────────────────────────────
    "알림":     ["알림", "푸시", "푸시알림", "알림설정", "알림끄기"],
    "푸시":     ["푸시", "알림", "푸시알림", "알림설정"],
    "설정":     ["설정", "환경설정", "화면설정", "알림설정"],
    "끄기":     ["해제", "알림해제", "알림끄기", "알림설정"],

    # ── 해외주식 계열 ──────────────────────────────────────────────────────
    "해외":     ["해외", "미국", "해외주식", "글로벌", "나스닥"],
    "미국":     ["미국", "해외", "나스닥", "NYSE", "미국주식", "해외주식"],
    "나스닥":   ["나스닥", "미국", "해외주식", "NYSE"],

    # ── 선물/옵션 계열 ────────────────────────────────────────────────────
    "선물":     ["선물", "선물주문", "선물옵션", "파생", "futures"],
    "옵션":     ["옵션", "선물옵션", "파생상품", "선물"],
}


class QueryExpander:
    """
    짧은 HTS 검색 쿼리를 도메인 동의어로 확장.

    BM25 검색 전처리에 사용. Dense 임베딩에는 사용하지 않음.

    Args:
        max_extra_terms: 추가할 최대 동의어 수 (기본 8개)
                         너무 많으면 BM25 IDF 점수 희석

    사용 예::
        expander = QueryExpander()
        print(expander.expand("주문 취소하고 싶어"))
        # "주문 취소하고 싶어 매수 매도 거래 정정 주문변경 미체결"

        print(expander.get_expansion_terms("차트 보여줘"))
        # ["일봉", "주봉", "캔들", "그래프", "시세차트"]
    """

    def __init__(self, max_extra_terms: int = 8):
        self.synonym_map = SYNONYM_MAP
        self.max_extra_terms = max_extra_terms

    def expand(self, query: str) -> str:
        """
        쿼리에서 등장하는 키워드를 동의어로 확장.

        원본 쿼리를 앞에 유지하고, 동의어를 뒤에 추가.
        공백 기반 토큰 단위로 처리 (형태소 분석 없음).

        Returns:
            확장된 쿼리 문자열
        """
        tokens = query.strip().split()
        extra_terms: List[str] = []
        seen = set(tokens)

        for token in tokens:
            if token in self.synonym_map:
                for synonym in self.synonym_map[token]:
                    if synonym not in seen:
                        extra_terms.append(synonym)
                        seen.add(synonym)
                        if len(extra_terms) >= self.max_extra_terms:
                            break
            if len(extra_terms) >= self.max_extra_terms:
                break

        if extra_terms:
            return query + " " + " ".join(extra_terms)
        return query

    def get_expansion_terms(self, query: str) -> List[str]:
        """확장 용어 목록만 반환 (디버깅/로깅용)"""
        original_tokens = set(query.strip().split())
        expanded = self.expand(query)
        all_tokens = expanded.split()
        return [t for t in all_tokens if t not in original_tokens]