#!/usr/bin/env python3
"""
개선된 updater 테스트
"""

import os
import sys
import tempfile
import shutil

def create_fake_pe_exe(path, content_prefix, size_mb=5):
    """가짜 PE 실행 파일 생성 (유효한 MZ 헤더 포함)"""
    # 기본 PE 헤더 (MZ signature + minimal DOS header)
    dos_header = bytearray(64)
    dos_header[0:2] = b'MZ'  # MZ signature
    dos_header[60:64] = b'\x80\x00\x00\x00'  # PE header offset
    
    # 가짜 PE header
    pe_header = b'PE\x00\x00' + b'\x00' * 60
    
    # 실제 내용
    content = content_prefix.encode('utf-8')
    
    # 지정된 크기까지 패딩
    target_size = size_mb * 1024 * 1024
    header_size = len(dos_header) + len(pe_header)
    content_size = len(content)
    padding_size = max(0, target_size - header_size - content_size)
    
    # 파일 생성
    with open(path, 'wb') as f:
        f.write(dos_header)
        f.write(pe_header)
        f.write(content)
        f.write(b'0' * padding_size)
    
    print(f"Created fake PE exe: {path} ({os.path.getsize(path):,} bytes)")

def test_improved_update():
    """개선된 업데이트 로직 테스트"""
    print("=" * 60)
    print("Improved Update Test")
    print("=" * 60)
    
    temp_dir = tempfile.mkdtemp(prefix="improved_update_")
    print(f"Test directory: {temp_dir}")
    
    try:
        # 가짜 PE 파일들 생성
        old_exe = os.path.join(temp_dir, "app.exe")
        new_exe = os.path.join(temp_dir, "app_new.exe")
        
        create_fake_pe_exe(old_exe, "OLD_VERSION_CONTENT", 8)
        create_fake_pe_exe(new_exe, "NEW_VERSION_CONTENT", 10)
        
        # 파일 크기 확인
        old_size = os.path.getsize(old_exe)
        new_size = os.path.getsize(new_exe)
        print(f"\nInitial sizes:")
        print(f"  Old exe: {old_size:,} bytes")
        print(f"  New exe: {new_size:,} bytes")
        
        # PE 헤더 검증 테스트
        print(f"\nPE header verification:")
        for path, name in [(old_exe, "old"), (new_exe, "new")]:
            with open(path, 'rb') as f:
                header = f.read(64)
                if header.startswith(b'MZ'):
                    print(f"✓ {name} exe has valid MZ signature")
                else:
                    print(f"✗ {name} exe has invalid signature")
        
        # 백업 생성
        backup_path = f"{old_exe}.backup"
        shutil.copy2(old_exe, backup_path)
        print(f"\n✓ Backup created: {backup_path}")
        
        # 안전한 파일 교체 시뮬레이션
        print(f"\nPerforming safe file replacement...")
        
        # 1. 임시 위치에 복사
        temp_target = f"{old_exe}.tmp"
        shutil.copy2(new_exe, temp_target)
        print(f"✓ Copied to temporary location")
        
        # 2. 복사된 파일 검증
        copied_size = os.path.getsize(temp_target)
        if copied_size == new_size:
            print(f"✓ Size verification passed: {copied_size:,} bytes")
        else:
            print(f"✗ Size mismatch: {copied_size} != {new_size}")
            return False
        
        # 3. PE 헤더 재검증
        with open(temp_target, 'rb') as f:
            header = f.read(64)
            if header.startswith(b'MZ'):
                print(f"✓ Temporary file PE header valid")
            else:
                print(f"✗ Temporary file PE header invalid")
                return False
        
        # 4. 원자적 교체
        os.replace(temp_target, old_exe)
        print(f"✓ Atomic replacement completed")
        
        # 5. 최종 검증
        final_size = os.path.getsize(old_exe)
        print(f"✓ Final verification: {final_size:,} bytes")
        
        # 6. 내용 확인
        with open(old_exe, 'rb') as f:
            f.seek(128)  # PE 헤더 이후
            content = f.read(50)
            content_str = content.decode('utf-8', errors='ignore')
            print(f"Content preview: {content_str}")
            
        if b"NEW_VERSION" in content:
            print("✓ Content verification successful")
            
            # 백업 정리
            os.remove(backup_path)
            print("✓ Backup cleaned up")
            
            # 원본 새 파일 삭제
            if os.path.exists(new_exe):
                os.remove(new_exe)
                print("✓ Original new file cleaned up")
            
            return True
        else:
            print("✗ Content verification failed")
            return False
            
    except Exception as e:
        print(f"✗ Error during test: {e}")
        import traceback
        traceback.print_exc()
        return False
    finally:
        shutil.rmtree(temp_dir, ignore_errors=True)
        print(f"Test directory cleaned up")

def test_launch_logic():
    """실행 로직 테스트"""
    print("\n" + "=" * 60)
    print("Launch Logic Test")
    print("=" * 60)
    
    temp_dir = tempfile.mkdtemp(prefix="launch_test_")
    
    try:
        # 가짜 실행 파일 생성
        test_exe = os.path.join(temp_dir, "test_app.exe")
        create_fake_pe_exe(test_exe, "TEST_EXECUTABLE", 2)
        
        print(f"Created test executable: {test_exe}")
        
        # 실행 가능성 확인
        if os.path.exists(test_exe):
            size = os.path.getsize(test_exe)
            print(f"✓ Executable exists: {size:,} bytes")
            
            # PE 헤더 확인
            with open(test_exe, 'rb') as f:
                header = f.read(2)
                if header == b'MZ':
                    print("✓ Valid PE executable")
                    return True
                else:
                    print("✗ Invalid executable format")
                    return False
        else:
            print("✗ Executable not found")
            return False
            
    except Exception as e:
        print(f"✗ Error in launch test: {e}")
        return False
    finally:
        shutil.rmtree(temp_dir, ignore_errors=True)

if __name__ == "__main__":
    print("Testing improved updater logic...\n")
    
    # 파일 교체 테스트
    update_success = test_improved_update()
    
    # 실행 테스트
    launch_success = test_launch_logic()
    
    print("\n" + "=" * 60)
    print("Improved Test Results")
    print("=" * 60)
    print(f"Safe file replacement: {'✓ PASS' if update_success else '✗ FAIL'}")
    print(f"Launch preparation: {'✓ PASS' if launch_success else '✗ FAIL'}")
    
    if update_success and launch_success:
        print("\n✓ All improved tests passed")
        print("The updater should now:")
        print("  - Verify PE headers before and after replacement")
        print("  - Use atomic file operations")
        print("  - Properly launch applications on Windows")
        print("  - Handle file corruption issues")
    else:
        print("\n✗ Some tests failed - check implementation")