src/vector_db_client.py

"""
벡터 데이터베이스 관리 모듈
"""

import os
import json
import logging
import re
import numpy as np
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Dict, List, Any, Optional, Union, Tuple
from datetime import datetime
from tqdm import tqdm

import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer

from src import config
from src.utils import (
    VectorDBError,
    DataProcessingError,
    get_latest_data_files,
    load_json_data,
    determine_optimal_chunk_size,
    split_text_into_chunks,
    check_index_status,
    structured_chunking
)
from src.bm25_client import BM25Client
from src.standard_codes import (
    STANDARD_CODE_PATTERN,
    canonical_doc_type,
    collection_for_doc_type,
    collection_names_for_doc_types,
    doc_type_matches,
    normalize_standard_code,
    standard_family_name,
)
from src.ranking_rules import score_ranking_rules

logger = config.setup_logger("vector_db")

STANDARD_CODE_LABEL_PATTERN = re.compile(r"(?:표준코드|기준코드)\W*[:：]?\s*", re.IGNORECASE)


def _normalize_standard_code(value: Any) -> str:
    return normalize_standard_code(value)


def _extract_standard_codes(value: Any) -> List[str]:
    codes: List[str] = []
    seen = set()
    for match in STANDARD_CODE_PATTERN.finditer(str(value or "")):
        code = f"{match.group(1).upper()} {match.group(2)} {match.group(3)} {match.group(4)}"
        if code not in seen:
            codes.append(code)
            seen.add(code)
    return codes


def _extract_labeled_standard_code(value: Any) -> str:
    text = str(value or "")
    label_match = STANDARD_CODE_LABEL_PATTERN.search(text)
    if not label_match:
        return ""
    codes = _extract_standard_codes(text[label_match.end(): label_match.end() + 100])
    return codes[0] if codes else ""


