# simple_test.py - 간단한 RAG 테스트 스크립트
import sys
import os

# 인코딩 설정
if sys.platform.startswith('win'):
    os.system('chcp 65001 > nul')

def test_rag_system():
    """RAG 시스템 간단 테스트"""
    print("RAG 시스템 테스트 시작")
    print("=" * 50)
    
    try:
        from rag_system import HFSpacesTaxRAG
        
        # RAG 시스템 초기화
        print("RAG 시스템 초기화 중...")
        rag = HFSpacesTaxRAG()
        
        if not rag.vector_db or not rag.documents:
            print("벡터 DB가 없습니다. 시스템을 재구축합니다...")
            success = rag.build_system()
            if not success:
                print("RAG 시스템 구축 실패")
                return False
        
        # 간단한 테스트 쿼리들
        test_queries = [
            "취득세율이 얼마인가요?",
            "1세대 1주택자 감면 혜택은?",
            "취득세 신고 기한은 언제인가요?"
        ]
        
        print("\n테스트 쿼리 실행:")
        print("-" * 30)
        
        for i, query in enumerate(test_queries, 1):
            print(f"\n{i}. {query}")
            
            try:
                result = rag.query(query)
                
                if result and result.get('answer'):
                    print(f"응답: {result['answer'][:100]}...")
                    print(f"소스 문서: {result.get('source_count', 0)}개")
                else:
                    print("응답 생성 실패")
                    
            except Exception as e:
                print(f"쿼리 처리 오류: {e}")
        
        print("\n" + "=" * 50)
        print("테스트 완료")
        return True
        
    except Exception as e:
        print(f"RAG 시스템 오류: {e}")
        return False

if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == '--rebuild':
        print("RAG 시스템을 재구축합니다...")
        # 기존 벡터 DB 파일 삭제
        for file in ['vector_db.faiss', 'documents.pkl', 'metadata.json']:
            if os.path.exists(file):
                os.remove(file)
                print(f"기존 파일 삭제: {file}")
    
    success = test_rag_system()
    sys.exit(0 if success else 1)