src/streamlit_app.py

import streamlit as st
import os
import json
import re
import numpy as np
from typing import List, Dict, Tuple, Optional
from pathlib import Path
import logging
from sentence_transformers import SentenceTransformer
import faiss
import json
from rank_bm25 import BM25Okapi

# 기본 로깅 설정
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 페이지 설정
st.set_page_config(
    page_title="하이브리드 차량 정비 검색 시스템",
    page_icon="🔧",
    layout="wide",
    initial_sidebar_state="expanded"
)

# CSS 스타일
st.markdown("""
<style>
.main-header {
    font-size: 2.5rem;
    color: #1f4e79;
    text-align: center;
    margin-bottom: 2rem;
    font-weight: bold;
}

.search-container {
    background-color: #f8f9fa;
    padding: 2rem;
    border-radius: 10px;
    margin-bottom: 2rem;
    border-left: 5px solid #1f4e79;
}

.result-card {
    background-color: white;
    padding: 1.5rem;
    border-radius: 8px;
    margin-bottom: 1rem;
    border: 1px solid #dee2e6;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}

.score-badge {
    background-color: #e3f2fd;
    color: #1565c0;
    padding: 0.25rem 0.75rem;
    border-radius: 15px;
    font-size: 0.8rem;
    font-weight: bold;
}

.category-badge {
    background-color: #f3e5f5;
    color: #7b1fa2;
    padding: 0.25rem 0.75rem;
    border-radius: 15px;
    font-size: 0.8rem;
    margin-right: 0.5rem;
}

.content-text {
    background-color: #f8f9fa;
    padding: 1rem;
    border-radius: 5px;
    border-left: 3px solid #28a745;
    margin-top: 1rem;
    line-height: 1.6;
}

.metric-card {
    background-color: #e8f5e8;
    padding: 1rem;
    border-radius: 5px;
    text-align: center;
    margin: 0.5rem;
}
</style>
""", unsafe_allow_html=True)

# 간단한 부품 사전 (실제 vocab.py 대신 사용)
PARTS = [
    "수동변속기", "클러치", "브레이크", "엔진", "타이어", "배터리",
    "오일", "필터", "벨트", "호스", "펌프", "센서", "트랜스미션",
    "디스크", "패드", "슈", "로터", "캘리퍼", "마스터실린더"
]

# 간단한 시스템 매핑 (실제 parts_config.py 대신 사용)
SYSTEM_PARTS_MAP = {
    "수동변속기": ["클러치", "변속기", "드라이브샤프트", "디퍼렌셜"],
    "엔진": ["피스톤", "실린더", "크랭크샤프트", "캠샤프트"],
    "브레이크": ["브레이크패드", "브레이크디스크", "캘리퍼", "마스터실린더"]
}

def get_specific_parts_for_system(system_name: str) -> list:
    return SYSTEM_PARTS_MAP.get(system_name, [])

def get_all_specific_parts() -> list:
    all_parts = []
    for parts in SYSTEM_PARTS_MAP.values():
        all_parts.extend(parts)
    return list(set(all_parts))

class SimpleMecab:
    """MeCab 대신 사용할 간단한 형태소 분석기"""
    def pos(self, text):
        # 간단한 명사/동사 추출 (실제 환경에서는 MeCab 사용)
        words = text.split()
        return [(word, 'NN') for word in words if len(word) > 1]

