# run_rag_tests.py - RAG 테스트 실행 스크립트 (허깅페이스 스페이스 최적화)
"""
RAG 시스템 테스트 및 튜닝을 위한 통합 실행 스크립트

사용법:
  python run_rag_tests.py                 # 기본 테스트
  python run_rag_tests.py --rebuild       # RAG 시스템 재구축 후 테스트
  python run_rag_tests.py --tune          # 설정 튜닝
  python run_rag_tests.py --benchmark     # 성능 벤치마크만
  python run_rag_tests.py --quick         # 빠른 테스트
"""

import os
import sys
import argparse
import subprocess
from datetime import datetime

def print_banner():
    """배너 출력"""
    print("="*80)
    print("RAG 시스템 테스트 및 튜닝 도구")
    print("실행 시간:", datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
    print("작업 디렉토리:", os.getcwd())
    print("="*80)

def check_dependencies():
    """종속성 확인"""
    print("종속성 확인 중...")
    
    required_files = [
        'rag_system.py',
        'config.py', 
        'law_fetcher.py',
        'test_rag.py',
        'tune_rag_config.py'
    ]
    
    missing_files = []
    for file in required_files:
        if not os.path.exists(file):
            missing_files.append(file)
    
    if missing_files:
        print(f"누락된 파일: {', '.join(missing_files)}")
        return False
    
    print("모든 필수 파일 확인 완료")
    return True

def run_basic_tests(rebuild=False):
    """기본 테스트 실행"""
    print("\n🧪 기본 RAG 테스트 실행")
    print("-" * 40)
    
    cmd = [sys.executable, 'test_rag.py']
    if rebuild:
        cmd.append('--rebuild')
    
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, encoding='utf-8', errors='replace')
        
        print("📤 테스트 출력:")
        if result.stdout:
            print(result.stdout)
        
        if result.stderr:
            print("⚠️ 경고/오류:")
            print(result.stderr)
        
        if result.returncode == 0:
            print("✅ 기본 테스트 성공")
            return True
        else:
            print(f"❌ 기본 테스트 실패 (exit code: {result.returncode})")
            return False
            
    except Exception as e:
        print(f"💥 테스트 실행 오류: {e}")
        return False

def run_config_tuning():
    """설정 튜닝 실행"""
    print("\n🎛️ RAG 설정 튜닝 실행")
    print("-" * 40)
    
    # 벡터 DB 존재 확인
    required_files = ['vector_db.faiss', 'documents.pkl', 'metadata.json']
    missing_files = [f for f in required_files if not os.path.exists(f)]
    
    if missing_files:
        print(f"❌ 벡터 DB 파일 누락: {', '.join(missing_files)}")
        print("📋 먼저 RAG 시스템을 구축해주세요:")
        print("   python run_rag_tests.py --rebuild")
        return False
    
    try:
        result = subprocess.run([sys.executable, 'tune_rag_config.py'], 
                              capture_output=True, text=True, encoding='utf-8', errors='replace')
        
        print("📤 튜닝 출력:")
        if result.stdout:
            print(result.stdout)
        
        if result.stderr:
            print("⚠️ 경고/오류:")
            print(result.stderr)
        
        if result.returncode == 0:
            print("✅ 설정 튜닝 성공")
            
            # 최적화된 설정 파일 확인
            if os.path.exists('optimized_config.py'):
                print("📄 최적화된 설정 파일 생성: optimized_config.py")
            
            return True
        else:
            print(f"❌ 설정 튜닝 실패 (exit code: {result.returncode})")
            return False
            
    except Exception as e:
        print(f"💥 튜닝 실행 오류: {e}")
        return False

def run_benchmark_only():
    """벤치마크만 실행"""
    print("\n🏃 성능 벤치마크 실행")
    print("-" * 40)
    
    try:
        # test_rag.py에서 벤치마크 부분만 실행하도록 별도 스크립트 작성
        benchmark_code = '''
import sys
sys.path.append('.')
from test_rag import RAGTester

tester = RAGTester()
if tester.initialize_rag():
    tester.performance_benchmark(iterations=10)
else:
    print("❌ RAG 시스템 초기화 실패")
'''
        
        result = subprocess.run([sys.executable, '-c', benchmark_code], 
                              capture_output=True, text=True, encoding='utf-8', errors='replace')
        
        print("📤 벤치마크 결과:")
        if result.stdout:
            print(result.stdout)
        
        if result.stderr:
            print("⚠️ 경고/오류:")
            print(result.stderr)
        
        return result.returncode == 0
        
    except Exception as e:
        print(f"💥 벤치마크 실행 오류: {e}")
        return False