class KCSCVectorDB:
    """
    ChromaDB를 활용한 벡터 데이터베이스 관리
    
    한국건설기준센터(KCSC) 데이터를 임베딩하고 검색하는 기능 제공
    """
    
    _instance = None
    _model = None
    
    def __new__(cls, *args, **kwargs):
        """싱글턴 패턴 구현"""
        if cls._instance is None:
            cls._instance = super(KCSCVectorDB, cls).__new__(cls)
        return cls._instance
    
    def __init__(self, data_dir=None, db_dir=None, embedding_model=None):
        """
        벡터 데이터베이스 초기화
        
        Args:
            data_dir: 데이터 디렉토리 경로
            db_dir: 벡터 DB 저장 경로
            embedding_model: 임베딩 모델 이름 또는 경로
        """
        # 이미 초기화된 경우 중복 초기화 방지
        if hasattr(self, 'initialized') and self.initialized:
            return
            
        self.data_dir = data_dir or config.DATA_DIR
        self.db_dir = db_dir or config.VECTOR_DB_DIR
        
        # 디렉토리 생성
        os.makedirs(self.data_dir, exist_ok=True)
        os.makedirs(self.db_dir, exist_ok=True)
        
        # 임베딩 모델 초기화
        self._init_embedding_model(embedding_model or config.EMBEDDING_MODEL)
        
        # BM25 클라이언트 초기화
        self.bm25_client = BM25Client()
        self.bm25_client.load_index("kcsc_bm25")
        
        # ChromaDB 클라이언트 초기화
        logger.info(f"ChromaDB 초기화 중... (경로: {self.db_dir})")
        
        try:
            self.client = chromadb.PersistentClient(
                path=self.db_dir,
                settings=Settings(
                    anonymized_telemetry=False,
                    allow_reset=True
                )
            )
            
            # 컬렉션 초기화
            self.collections = {}
            self._init_collections()
            
            self.initialized = True
            logger.info("벡터 데이터베이스 초기화 완료")
            
        except Exception as e:
            logger.error(f"ChromaDB 초기화 실패: {str(e)}")
            raise VectorDBError(f"벡터 데이터베이스 초기화 실패: {str(e)}", "init")
    
    @classmethod
    def _init_embedding_model(cls, model_name):
        """
        임베딩 모델 초기화 (클래스 메서드로 싱글턴 구현)
        
        Args:
            model_name: 임베딩 모델 이름 또는 경로
        """
        if cls._model is None:
            try:
                logger.info(f"임베딩 모델 로드 중... (모델: {model_name})")
                cls._model = SentenceTransformer(model_name)
                logger.info("임베딩 모델 로드 완료")
            except Exception as e:
                logger.error(f"임베딩 모델 로드 실패: {str(e)}")
                raise VectorDBError(f"임베딩 모델 '{model_name}' 로드 실패: {str(e)}", "model_load")
    
    def _init_collections(self):
        """컬렉션 초기화 또는 로드"""
        for doc_type in ["KDS", "KCS"]:
            try:
                self.collections[doc_type] = self.client.get_or_create_collection(
                    name=doc_type,
                    metadata={"description": f"{doc_type} 설계기준 문서"}
                )
                count = self.collections[doc_type].count()
                logger.info(f"{doc_type} 컬렉션 초기화 완료 (항목 수: {count})")
            except Exception as e:
                logger.error(f"{doc_type} 컬렉션 초기화 실패: {str(e)}")
                raise VectorDBError(f"{doc_type} 컬렉션 초기화 실패: {str(e)}", "collection_init")
    
    def process_document(self, doc: Dict[str, Any], doc_type: str) -> Optional[Dict[str, Any]]:
        """
        문서를 처리하여 인덱싱 가능한 형태로 변환
        
        Args:
            doc: 문서 데이터
            doc_type: 문서 유형
            
        Returns:
            처리된 문서 데이터 또는 None (내용이 없는 경우)
        """
        # 기본 필드 설정
        processed = {
            "id": f"{doc_type}-{doc.get('Code', '')}-{doc.get('No', '0')}",
            "code": doc.get('Code', ''),
            "full_code": doc.get('FullCode', ''),
            "name": doc.get('Name', ''),
            "title": doc.get('Title', ''),
            "content": doc.get('Contents', ''),
            "version": doc.get('Version', ''),
            "update_date": doc.get('UpdateDate', ''),
            "doc_type": doc_type
        }
        
        # 내용이 있는 경우에만 반환
        if processed["content"]:
            return processed
        
        logger.warning(f"문서 {processed['id']}에 내용이 없습니다.")
        return None
    
    def create_document_chunks(
        self, 
        doc: Dict[str, Any], 
        chunk_size: int = None, 
        overlap: int = None
    ) -> List[Dict[str, Any]]:
        """
        문서를 청크로 나누기
        
        Args:
            doc: 처리된 문서 데이터
            chunk_size: 청크 크기 (None이면 자동 결정)
            overlap: 청크 간 겹치는 문자 수
            
        Returns:
            청크 목록
        """
        content = doc.get("content", "")
        if not content:
            return [doc]
        
        # 청크 크기 결정
        if chunk_size is None:
            chunk_size = determine_optimal_chunk_size(content)
        
        overlap = overlap or config.CHUNK_OVERLAP
        
        # 청크로 나누기
        if len(content) <= chunk_size:
            # 작은 문서는 그대로 사용
            doc["chunk_id"] = f"{doc['id']}-0"
            return [doc]
        
        # 청크 생성
        text_chunks = split_text_into_chunks(content, chunk_size, overlap)
        
        chunks = []
        for i, chunk_content in enumerate(text_chunks):
            chunk = doc.copy()
            chunk["content"] = chunk_content
            chunk["chunk_id"] = f"{doc['id']}-{i}//{len(text_chunks)}"
            chunks.append(chunk)
        
        return chunks
    
    def index_documents(
        self, 
        doc_type: str, 
        documents: List[Dict[str, Any]], 
        chunk_size: int = None,
        overlap: int = None,
        batch_size: int = 100,
        reset: bool = True
    ) -> Tuple[int, int]:
        """
        문서를 인덱싱
        
        Args:
            doc_type: 문서 유형
            documents: 문서 목록
            chunk_size: 청크 크기
            overlap: 청크 간 겹치는 문자 수
            batch_size: 배치 크기
            reset: 기존 데이터 리셋 여부
            
        Returns:
            (인덱싱된 문서 수, 청크 수) 튜플
        """
        if doc_type not in self.collections:
            raise VectorDBError(f"{doc_type} 컬렉션이 없습니다.", "collection_missing")
        
        collection = self.collections[doc_type]
        
        # 청크 생성
        chunks = []
        processed_docs = 0
        
        for doc in documents:
            processed = self.process_document(doc, doc_type)
            if not processed:
                continue
                
            # 문서를 청크로 나누기
            doc_chunks = self.create_document_chunks(processed, chunk_size, overlap)
            chunks.extend(doc_chunks)
            processed_docs += 1
        
        logger.info(f"{doc_type}: {processed_docs}개 문서에서 {len(chunks)}개 청크 생성됨")
        
        # 컬렉션 리셋 (필요시)
        if reset and collection.count() > 0:
            logger.info(f"{doc_type} 컬렉션 초기화 중... (기존 {collection.count()}개 항목)")
            collection.delete(where={})
        
        # 청크 데이터 준비
        if not chunks:
            logger.warning(f"{doc_type}: 인덱싱할 청크가 없습니다.")
            return 0, 0
        
        # 배치 처리
        total_batches = (len(chunks) + batch_size - 1) // batch_size
        
        for batch_idx in range(total_batches):
            start_idx = batch_idx * batch_size
            end_idx = min((batch_idx + 1) * batch_size, len(chunks))
            batch_chunks = chunks[start_idx:end_idx]
            
            try:
                ids = []
                documents = []
                metadatas = []
                embeddings = []
                
                for chunk in batch_chunks:
                    # 임베딩할 텍스트 준비
                    text = f"{chunk['name']} {chunk['title'] or ''}: {chunk['content']}"
                    
                    # 메타데이터 준비
                    metadata = {
                        "id": chunk["id"],
                        "code": chunk["code"],
                        "full_code": chunk["full_code"],
                        "name": chunk["name"],
                        "title": chunk["title"],
                        "doc_type": chunk["doc_type"],
                        "version": chunk["version"]
                    }
                    
                    # 임베딩 생성
                    if self._model is None:
                        raise VectorDBError("임베딩 모델이 초기화되지 않았습니다.", "model_missing")
                        
                    embedding = self._model.encode(text)
                    
                    ids.append(chunk["chunk_id"])
                    documents.append(text)
                    metadatas.append(metadata)
                    embeddings.append(embedding.tolist())
                
                # 데이터 추가
                collection.add(
                    ids=ids,
                    documents=documents,
                    metadatas=metadatas,
                    embeddings=embeddings
                )
                
                logger.info(f"{doc_type} 배치 {batch_idx+1}/{total_batches} 인덱싱 완료 ({len(ids)}개 항목)")
                
            except Exception as e:
                logger.error(f"{doc_type} 배치 {batch_idx+1} 인덱싱 실패: {str(e)}")
                raise VectorDBError(f"{doc_type} 배치 인덱싱 실패: {str(e)}", "batch_indexing")
        
        return processed_docs, len(chunks)
    
    def build_hybrid_index_from_standards(self, file_path: str = None, chunk_size: int = 4000, limit: int = None) -> Dict[str, Any]:
        """
        standards.json을 로드하여 필드 기반 청킹 후 하이브리드 인덱스 구축.

        전략: scope / materials / construction / quality / safety 5개 핵심 필드의
              발췌문을 합쳐 문서당 1개 대표 청크를 만든다.
        """
        file_path = file_path or os.path.join(config.DATA_DIR, "raw", "standards.json")

        if not os.path.exists(file_path):
            logger.warning(f"표준 데이터 파일을 찾을 수 없습니다: {file_path}")
            return {"status": "error", "message": "표준 데이터 파일 없음"}

        logger.info(f"표준 데이터 로드 중... ({file_path})")
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                standards_data = json.load(f)
        except Exception as e:
            return {"status": "error", "message": f"로드 실패: {e}"}

        items = list(standards_data.values()) if isinstance(standards_data, dict) else standards_data
        if limit:
            items = items[:limit]

        # 문서당 1개 대표 청크 전략:
        # standard_code + title + 각 핵심 필드 앞 N자를 합산 → 하나의 검색 가능한 단위 생성
        # 목표: 5,233 문서 = 5,233 청크 → 임베딩 배치 11개 → ~2분 완료
        FIELD_EXCERPT = [
            ("scope",        "적용범위",  600),
            ("materials",    "자재",      600),
            ("construction", "시공",      600),
            ("quality",      "품질관리",  400),
            ("safety",       "안전",      400),
        ]
        SKIP_VALUES = {"내용 없음", "", None}

        all_chunks: List[Dict[str, Any]] = []
        processed_docs = 0
        fallback_code_docs = []
        skipped_docs = []
        failed_docs = []

        def resolve_standard_code(doc: Dict[str, Any], doc_idx: int) -> Tuple[str, str]:
            raw_code = doc.get("standard_code") or doc.get("code") or ""
            code = _normalize_standard_code(raw_code)
            if code:
                return code, "standard_code"

            title_codes = _extract_standard_codes(doc.get("title", ""))
            if len(title_codes) == 1:
                return title_codes[0], "title"
            if len(title_codes) > 1:
                failed_docs.append({
                    "doc_idx": doc_idx,
                    "title": str(doc.get("title", ""))[:160],
                    "reason": "multiple_codes_in_title",
                    "candidates": title_codes,
                })
                return "", ""

            for field_name in ("full_content", "summary", "scope"):
                labeled_code = _extract_labeled_standard_code(doc.get(field_name, ""))
                if labeled_code:
                    return labeled_code, f"{field_name}_label"

            for field_name in ("summary", "scope", "full_content"):
                field_codes = _extract_standard_codes(doc.get(field_name, ""))
                if len(field_codes) == 1:
                    return field_codes[0], field_name
                if len(field_codes) > 1:
                    failed_docs.append({
                        "doc_idx": doc_idx,
                        "title": str(doc.get("title", ""))[:160],
                        "reason": f"multiple_codes_in_{field_name}",
                        "candidates": field_codes[:10],
                    })
                    return "", ""

            failed_docs.append({
                "doc_idx": doc_idx,
                "title": str(doc.get("title", ""))[:160],
                "reason": "code_not_found",
                "candidates": [],
            })
            return "", ""

        for doc_idx, doc in enumerate(tqdm(items, desc="문서 대표 청크 생성 중")):
            code, code_source = resolve_standard_code(doc, doc_idx)
            title = doc.get("title", "")
            if not code:
                continue
            if code_source != "standard_code":
                fallback_code_docs.append({
                    "doc_idx": doc_idx,
                    "code": code,
                    "source": code_source,
                    "title": str(title)[:160],
                })

            doc_type = canonical_doc_type(code)
            collection_type = collection_for_doc_type(doc_type)
            standard_family = standard_family_name(doc_type)

            # 짧은 단어형 제목은 제목만으로 의미가 모호하므로 기준코드를 청크 앞에 명시한다.
            parts = [
                f"[기준코드] {code}",
                f"[기준종류] {standard_family}",
                f"[제목] {code} {title}".strip(),
            ]
            for field_key, field_label, max_len in FIELD_EXCERPT:
                val = str(doc.get(field_key, "") or "").strip()
                if val and val not in SKIP_VALUES:
                    effective_max_len = 1000 if field_key == "scope" and "(전문)" in title else max_len
                    excerpt = val[:effective_max_len]
                    parts.append(f"[{field_label}] {excerpt}")

            page_content = "\n".join(parts)
            if len(page_content) < 50:   # 내용이 거의 없는 문서 건너뜀
                skipped_docs.append({
                    "doc_idx": doc_idx,
                    "code": code,
                    "title": str(title)[:160],
                    "reason": "representative_chunk_too_short",
                })
                continue

            chunk_id = f"{code}_main_{doc_idx}"
            all_chunks.append({
                "id": chunk_id,
                "page_content": page_content,
                "metadata": {
                    "code": code,
                    "full_code": code,
                    "name": title,
                    "title": title,
                    "doc_type": doc_type,
                    "collection_type": collection_type,
                    "standard_family": standard_family,
                    "standard_code_source": code_source,
                    "field": "combined",
                    "chunk_index": 0,
                }
            })
            processed_docs += 1

        logger.info(f"총 {processed_docs}개 문서에서 {len(all_chunks)}개 청크 생성됨.")
        if fallback_code_docs:
            logger.info(f"표준코드 fallback 적용: {len(fallback_code_docs)}개 문서")
        if failed_docs:
            logger.error(f"표준코드 추출 실패: {len(failed_docs)}개 문서")
            return {
                "status": "error",
                "message": "표준코드 추출 실패 문서가 있어 인덱스 빌드를 중단했습니다.",
                "processed_docs": processed_docs,
                "total_input_docs": len(items),
                "fallback_code_docs": len(fallback_code_docs),
                "skipped_docs": len(skipped_docs),
                "failed_docs": len(failed_docs),
                "failed_doc_samples": failed_docs[:20],
            }


        # ── BM25 인덱스 빌드 ─────────────────────────────────────────
        self.bm25_client.build_index(all_chunks)
        self.bm25_client.save_index("kcsc_bm25")

        # ── ChromaDB 인덱스 빌드 ──────────────────────────────────────
        for doc_type in ["KDS", "KCS"]:
            try:
                self.client.delete_collection(name=doc_type)
            except Exception:
                pass

            collection = self.client.create_collection(
                name=doc_type,
                metadata={"description": f"{doc_type} 설계기준 문서"}
            )
            self.collections[doc_type] = collection

            type_chunks = [c for c in all_chunks if c["metadata"].get("collection_type") == doc_type]
            if not type_chunks:
                logger.info(f"{doc_type}: 인덱싱할 청크 없음 (건너뜀)")
                continue

            batch_size = 512  # 임베딩 배치 크기 (속도 최적화)
            total_batches = (len(type_chunks) + batch_size - 1) // batch_size
            logger.info(f"{doc_type} ChromaDB 인덱싱 시작: {len(type_chunks)}개 청크, {total_batches}개 배치")

            for batch_idx in tqdm(range(total_batches), desc=f"{doc_type} 임베딩"):
                s = batch_idx * batch_size
                e = min(s + batch_size, len(type_chunks))
                batch = type_chunks[s:e]

                ids       = [c["id"] for c in batch]
                documents = [c["page_content"] for c in batch]
                metadatas = [c["metadata"] for c in batch]
                embeddings = self._model.encode(documents, show_progress_bar=False).tolist()

                collection.add(ids=ids, documents=documents, metadatas=metadatas, embeddings=embeddings)

            logger.info(f"ChromaDB {doc_type} 인덱싱 완료 ({len(type_chunks)}개 청크)")

        return {
            "status": "success",
            "processed_docs": processed_docs,
            "total_input_docs": len(items),
            "fallback_code_docs": len(fallback_code_docs),
            "skipped_docs": len(skipped_docs),
            "failed_docs": len(failed_docs),
            "total_chunks": len(all_chunks),
            "kds_chunks": sum(1 for c in all_chunks if c["metadata"].get("doc_type") == "KDS"),
            "kcs_chunks": sum(1 for c in all_chunks if c["metadata"].get("doc_type") == "KCS"),
            "excs_chunks": sum(1 for c in all_chunks if c["metadata"].get("doc_type") == "EXCS"),
            "kcs_collection_chunks": sum(1 for c in all_chunks if c["metadata"].get("collection_type") == "KCS"),
        }


    def build_index_from_files(
        self, 
        chunk_size: int = None, 
        overlap: int = None,
        reset: bool = True
    ) -> Dict[str, Dict[str, int]]:
        """
        모든 문서 파일을 인덱싱
        
        Args:
            chunk_size: 청크 크기
            overlap: 청크 간 겹치는 문자 수
            reset: 기존 데이터 리셋 여부
            
        Returns:
            인덱싱 통계 정보
        """
        stats = {}
        
        # 세부 정보 파일 가져오기
        detail_files = get_latest_data_files("KDS_details", self.data_dir) + \
                      get_latest_data_files("KCS_details", self.data_dir)
        
        if not detail_files:
            logger.warning("인덱싱할 데이터 파일이 없습니다.")
            return {"status": "warning", "message": "인덱싱할 데이터 파일이 없습니다."}
        
        for file_path in detail_files:
            try:
                # 파일명에서 문서 유형 추출
                file_name = os.path.basename(file_path)
                doc_type = file_name.split('_')[0]
                
                logger.info(f"{file_path} 파일 인덱싱 시작")
                
                # 데이터 로드
                documents = load_json_data(file_path)
                
                if documents:
                    # 인덱싱
                    doc_count, chunk_count = self.index_documents(
                        doc_type, 
                        documents, 
                        chunk_size=chunk_size, 
                        overlap=overlap,
                        reset=reset
                    )
                    
                    stats[doc_type] = {
                        "documents": doc_count,
                        "chunks": chunk_count,
                        "file": file_path
                    }
                    
                    logger.info(f"{doc_type} 인덱싱 완료: {doc_count}개 문서, {chunk_count}개 청크")
                else:
                    logger.warning(f"{file_path} 파일에 인덱싱할 문서가 없습니다.")
                    stats[doc_type] = {"documents": 0, "chunks": 0, "file": file_path}
                    
            except Exception as e:
                logger.error(f"{file_path} 파일 인덱싱 실패: {str(e)}")
                stats[os.path.basename(file_path)] = {"error": str(e)}
        
        return {
            "status": "success" if stats else "error",
            "stats": stats
        }
    
    def search(
        self, 
        query: str, 
        doc_types: List[str] = None, 
        limit: int = None,
        min_relevance: float = 0.3
    ) -> List[Dict[str, Any]]:
        """
        자연어 쿼리로 관련 문서 검색
        
        Args:
            query: 검색 쿼리
            doc_types: 검색할 문서 유형 목록
            limit: 결과 제한 수
            min_relevance: 최소 관련성 점수
            
        Returns:
            검색 결과 목록
        """
        if not query:
            logger.warning("검색 쿼리가 비어 있습니다.")
            return []
        
        limit = limit or config.SEARCH_LIMIT
        
        allowed_doc_types = {str(dt).upper() for dt in doc_types} if doc_types else None
        collection_doc_types = (
            sorted(collection_names_for_doc_types(doc_types))
            if doc_types
            else list(self.collections.keys())
        )
        
        logger.debug(f"검색 쿼리: '{query}', 대상 문서 유형: {doc_types or collection_doc_types}")
        
        try:
            standard_code_pattern = STANDARD_CODE_PATTERN

            def clean_title_text(value: Any) -> str:
                title = standard_code_pattern.sub(" ", str(value or ""))
                title = re.sub(r"[_\-]+", " ", title)
                title = re.sub(r"\([^)]*(?:\d{2,4}(?:[.\-]\d{1,2})?|전문|수정|제정)[^)]*\)", " ", title)
                title = re.sub(r"(?:수정|제정|전문)", " ", title)
                return re.sub(r"\s+", " ", title).strip()

            def normalize_title_match_text(value: Any) -> str:
                return "".join(ch for ch in str(value or "").lower() if ch.isalnum())

            normalized_query = normalize_title_match_text(query)
            clean_normalized_query = normalize_title_match_text(clean_title_text(query))
            query_tokens = set(self.bm25_client._tokenize(query)) if self.bm25_client.documents else set()
            wants_excs = (
                "고속도로공사전문시방서" in normalized_query
                or "전문시방서" in normalized_query
                or "고속도로공사" in normalized_query
            )
            wants_kds = "설계기준" in normalized_query
            wants_kcs = "표준시방서" in normalized_query
            generic_signal_tokens = {
                "기준", "설계", "설계기준", "공사", "구조", "일반", "수정", "제정",
                "검토", "검토할", "기준으로", "찾아줘", "알려주세요", "콘크", "크리",
                "리트", "콘크리", "크리트", "트구", "리트구", "트구조", "크리트구",
                "리트구조", "계기", "계기준", "길이", "조건", "확인", "결과",
                "어떤", "하는", "하는데", "해야", "하나요", "검토해", "검토해야",
                "토해", "토해야", "적합", "적합한", "충분", "충분한",
            }
            important_short_terms = {
                "내진", "지진", "정착", "이음", "철근", "수압", "시험", "검사",
                "안전", "보건", "관리", "계획", "합성", "강도", "피복", "전단",
                "압축", "볼트", "접합", "용접", "방수", "지반", "기초", "품질",
                "앵커", "가설", "슬래브", "말뚝", "교량", "터널", "배관",
                "네일", "옹벽", "한중", "조임", "체결",
            }
            query_signal_tokens = {
                token for token in query_tokens
                if token not in generic_signal_tokens and (len(token) >= 3 or token in important_short_terms)
            }
            token_idf = getattr(self.bm25_client.bm25, "idf", {}) if self.bm25_client.bm25 else {}

            def token_importance(token: str) -> float:
                try:
                    return max(float(token_idf.get(token, 0.0)), 0.1)
                except Exception:
                    return 0.1
            vector_candidate_count = max(limit, 5)
            lexical_candidate_count = max(limit * 30, 150)

            # 쿼리 임베딩. 모델/Chroma 쪽이 메모리 부족으로 실패해도
            # 아래 BM25 경로는 계속 사용할 수 있어야 한다.
            query_embedding = None
            if self._model is None:
                logger.warning("임베딩 모델이 초기화되지 않아 벡터 검색을 건너뜁니다.")
            else:
                try:
                    query_embedding = self._model.encode(query)
                except Exception as e:
                    logger.warning(f"쿼리 임베딩 실패, BM25 검색만 수행합니다: {e}")
            
            all_results = []
            
            # 각 컬렉션에서 검색
            if query_embedding is not None:
                for doc_type in collection_doc_types:
                    if doc_type not in self.collections:
                        logger.warning(f"{doc_type} 컬렉션이 없습니다.")
                        continue
                    
                    collection = self.collections[doc_type]
                    
                    try:
                        collection_count = collection.count()
                    except Exception as e:
                        logger.warning(f"{doc_type} 컬렉션 count 실패, 벡터 검색 건너뜀: {e}")
                        continue
                    
                    # 컬렉션에 데이터가 없으면 건너뛰기
                    if collection_count == 0:
                        logger.warning(f"{doc_type} 컬렉션에 데이터가 없습니다.")
                        continue
                    
                    try:
                        # 검색 실행
                        results = collection.query(
                            query_embeddings=[query_embedding.tolist()],
                            n_results=min(vector_candidate_count, collection_count)  # 필터링 전에 더 많이 가져옴
                        )
                    except Exception as e:
                        logger.warning(f"{doc_type} Chroma 벡터 검색 실패, BM25 검색으로 계속 진행: {e}")
                        continue
                    
                    # 결과 형식 변환
                    for i in range(len(results['ids'][0])):
                        distance = float(results['distances'][0][i])
                        relevance = float(1.0 / (1.0 + max(distance, 0.0)))  # 거리값을 0~1 유사도로 정규화
                        
                        result = {
                            "id": results['ids'][0][i],
                            "text": results['documents'][0][i],
                            "metadata": results['metadatas'][0][i],
                            "distance": distance,
                            "relevance": relevance
                        }
                        all_results.append(result)
            
            # BM25 검색 (Lexical)
            bm25_results = self.bm25_client.search(query, top_k=lexical_candidate_count) if self.bm25_client.bm25 else []
            title_results = []
            if self.bm25_client.documents and normalized_query:
                for doc in self.bm25_client.documents:
                    metadata = doc.get("metadata", {})
                    title = metadata.get("title") or metadata.get("name")
                    clean_title = clean_title_text(title)
                    normalized_title = normalize_title_match_text(clean_title)
                    if not normalized_title:
                        continue

                    code = metadata.get("code", "")
                    raw_type = canonical_doc_type(code, metadata.get("doc_type", ""))
                    title_score = 0.0

                    if normalized_title in clean_normalized_query:
                        title_score += 80.0 + min(len(normalized_title), 30)
                    else:
                        title_tokens = set(self.bm25_client._tokenize(clean_title))
                        overlap = title_tokens & query_tokens
                        if overlap:
                            title_score += float(len(overlap))
                            title_score += sum(
                                min(token_importance(token), 5.0)
                                for token in overlap & query_signal_tokens
                            )
                            if len(overlap) / max(len(title_tokens), 1) >= 0.5:
                                title_score += 3.0

                    if title_score > 0:
                        if wants_excs and raw_type == "EXCS":
                            title_score += 10.0
                        elif wants_kds and raw_type == "KDS":
                            title_score += 8.0
                        elif wants_kcs and raw_type == "KCS":
                            title_score += 8.0
                        title_results.append((doc, title_score))

                title_results.sort(key=lambda item: item[1], reverse=True)
                title_results = title_results[:lexical_candidate_count]

            substring_results = []
            if self.bm25_client.documents:
                query_terms = [normalize_title_match_text(term) for term in query.split()]
                query_terms = [term for term in query_terms if len(term) >= 2]
                for doc in self.bm25_client.documents:
                    metadata = doc.get("metadata", {})
                    title_text = normalize_title_match_text(clean_title_text(metadata.get("title") or metadata.get("name")))
                    content_text = normalize_title_match_text(doc.get("page_content", ""))
                    score = 0
                    for term in query_terms:
                        if term in title_text:
                            score += 3
                        elif term in content_text:
                            score += 1
                    if score > 0:
                        substring_results.append((doc, score))
                substring_results.sort(key=lambda item: item[1], reverse=True)
                substring_results = substring_results[:lexical_candidate_count]
            
            # RRF (Reciprocal Rank Fusion) 결합
            rrf_k = 60
            fused_scores = {}
            fused_results = {}
            
            # ChromaDB 점수 반영
            for rank, res in enumerate(sorted(all_results, key=lambda x: x['distance'])):
                doc_id = res["id"]
                if doc_id not in fused_scores:
                    fused_scores[doc_id] = 0
                    fused_results[doc_id] = res
                fused_scores[doc_id] += 1.0 / (rrf_k + rank + 1)
                
            # BM25 점수 반영
            for rank, (bm25_doc, score) in enumerate(bm25_results):
                # BM25 인덱스는 같은 청크 id를 보유하므로 문자열 비교 없이 id로 결합한다.
                matched_id = bm25_doc.get("id") or f"bm25_only_{rank}"
                        
                if matched_id not in fused_results:
                    fused_results[matched_id] = {
                        "id": matched_id,
                        "text": bm25_doc["page_content"],
                        "metadata": bm25_doc["metadata"],
                        "distance": 0,
                        "relevance": min(score / 10.0, 1.0) # 임의 정규화
                    }
                    fused_scores[matched_id] = 0
                    
                fused_scores[matched_id] += 1.0 / (rrf_k + rank + 1)

            # 자연어 질의에 포함된 제목구 후보는 BM25 상위권에 없어도 id 기준으로 결합한다.
            for rank, (title_doc, score) in enumerate(title_results):
                matched_id = title_doc.get("id") or f"title_only_{rank}"
                if matched_id not in fused_results:
                    fused_results[matched_id] = {
                        "id": matched_id,
                        "text": title_doc["page_content"],
                        "metadata": title_doc["metadata"],
                        "distance": 0,
                        "relevance": min(score / 50.0, 1.0)
                    }
                    fused_scores[matched_id] = 0
                fused_scores[matched_id] += 4.0 / (rrf_k + rank + 1)

            # BM25가 공백 기반 토큰화로 복합어를 놓칠 때만 부분문자열 후보를 보강한다.
            for rank, (substring_doc, score) in enumerate(substring_results):
                doc_text = substring_doc["page_content"]
                matched_id = substring_doc.get("id")
                if matched_id not in fused_results:
                    fused_results[matched_id] = {
                        "id": matched_id,
                        "text": doc_text,
                        "metadata": substring_doc["metadata"],
                        "distance": 0,
                        "relevance": min(score / 10.0, 1.0)
                    }
                    fused_scores[matched_id] = 0
                fused_scores[matched_id] += 2.0 / (rrf_k + rank + 1)
                 
            # 제목 그대로 검색하는 질의는 같은 제목의 기준을 우선한다.
            if normalized_query:
                for doc_id, res in fused_results.items():
                    metadata = res.get("metadata", {})
                    clean_title = clean_title_text(metadata.get("title") or metadata.get("name"))
                    normalized_title = normalize_title_match_text(clean_title)
                    if not normalized_title:
                        continue
                    raw_type = canonical_doc_type(metadata.get("code", ""), metadata.get("doc_type", ""))
                    if clean_normalized_query == normalized_title:
                        fused_scores[doc_id] += 1.0
                    elif clean_normalized_query in normalized_title or normalized_title in clean_normalized_query:
                        fused_scores[doc_id] += 0.75

                    title_tokens = set(self.bm25_client._tokenize(clean_title))
                    query_tokens = set(self.bm25_client._tokenize(query))
                    meaningful_overlap = {
                        token for token in title_tokens & query_tokens
                        if token in query_signal_tokens
                    }
                    if meaningful_overlap:
                        fused_scores[doc_id] += min(
                            0.08 * sum(token_importance(token) for token in meaningful_overlap),
                            0.65,
                        )

                    content_text = res.get("text", "")
                    content_tokens = set(self.bm25_client._tokenize(content_text[:2500]))
                    content_overlap = {
                        token for token in query_tokens & content_tokens
                        if token in query_signal_tokens
                    }
                    if content_overlap:
                        fused_scores[doc_id] += min(
                            0.05 * sum(token_importance(token) for token in content_overlap),
                            0.75,
                        )

                    rare_title_hits = {
                        token for token in meaningful_overlap
                        if token_importance(token) >= 2.0
                    }
                    rare_content_hits = {
                        token for token in content_overlap
                        if token_importance(token) >= 2.0
                    }
                    if rare_title_hits:
                        fused_scores[doc_id] += min(
                            0.18 * sum(token_importance(token) for token in rare_title_hits),
                            0.85,
                        )
                    if rare_content_hits:
                        fused_scores[doc_id] += min(
                            0.12 * sum(token_importance(token) for token in rare_content_hits),
                            0.75,
                        )

                    code = metadata.get("code", "")
                    fused_scores[doc_id] += score_ranking_rules(normalized_query, code, "vector")
                    if wants_excs:
                        if raw_type == "EXCS":
                            fused_scores[doc_id] += 0.75
                        elif raw_type in {"KCS", "KDS"}:
                            fused_scores[doc_id] -= 0.35

            # 최종 RRF 스코어로 정렬
            final_results = sorted(list(fused_results.values()), key=lambda x: fused_scores[x["id"]], reverse=True)
            deduped_results = []
            seen_codes = set()
            for res in final_results:
                res = dict(res)
                raw_fusion_score = float(fused_scores.get(res["id"], 0.0))
                res["fusion_score"] = raw_fusion_score
                res["relevance"] = min(raw_fusion_score / 0.08, 1.0)
                if res["relevance"] < min_relevance:
                    continue
                metadata = dict(res.get("metadata", {}) or {})
                if not doc_type_matches(metadata, allowed_doc_types):
                    continue
                metadata["doc_type"] = canonical_doc_type(metadata.get("code", ""), metadata.get("doc_type", ""))
                metadata.setdefault("collection_type", collection_for_doc_type(metadata["doc_type"]))
                res["metadata"] = metadata
                code = metadata.get("code") or res.get("id")
                if code in seen_codes:
                    continue
                seen_codes.add(code)
                deduped_results.append(res)
            
            # 결과 수 제한
            limited_results = deduped_results[:limit]
            
            logger.debug(f"하이브리드 RRF 검색 결과: {len(limited_results)}개 항목")
            return limited_results
            
        except Exception as e:
            logger.error(f"검색 중 오류 발생: {str(e)}")
            raise VectorDBError(f"검색 중 오류 발생: {str(e)}", "search")
    
    def get_document_by_id(self, doc_type: str, doc_id: str) -> Optional[Dict[str, Any]]:
        """
        ID로 문서 가져오기
        
        Args:
            doc_type: 문서 유형
            doc_id: 문서 ID
            
        Returns:
            문서 정보 또는 None (찾지 못한 경우)
        """
        if doc_type not in self.collections:
            logger.warning(f"{doc_type} 컬렉션이 없습니다.")
            return None
        
        collection = self.collections[doc_type]
        
        try:
            results = collection.get(ids=[doc_id])
            
            if not results or not results["ids"]:
                return None
            
            return {
                "id": results["ids"][0],
                "text": results["documents"][0],
                "metadata": results["metadatas"][0]
            }
            
        except Exception as e:
            logger.error(f"ID {doc_id}로 문서 조회 중 오류 발생: {str(e)}")
            return None
    
    def get_collection_stats(self) -> Dict[str, int]:
        """
        컬렉션 통계 정보 조회
        
        Returns:
            컬렉션별 문서 수
        """
        stats = {}
        
        for doc_type, collection in self.collections.items():
            try:
                stats[doc_type] = collection.count()
            except Exception as e:
                logger.error(f"{doc_type} 컬렉션 통계 조회 실패: {str(e)}")
                stats[doc_type] = -1
        
        return stats
    
    def similarity_search(
        self, 
        text: str, 
        doc_types: List[str] = None,
        limit: int = None
    ) -> List[Dict[str, Any]]:
        """
        텍스트 유사성 검색 (자연어 검색 대신 직접 텍스트 비교)
        
        Args:
            text: 비교할 텍스트
            doc_types: 검색할 문서 유형 목록
            limit: 결과 제한 수
            
        Returns:
            유사한 문서 목록
        """
        # 일반 검색 메서드 활용
        return self.search(text, doc_types, limit)
    
    def reset_collection(self, doc_type: str = None) -> Dict[str, Any]:
        """
        컬렉션 데이터 초기화
        
        Args:
            doc_type: 초기화할 문서 유형 (None이면 모든 컬렉션)
            
        Returns:
            초기화 결과 통계
        """
        results = {}
        
        if doc_type:
            if doc_type not in self.collections:
                raise VectorDBError(f"{doc_type} 컬렉션이 없습니다.", "collection_missing")
                
            # 단일 컬렉션 초기화
            collection = self.collections[doc_type]
            count_before = collection.count()
            collection.delete(where={})
            results[doc_type] = {"before": count_before, "after": collection.count()}
            
        else:
            # 모든 컬렉션 초기화
            for type_name, collection in self.collections.items():
                count_before = collection.count()
                collection.delete(where={})
                results[type_name] = {"before": count_before, "after": collection.count()}
        
        return {
            "status": "success",
            "reset_stats": results
        }

# 실행 예시
if __name__ == "__main__":
    # 벡터 DB 초기화
    vector_db = KCSCVectorDB()
    
    # 기존 데이터 통계 확인
    print("벡터 DB 초기 상태:", vector_db.get_collection_stats())
    
    # 테스트 검색 (데이터가 있는 경우)
    for collection_name, count in vector_db.get_collection_stats().items():
        if count > 0:
            print(f"\n{collection_name} 컬렉션에서 테스트 검색:")
            results = vector_db.search("철근콘크리트 기둥 설계", [collection_name], 3)
            
            for i, result in enumerate(results):
                print(f"\n--- 결과 {i+1} ---")
                print(f"문서: {result['metadata']['name']} ({result['metadata']['code']})")
                print(f"유사도: {result['relevance']:.4f}")
                print(f"내용 미리보기: {result['text'][:150]}...")