class HybridMultiCollectionSearcher:
    def __init__(self, model_name: str = "upskyy/bge-m3-korean", target_system: str = None):
        """
        하이브리드 다중 컬렉션 검색기 (벡터 + 키워드 검색)
        """
        self.model = None  # 나중에 로드
        self.collections = {}
        self.bm25_indexes = {}
        self.target_system = target_system
        self.mecab = SimpleMecab()  # 간단한 분석기 사용
        self.model_name = model_name
        
    @st.cache_resource
    def load_model(_self):
        """모델을 캐시와 함께 로드"""
        try:
            return SentenceTransformer(_self.model_name)
        except Exception as e:
            st.error(f"모델 로드 실패: {e}")
            return None
    
    def _extract_nouns_and_verbs(self, text: str) -> str:
        """간단한 명사와 동사 추출"""
        try:
            # 부품명 우선 처리
            for part in PARTS:
                if part in text:
                    text = text.replace(part, f" {part} ")
            
            # 간단한 명사 추출 (실제로는 MeCab 사용)
            morphs = self.mecab.pos(text)
            meaningful_words = []
            
            for word, pos in morphs:
                if len(word) > 1 and not word.isspace():
                    meaningful_words.append(word)
            
            return ' '.join(meaningful_words)
        except Exception as e:
            return text
    
    def _normalize_text_for_matching(self, text: str) -> str:
        normalized = text.lower()
        normalized = re.sub(r'[.]', '', normalized)
        return normalized
    
    def _normalize_scores(self, scores: np.ndarray) -> np.ndarray:
        """점수를 0-1 범위로 정규화"""
        scores = np.array(scores)
        if len(scores) == 0 or scores.max() == scores.min():
            return np.ones_like(scores) * 0.5
        return (scores - scores.min()) / (scores.max() - scores.min())
    
    def _calculate_boost_score(self, original_query: str, processed_query: str, metadata: Dict, content: str) -> float:
        """간단한 부스팅 점수 계산"""
        boost_score = 0
        query_lower = original_query.lower()
        
        # 콘텐츠 타입 매칭
        content_type = metadata.get('content_type', '')
        if '탈거' in query_lower and '탈거' in content_type:
            boost_score += 0.5
        if '장착' in query_lower and '장착' in content_type:
            boost_score += 0.5
        if '점검' in query_lower and '점검' in content_type:
            boost_score += 0.5
        
        # 시스템 매칭
        system = metadata.get('vehicle_info', {}).get('system', '')
        if system and any(word in system.lower() for word in query_lower.split()):
            boost_score += 0.3
        
        return boost_score
    
    def create_sample_collection(self, collection_name: str):
        """샘플 데이터로 컬렉션 생성"""
        try:
            if self.model is None:
                self.model = self.load_model()
                if self.model is None:
                    return False
            
            # 샘플 데이터
            sample_data = [
                {
                    'chunk_id': 'sample_001',
                    'content': '수동변속기 탈거 시에는 먼저 엔진을 정지하고 변속기 오일을 배출합니다. 클러치를 분리한 후 변속기를 탈거합니다.',
                    'metadata': {
                        'chunk_id': 'sample_001',
                        'content_type': '탈거방법',
                        'main_topic': '수동변속기 탈거',
                        'vehicle_info': {'system': '수동변속기', 'model': '에어로시티'},
                        'category_levels': ['변속기', '수동변속기', '탈거방법'],
                        'extracted_components': ['변속기', '클러치']
                    }
                },
                {
                    'chunk_id': 'sample_002', 
                    'content': '수동변속기 장착은 탈거의 역순으로 진행합니다. 변속기를 정확한 위치에 고정하고 클러치를 연결합니다.',
                    'metadata': {
                        'chunk_id': 'sample_002',
                        'content_type': '장착방법',
                        'main_topic': '수동변속기 장착',
                        'vehicle_info': {'system': '수동변속기', 'model': '에어로시티'},
                        'category_levels': ['변속기', '수동변속기', '장착방법'],
                        'extracted_components': ['변속기', '클러치']
                    }
                },
                {
                    'chunk_id': 'sample_003',
                    'content': '변속기 오일 점검 시 오일 레벨과 오일 상태를 확인합니다. 규정량은 2.5L이며 오일 온도는 80°C에서 측정합니다.',
                    'metadata': {
                        'chunk_id': 'sample_003',
                        'content_type': '점검절차',
                        'main_topic': '오일 점검',
                        'vehicle_info': {'system': '수동변속기', 'model': '에어로시티'},
                        'category_levels': ['변속기', '수동변속기', '점검절차'],
                        'extracted_components': ['오일']
                    }
                }
            ]
            
            # 검색 텍스트 생성
            search_texts = []
            metadata_list = []
            content_dict = {}
            
            for data in sample_data:
                metadata = data['metadata']
                content = data['content']
                
                # 검색용 텍스트 구성
                search_components = [
                    metadata.get('content_type', ''),
                    metadata.get('main_topic', ''),
                    ' '.join(metadata.get('category_levels', [])),
                    content
                ]
                
                search_text = self._extract_nouns_and_verbs(' '.join(search_components))
                search_texts.append(search_text)
                metadata_list.append(metadata)
                content_dict[metadata['chunk_id']] = content
            
            # 벡터 임베딩 생성
            embeddings = self.model.encode(search_texts, show_progress_bar=False)
            
            # FAISS 인덱스 생성
            embedding_dim = embeddings.shape[1]
            faiss.normalize_L2(embeddings)
            faiss_index = faiss.IndexFlatIP(embedding_dim)
            faiss_index.add(embeddings.astype(np.float32))
            
            # BM25 인덱스 생성
            tokenized_docs = [text.split() for text in search_texts]
            bm25_index = BM25Okapi(tokenized_docs)
            
            # 컬렉션 저장
            self.collections[collection_name] = {
                'metadata_list': metadata_list,
                'content_dict': content_dict,
                'search_texts': search_texts,
                'faiss_index': faiss_index
            }
            
            self.bm25_indexes[collection_name] = bm25_index
            return True
            
        except Exception as e:
            logger.error(f"샘플 컬렉션 생성 실패: {e}")
            return False
        """저장된 하이브리드 컬렉션들 로드 (FAISS + BM25) - pickle 없이"""
        save_dir = Path(save_dir)
        
        if not save_dir.exists():
            logger.warning(f"컬렉션 디렉토리가 존재하지 않습니다: {save_dir}")
            return False
        
        loaded_collections = []
        
        for collection_dir in save_dir.iterdir():
            if collection_dir.is_dir():
                collection_name = collection_dir.name
                
                try:
                    # 1. FAISS 인덱스 로드
                    faiss_path = collection_dir / "faiss.index"
                    if not faiss_path.exists():
                        logger.warning(f"FAISS 인덱스가 없습니다: {faiss_path}")
                        continue
                        
                    faiss_index = faiss.read_index(str(faiss_path))
                    
                    # 2. BM25 토큰 데이터 로드 (JSON)
                    bm25_tokens_path = collection_dir / "bm25_tokens.json"
                    if not bm25_tokens_path.exists():
                        logger.warning(f"BM25 토큰 데이터가 없습니다: {bm25_tokens_path}")
                        continue
                        
                    with open(bm25_tokens_path, 'r', encoding='utf-8') as f:
                        tokenized_docs = json.load(f)
                    
                    # BM25 인덱스 재생성
                    bm25_index = BM25Okapi(tokenized_docs)
                    
                    # 3. 메타데이터 로드 (JSON)
                    metadata_path = collection_dir / "metadata.json"
                    if not metadata_path.exists():
                        logger.warning(f"메타데이터가 없습니다: {metadata_path}")
                        continue
                        
                    with open(metadata_path, 'r', encoding='utf-8') as f:
                        save_data = json.load(f)
                    
                    # 컬렉션 복원
                    self.collections[collection_name] = {
                        'faiss_index': faiss_index,
                        **save_data
                    }
                    
                    self.bm25_indexes[collection_name] = bm25_index
                    loaded_collections.append(collection_name)
                    
                    logger.info(f"컬렉션 '{collection_name}' 로드 완료")
                    
                except Exception as e:
                    logger.error(f"컬렉션 '{collection_name}' 로드 실패: {e}")
                    continue
        
        if loaded_collections:
            logger.info(f"하이브리드 컬렉션 로드 완료: {loaded_collections}")
            return True
        else:
            logger.error("로드된 컬렉션이 없습니다.")
            return False
    
    def list_collections(self) -> List[str]:
        """등록된 컬렉션 목록 반환"""
        return list(self.collections.keys())
    
    def search_collection(self, collection_name: str, query: str, top_k: int = 5, alpha: float = 0.7) -> List[Dict]:
        """하이브리드 검색 수행"""
        if collection_name not in self.collections:
            return []
        
        if self.model is None:
            self.model = self.load_model()
            if self.model is None:
                return []
        
        collection = self.collections[collection_name]
        faiss_index = collection['faiss_index']
        metadata_list = collection['metadata_list']
        content_dict = collection['content_dict']
        bm25_index = self.bm25_indexes[collection_name]
        
        # 쿼리 처리
        processed_query = self._extract_nouns_and_verbs(query)
        
        # 벡터 검색
        query_embedding = self.model.encode([processed_query])
        faiss.normalize_L2(query_embedding)
        
        search_k = min(len(metadata_list), top_k * 3)
        dense_similarities, dense_indices = faiss_index.search(
            query_embedding.astype(np.float32), search_k
        )
        
        # 키워드 검색
        query_tokens = processed_query.split()
        sparse_scores = bm25_index.get_scores(query_tokens)
        
        # 점수 정규화
        dense_scores_norm = self._normalize_scores(dense_similarities[0])
        sparse_scores_norm = self._normalize_scores(sparse_scores)
        
        # 결과 생성
        results = []
        for i, (similarity, idx) in enumerate(zip(dense_similarities[0], dense_indices[0])):
            if idx == -1:
                continue
            
            metadata = metadata_list[idx]
            chunk_id = metadata['chunk_id']
            content = content_dict.get(chunk_id, '')
            
            dense_score = dense_scores_norm[i]
            sparse_score = sparse_scores_norm[idx] if idx < len(sparse_scores_norm) else 0
            boost_score = self._calculate_boost_score(query, processed_query, metadata, content)
            
            hybrid_score = (alpha * dense_score + (1 - alpha) * sparse_score + boost_score)
            
            category_levels = metadata.get('category_levels', [])
            category_path = ' > '.join(category_levels)
            
            result = {
                'chunk_id': chunk_id,
                'content': content,
                'metadata': metadata,
                'dense_similarity': float(similarity),
                'dense_score': dense_score,
                'sparse_score': sparse_score,
                'boost_score': boost_score,
                'hybrid_score': hybrid_score,
                'vehicle_info': metadata.get('vehicle_info', {}),
                'content_type': metadata.get('content_type', ''),
                'main_topic': metadata.get('main_topic', ''),
                'category_path': category_path,
                'processed_query': processed_query,
            }
            results.append(result)
        
        results.sort(key=lambda x: x['hybrid_score'], reverse=True)
        return results[:top_k]

