# 안전한 파일 로딩으로 input stream 에러 해결
import streamlit as st
import pandas as pd
import numpy as np
import pickle
import json
import os

st.set_page_config(
    page_title="슬로레터 지능형 검색",
    page_icon="🔍",
    layout="wide"
)

class SafeSlowletterSearch:
    def __init__(self):
        self.bm25 = None
        self.processed_data = None
        self.metadata = {}
        self.system_loaded = False
    
    def safe_load_json(self, filepath):
        """JSON 파일 안전 로딩"""
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            st.error(f"JSON 로딩 실패 ({filepath}): {e}")
            return {}
    
    def safe_load_pickle(self, filepath):
        """Pickle 파일 안전 로딩"""
        try:
            with open(filepath, 'rb') as f:
                return pickle.load(f)
        except Exception as e:
            st.error(f"Pickle 로딩 실패 ({filepath}): {e}")
            return None
    
    def safe_load_numpy(self, filepath):
        """NumPy 파일 안전 로딩"""
        try:
            return np.load(filepath, allow_pickle=True)
        except Exception as e:
            st.error(f"NumPy 로딩 실패 ({filepath}): {e}")
            return None
    
    def load_system_step_by_step(self):
        """단계별 시스템 로딩"""
        st.write("🔄 시스템 파일 로딩 시작...")
        
        # 1. 파일 존재 확인
        required_files = {
            'system_metadata.json': '메타데이터',
            'processed_data.pkl': '처리된 데이터', 
            'bm25_index.pkl': 'BM25 인덱스',
            'embeddings.npy': '임베딩 데이터'
        }
        
        missing_files = []
        for file, desc in required_files.items():
            if os.path.exists(file):
                size = os.path.getsize(file) / 1024 / 1024
                st.success(f"✅ {desc}: {file} ({size:.1f}MB)")
            else:
                st.error(f"❌ {desc}: {file} (없음)")
                missing_files.append(file)
        
        if missing_files:
            st.error(f"필수 파일 부족: {missing_files}")
            return False
        
        # 2. JSON 파일 먼저 로딩 (가장 안전)
        st.write("📋 메타데이터 로딩 중...")
        self.metadata = self.safe_load_json('system_metadata.json')
        if self.metadata:
            st.success("✅ 메타데이터 로딩 완료")
            st.json(self.metadata)
        
        # 3. 처리된 데이터 로딩
        st.write("📊 데이터 로딩 중...")
        self.processed_data = self.safe_load_pickle('processed_data.pkl')
        if self.processed_data is not None:
            st.success(f"✅ 데이터 로딩 완료: {len(self.processed_data)}개 행")
        else:
            st.error("❌ 데이터 로딩 실패")
            return False
        
        # 4. BM25 인덱스 로딩
        st.write("🔍 BM25 인덱스 로딩 중...")
        self.bm25 = self.safe_load_pickle('bm25_index.pkl')
        if self.bm25 is not None:
            st.success("✅ BM25 인덱스 로딩 완료")
        else:
            st.error("❌ BM25 인덱스 로딩 실패")
            return False
        
        # 5. 임베딩은 일단 스킵 (가장 문제가 될 가능성 높음)
        st.info("ℹ️ 임베딩 파일은 일단 스킵 (BM25만 사용)")
        
        self.system_loaded = True
        return True
    
    def tokenize_korean(self, text):
        import re
        tokens = re.findall(r'\b\w+\b', str(text).lower())
        return [token for token in tokens if len(token) > 1]
    
    def search(self, query, top_k=10):
        """BM25 전용 검색"""
        if not self.system_loaded or self.bm25 is None:
            return []
        
        try:
            tokenized_query = self.tokenize_korean(query)
            bm25_scores = self.bm25.get_scores(tokenized_query)
            
            # 상위 결과 선택
            top_indices = np.argsort(bm25_scores)[::-1][:top_k]
            
            results = []
            for rank, idx in enumerate(top_indices, 1):
                if bm25_scores[idx] > 0:
                    row = self.processed_data.iloc[idx]
                    results.append({
                        'rank': rank,
                        'title': row.get('h3_title', ''),
                        'date': row.get('date', ''),
                        'newsletter': row.get('h1_title', ''),
                        'entities': row.get('entities', ''),
                        'events': row.get('events', ''),
                        'content': row.get('h3_content_text', ''),
                        'score': float(bm25_scores[idx])
                    })
            
            return results
            
        except Exception as e:
            st.error(f"검색 중 오류: {e}")
            return []

def main():
    st.title("🔍 슬로레터 지능형 검색")
    st.markdown("**Input Stream 에러 해결 버전**")
    
    # 시스템 초기화
    if 'search_system' not in st.session_state:
        st.session_state.search_system = SafeSlowletterSearch()
        st.session_state.system_loaded = False
    
    search_system = st.session_state.search_system
    
    # 시스템 로딩 (한 번만)
    if not st.session_state.system_loaded:
        with st.container():
            if st.button("🚀 시스템 로딩 시작", type="primary"):
                success = search_system.load_system_step_by_step()
                st.session_state.system_loaded = success
                if success:
                    st.balloons()
    
    # 검색 인터페이스 (시스템 로딩 후)
    if st.session_state.system_loaded:
        st.divider()
        st.subheader("🔍 검색")
        
        # 빠른 검색 버튼
        quick_queries = ["윤석열", "이재명", "명태균", "탄핵", "경제", "북한"]
        
        cols = st.columns(3)
        selected_query = None
        
        for i, query in enumerate(quick_queries):
            col = cols[i % 3]
            if col.button(query, key=f"quick_{i}"):
                selected_query = query
        
        # 검색창
        search_query = st.text_input(
            "검색어:",
            value=selected_query or "",
            placeholder="예: 윤석열 대통령"
        )
        
        # 검색 실행
        if search_query:
            with st.spinner("검색 중..."):
                results = search_system.search(search_query, top_k=10)
            
            if results:
                st.success(f"🎉 {len(results)}개 결과 발견!")
                
                for result in results:
                    with st.expander(f"🏆 {result['rank']}. {result['title']} ({result['date']})"):
                        col1, col2 = st.columns([3, 1])
                        
                        with col1:
                            if result['entities']:
                                entities = [e.strip() for e in result['entities'].split(';')[:3] if e.strip()]
                                if entities:
                                    st.write("**👥 관련 인물:**", ", ".join(entities))
                            
                            if result['content']:
                                st.write("**📝 내용:**")
                                st.write(result['content'][:200] + "..." if len(result['content']) > 200 else result['content'])
                        
                        with col2:
                            st.metric("점수", f"{result['score']:.2f}")
            else:
                st.warning("검색 결과가 없습니다.")
    
    else:
        st.info("👆 '시스템 로딩 시작' 버튼을 클릭해서 시작하세요!")

if __name__ == "__main__":
    main()