def run_quick_test():
    """빠른 테스트 실행"""
    print("\n⚡ 빠른 RAG 테스트 실행")
    print("-" * 40)
    
    try:
        quick_test_code = '''
import sys
sys.path.append('.')
from test_rag import RAGTester

tester = RAGTester()
if tester.initialize_rag():
    # 3개 쿼리만 테스트
    tester.test_queries = tester.test_queries[:3]
    success = tester.run_basic_tests()
    print(f"\\n빠른 테스트 결과: {'성공' if success else '실패'}")
else:
    print("❌ RAG 시스템 초기화 실패")
'''
        
        result = subprocess.run([sys.executable, '-c', quick_test_code], 
                              capture_output=True, text=True, encoding='utf-8')
        
        print("📤 빠른 테스트 결과:")
        print(result.stdout)
        
        if result.stderr:
            print("⚠️ 경고/오류:")
            print(result.stderr)
        
        return result.returncode == 0
        
    except Exception as e:
        print(f"💥 빠른 테스트 실행 오류: {e}")
        return False

def show_system_info():
    """시스템 정보 표시"""
    print("\n💻 시스템 정보")
    print("-" * 40)
    
    print(f"🐍 Python 버전: {sys.version}")
    print(f"📁 현재 디렉토리: {os.getcwd()}")
    
    # 벡터 DB 상태 확인
    vector_files = ['vector_db.faiss', 'documents.pkl', 'metadata.json']
    vector_status = []
    
    for file in vector_files:
        if os.path.exists(file):
            size = os.path.getsize(file) / 1024  # KB
            vector_status.append(f"✅ {file} ({size:.1f}KB)")
        else:
            vector_status.append(f"❌ {file} (없음)")
    
    print("📊 벡터 DB 상태:")
    for status in vector_status:
        print(f"   {status}")
    
    # 캐시 상태 확인
    cache_dir = './law_cache'
    if os.path.exists(cache_dir):
        cache_files = len([f for f in os.listdir(cache_dir) if f.endswith('.json')])
        print(f"💾 법령 캐시: {cache_files}개 파일")
    else:
        print("💾 법령 캐시: 없음")

def main():
    """메인 함수"""
    parser = argparse.ArgumentParser(description='RAG 시스템 테스트 및 튜닝 도구')
    parser.add_argument('--rebuild', action='store_true', help='RAG 시스템 재구축')
    parser.add_argument('--tune', action='store_true', help='설정 튜닝 실행')
    parser.add_argument('--benchmark', action='store_true', help='성능 벤치마크만 실행')
    parser.add_argument('--quick', action='store_true', help='빠른 테스트 실행')
    parser.add_argument('--info', action='store_true', help='시스템 정보만 표시')
    
    args = parser.parse_args()
    
    print_banner()
    
    # 시스템 정보만 표시하고 종료
    if args.info:
        show_system_info()
        return
    
    # 종속성 확인
    if not check_dependencies():
        print("❌ 종속성 확인 실패")
        sys.exit(1)
    
    show_system_info()
    
    # 실행할 작업 결정
    tasks = []
    
    if args.quick:
        tasks.append(('빠른 테스트', run_quick_test))
    elif args.benchmark:
        tasks.append(('벤치마크', run_benchmark_only))
    elif args.tune:
        tasks.append(('설정 튜닝', run_config_tuning))
    else:
        # 기본 테스트
        tasks.append(('기본 테스트', lambda: run_basic_tests(args.rebuild)))
        
        # 튜닝도 함께 실행 (벡터 DB가 있는 경우)
        if all(os.path.exists(f) for f in ['vector_db.faiss', 'documents.pkl', 'metadata.json']):
            print("\n📋 벡터 DB가 존재하므로 설정 튜닝도 실행합니다.")
            tasks.append(('설정 튜닝', run_config_tuning))
    
    # 작업 실행
    success_count = 0
    for task_name, task_func in tasks:
        print(f"\n🚀 {task_name} 시작...")
        try:
            if task_func():
                success_count += 1
                print(f"✅ {task_name} 완료")
            else:
                print(f"❌ {task_name} 실패")
        except KeyboardInterrupt:
            print(f"\n🛑 사용자에 의해 {task_name}이 중단되었습니다.")
            break
        except Exception as e:
            print(f"💥 {task_name} 중 오류: {e}")
    
    # 결과 요약
    print("\n" + "="*80)
    print("📊 실행 결과 요약")
    print("="*80)
    print(f"✅ 성공한 작업: {success_count}/{len(tasks)}")
    
    if success_count == len(tasks):
        print("🎉 모든 작업이 성공적으로 완료되었습니다!")
        
        # 생성된 파일들 안내
        result_files = []
        for pattern in ['test_results_*.json', 'rag_tuning_results_*.json', 'optimized_config.py']:
            import glob
            result_files.extend(glob.glob(pattern))
        
        if result_files:
            print(f"📄 생성된 결과 파일: {', '.join(result_files)}")
        
        sys.exit(0)
    else:
        print("⚠️ 일부 작업에서 문제가 발생했습니다.")
        sys.exit(1)

if __name__ == "__main__":
    main()