# Streamlit 앱 시작
def main():
    # 제목
    st.markdown('<h1 class="main-header">🔧 하이브리드 차량 정비 검색 시스템</h1>', unsafe_allow_html=True)
    
    # 사이드바
    with st.sidebar:
        st.header("⚙️ 설정")
        
        # 검색 파라미터
        st.subheader("검색 설정")
        top_k = st.slider("결과 개수", min_value=1, max_value=10, value=5)
        alpha = st.slider("벡터 검색 가중치", min_value=0.0, max_value=1.0, value=0.7, step=0.1)
        
        st.info(f"벡터 검색: {alpha:.1f}, 키워드 검색: {1-alpha:.1f}")
        
        # 시스템 선택
        st.subheader("대상 시스템")
        target_system = st.selectbox(
            "시스템 선택",
            ["수동변속기", "엔진", "브레이크"],
            index=0
        )
    
    # 메인 영역
    # 검색기 초기화
    if 'searcher' not in st.session_state:
        with st.spinner('검색 시스템 초기화 중...'):
            try:
                st.session_state.searcher = HybridMultiCollectionSearcher(target_system=target_system)
                
                # 먼저 샘플 데이터로 테스트
                st.info("🧪 샘플 데이터로 테스트 중...")
                success = st.session_state.searcher.create_sample_collection("테스트")
                
                if success:
                    st.success("✅ 샘플 검색 시스템이 준비되었습니다!")
                    st.info("💡 실제 컬렉션을 사용하려면 `saved_collections` 폴더를 업로드하세요.")
                else:
                    st.error("❌ 시스템 초기화에 실패했습니다.")
                    
            except Exception as e:
                st.error(f"❌ 초기화 오류: {str(e)}")
                st.info("🔧 문제를 해결하는 중입니다...")
    
    # 검색기가 있는 경우에만 진행
    if 'searcher' in st.session_state:
        available_collections = st.session_state.searcher.list_collections()
        
        # 컬렉션이 있는 경우에만 검색 인터페이스 표시
        if available_collections:
            # 컬렉션 선택
            st.subheader("📚 검색 대상 컬렉션")
            selected_collection = st.selectbox(
                "컬렉션 선택",
                available_collections,
                help="검색할 컬렉션을 선택하세요"
            )
            
            # 검색 인터페이스
            with st.container():
                st.markdown('<div class="search-container">', unsafe_allow_html=True)
                
                # 검색어 입력
                query = st.text_input(
                    "🔍 질문을 입력하세요",
                    placeholder="예: 수동변속기 탈거는 어떻게 하나요?",
                    help="차량 정비에 관한 질문을 자유롭게 입력하세요."
                )
                
                # 검색 버튼
                col1, col2, col3 = st.columns([1, 2, 1])
                with col2:
                    search_button = st.button("🔍 검색하기", type="primary", use_container_width=True)
                
                st.markdown('</div>', unsafe_allow_html=True)
            
            # 검색 실행
            if search_button and query:
                with st.spinner('검색 중...'):
                    results = st.session_state.searcher.search_collection(
                        selected_collection, 
                        query, 
                        top_k=top_k,
                        alpha=alpha
                    )
                
                if results:
                    st.success(f"✅ {len(results)}개의 검색 결과를 찾았습니다.")
                    
                    # 검색 통계
                    col1, col2, col3, col4 = st.columns(4)
                    with col1:
                        st.markdown('<div class="metric-card"><strong>검색 결과</strong><br>' + f'{len(results)}개</div>', unsafe_allow_html=True)
                    with col2:
                        avg_score = np.mean([r['hybrid_score'] for r in results])
                        st.markdown('<div class="metric-card"><strong>평균 점수</strong><br>' + f'{avg_score:.3f}</div>', unsafe_allow_html=True)
                    with col3:
                        max_score = max([r['hybrid_score'] for r in results])
                        st.markdown('<div class="metric-card"><strong>최고 점수</strong><br>' + f'{max_score:.3f}</div>', unsafe_allow_html=True)
                    with col4:
                        st.markdown('<div class="metric-card"><strong>컬렉션</strong><br>' + f'{selected_collection}</div>', unsafe_allow_html=True)
                    
                    st.markdown("---")
                    
                    # 검색 결과 표시
                    for i, result in enumerate(results, 1):
                        st.markdown('<div class="result-card">', unsafe_allow_html=True)
                        
                        # 헤더
                        col1, col2 = st.columns([3, 1])
                        with col1:
                            st.markdown(f"### 📄 결과 {i}: {result['main_topic']}")
                        with col2:
                            st.markdown(f'<span class="score-badge">점수: {result["hybrid_score"]:.3f}</span>', unsafe_allow_html=True)
                        
                        # 메타데이터
                        col1, col2 = st.columns(2)
                        with col1:
                            st.markdown(f'<span class="category-badge">{result["content_type"]}</span>', unsafe_allow_html=True)
                            st.markdown(f"**경로:** {result['category_path']}")
                        
                        with col2:
                            if result['vehicle_info']:
                                vehicle = result['vehicle_info']
                                st.markdown(f"**차량:** {vehicle.get('model', 'N/A')}")
                                st.markdown(f"**시스템:** {vehicle.get('system', 'N/A')}")
                        
                        # 내용
                        st.markdown('<div class="content-text">', unsafe_allow_html=True)
                        st.markdown(f"**📋 내용:**\n\n{result['content']}")
                        st.markdown('</div>', unsafe_allow_html=True)
                        
                        # 상세 점수 (확장 가능)
                        with st.expander("🔍 상세 점수 보기"):
                            score_col1, score_col2, score_col3 = st.columns(3)
                            with score_col1:
                                st.metric("벡터 점수", f"{result['dense_score']:.3f}")
                            with score_col2:
                                st.metric("키워드 점수", f"{result['sparse_score']:.3f}")
                            with score_col3:
                                st.metric("부스팅 점수", f"{result['boost_score']:.3f}")
                            
                            st.markdown(f"**처리된 쿼리:** `{result['processed_query']}`")
                            st.markdown(f"**청크 ID:** `{result['chunk_id']}`")
                        
                        st.markdown('</div>', unsafe_allow_html=True)
                        st.markdown("---")
                
                else:
                    st.warning("🤔 검색 결과가 없습니다. 다른 키워드로 검색해보세요.")
            
            elif search_button and not query:
                st.warning("⚠️ 검색어를 입력해주세요.")
        
        else:
            # 컬렉션이 없는 경우
            st.warning("⚠️ 로드된 컬렉션이 없습니다.")
            st.markdown("""
            ### 📁 컬렉션 파일 업로드 방법
            
            1. **로컬에서 컬렉션 생성**:
            ```python
            # 원본 코드 사용
            searcher = HybridMultiCollectionSearcher()
            searcher.add_collection("수동변속기", metadata_dir, chunks_dir)
            searcher.save_collections("./saved_collections")
            ```
            
            2. **생성된 파일들을 허깅페이스 Space에 업로드**:
            - `saved_collections/` 폴더 전체를 업로드
            - 각 컬렉션별로 `.pkl`, `.index` 파일들이 포함됨
            
            3. **앱 재시작** 후 검색 가능
            """)
    
    # 사용 가이드 (컬렉션이 있을 때만 표시)
    if 'searcher' in st.session_state and st.session_state.searcher.list_collections() and not query:
        st.markdown("### 💡 사용 가이드")
        
        col1, col2 = st.columns(2)
        with col1:
            st.markdown("""
            **🔧 정비 작업 질문:**
            - "수동변속기 탈거는 어떻게 하나요?"
            - "클러치 점검 방법을 알려주세요"
            - "변속기 오일 교환 절차는?"
            """)
        
        with col2:
            st.markdown("""
            **⚙️ 부품 정보 질문:**
            - "브레이크 패드 사양은?"
            - "엔진 오일 용량은 얼마인가요?"
            - "타이어 공기압 기준치는?"
            """)

if __name__ == "__main__":
    main()