#!/usr/bin/env python3
"""
통합 학습 시스템 테스트
train_ttrlvr_azr.py가 올바르게 작동하는지 확인
"""

import os
import sys
import subprocess
import time

# 환경 변수 설정
os.environ['CUDA_VISIBLE_DEVICES'] = '5'
os.environ['VLLM_ATTENTION_BACKEND'] = 'FLASH_ATTN'
os.environ['RAY_memory_monitor_refresh_ms'] = '0'
os.environ['RAY_LOGGING_LEVEL'] = 'DEBUG'
os.environ['HYDRA_FULL_ERROR'] = '1'
os.environ['PYTHONPATH'] = f"{os.environ.get('PYTHONPATH', '')}:/home/ubuntu/RLVR/verl:/home/ubuntu/RLVR/TestTime-RLVR-v2"


def test_small_run():
    """작은 규모로 통합 테스트"""
    
    print("=" * 80)
    print("TTRLVR + AZR 통합 학습 테스트")
    print("=" * 80)
    print()
    
    # 테스트 파라미터
    cmd = [
        sys.executable,
        "/home/ubuntu/RLVR/TestTime-RLVR-v2/test/train_ttrlvr_azr.py",
        "--benchmark", "mbpp",
        "--problems", "2",      # 2문제만 테스트
        "--rounds", "2",        # 2라운드만 테스트
        "--gpu", "5",
        "--batch-size", "1",    # 메모리 절약
        "--model", "Qwen/Qwen2.5-7B",
        "--debug"               # 디버그 모드
    ]
    
    print(f"명령어: {' '.join(cmd)}")
    print()
    print("환경 변수:")
    for key in ['CUDA_VISIBLE_DEVICES', 'VLLM_ATTENTION_BACKEND', 'PYTHONPATH']:
        print(f"  {key}: {os.environ.get(key, 'NOT SET')}")
    print()
    
    start_time = time.time()
    
    try:
        # subprocess 실행
        result = subprocess.run(
            cmd,
            capture_output=False,  # 실시간 출력 보기
            text=True
        )
        
        elapsed = time.time() - start_time
        
        print()
        print("=" * 80)
        print(f"테스트 완료 (소요 시간: {elapsed:.1f}초)")
        print(f"종료 코드: {result.returncode}")
        print("=" * 80)
        
        if result.returncode == 0:
            print("✅ 성공!")
            
            # 생성된 데이터 확인
            check_generated_data()
        else:
            print("❌ 실패!")
            
    except Exception as e:
        print(f"💥 예외 발생: {e}")


def check_generated_data():
    """생성된 데이터 확인"""
    
    print("\n생성된 데이터 확인:")
    
    # 첫 번째 라운드 데이터 확인
    round1_path = "/tmp/ttrlvr_azr_training/round_1"
    if os.path.exists(round1_path):
        print(f"\n📁 {round1_path}:")
        for file in os.listdir(round1_path):
            file_path = os.path.join(round1_path, file)
            if os.path.isfile(file_path):
                size = os.path.getsize(file_path)
                print(f"  - {file}: {size:,} bytes")
                
                if file.endswith('.parquet'):
                    try:
                        import pandas as pd
                        df = pd.read_parquet(file_path)
                        print(f"    행 수: {len(df)}")
                    except:
                        pass
    
    # 체크포인트 확인
    checkpoint_base = "/data/RLVR/checkpoints/ttrlvr_azr"
    if os.path.exists(checkpoint_base):
        print(f"\n📁 체크포인트 디렉토리: {checkpoint_base}")
        for exp_dir in os.listdir(checkpoint_base):
            exp_path = os.path.join(checkpoint_base, exp_dir)
            if os.path.isdir(exp_path):
                print(f"  - {exp_dir}/")
                for checkpoint in os.listdir(exp_path):
                    print(f"    - {checkpoint}")


def test_subprocess_azr():
    """subprocess AZR 호출만 테스트"""
    
    print("\n" + "=" * 80)
    print("Subprocess AZR 호출 테스트")
    print("=" * 80)
    
    # 테스트 데이터가 있는지 확인
    test_data_path = "/home/ubuntu/RLVR/TestTime-RLVR-v2/tmp/batch_results/ttrlvr_azr_20250730_155352/mbpp/Mbpp_7/round_1/azr_training_data"
    
    if not os.path.exists(test_data_path):
        print(f"테스트 데이터 없음: {test_data_path}")
        return
    
    sys.path.append('/home/ubuntu/RLVR/TestTime-RLVR-v2/test')
    from utils.run_azr_subprocess import run_azr_training
    
    success = run_azr_training(
        training_data_path=test_data_path,
        experiment_name="test_subprocess",
        gpu_id=5
    )
    
    print(f"\nSubprocess 결과: {'성공' if success else '실패'}")


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser()
    parser.add_argument('--mode', choices=['full', 'subprocess'], default='full',
                      help='테스트 모드 선택')
    
    args = parser.parse_args()
    
    if args.mode == 'full':
        test_small_run()
    else:
        test_subprocess_azr()