#!/usr/bin/env python3
"""
Real Functionality Test: Fresh Encryption Cycle

Tests the complete encryption/decryption workflow using:
- Real RawSession data with actual events
- Real file system operations
- Real encryption algorithms (AES-256)
- Real ZIP file creation and validation
- Actual environment variable loading

This test validates the entire encryption pipeline as it would work in production.
"""

import os
import sys
import tempfile
import json
import hashlib
from pathlib import Path
from datetime import datetime

# Add agent_v0 to path
sys.path.insert(0, str(Path(__file__).parent / "agent_v0"))

def test_fresh_encryption():
    """Test complete encryption cycle with real data and operations."""
    
    print("=== Real Functionality Test: Fresh Encryption Cycle ===")
    
    # Load real environment configuration
    env_local_path = Path(".env.local")
    if env_local_path.exists():
        with open(env_local_path, 'r') as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith('#') and '=' in line:
                    key, value = line.split('=', 1)
                    os.environ[key.strip()] = value.strip()
    
    password = os.getenv("ENCRYPTION_PASSWORD")
    if not password:
        print("ERROR: ENCRYPTION_PASSWORD not found in environment")
        return False
        
    print(f"Using encryption password: {password[:4]}{'*' * (len(password) - 4)}")
    
    # Import real agent modules (no mocks)
    from raw_session import RawSession
    from storage_encrypted import create_session_zip_encrypted, decrypt_session_file
    
    # Create realistic test session with actual user interaction data
    session = RawSession.create(
        user_id="test_encryption_user",
        platform="linux",
        hostname="test_workstation",
        screen_resolution=[1920, 1080]
    )
    
    # Add realistic interaction events (simulating real user actions)
    base_time = datetime.now()
    
    # Simulate a realistic workflow: login sequence
    session.add_mouse_click_event("left", [450, 320], "Login Form", "firefox", None)
    session.add_key_combo_event(["CTRL", "A"], "Login Form", "firefox", None)
    session.add_text_input_event("user@example.com", "Login Form", "firefox", None)
    session.add_key_press_event("Tab", "Login Form", "firefox", None)
    session.add_text_input_event("password123", "Login Form", "firefox", None)
    session.add_mouse_click_event("left", [500, 400], "Login Form", "firefox", None)
    
    print(f"Created realistic session: {session.session_id}")
    print(f"Session contains {len(session.events)} interaction events")
    
    with tempfile.TemporaryDirectory() as tmpdir:
        tmpdir_path = Path(tmpdir)
        
        # Save session using real file operations
        session_dir = session.save_json(tmpdir)
        print(f"Session saved to: {session_dir}")
        
        # Verify session files exist and are valid
        session_json_path = Path(session_dir) / f"{session.session_id}.json"
        assert session_json_path.exists(), "Session JSON file not created"
        
        # Validate JSON content
        with open(session_json_path, 'r') as f:
            session_data = json.load(f)
            assert session_data["session_id"] == session.session_id
            assert len(session_data["events"]) == 6
            print("✓ Session JSON validation passed")
        
        # Test encryption with real AES-256 encryption
        print("Encrypting session with real AES-256...")
        encrypted_path = create_session_zip_encrypted(
            session,
            password,
            tmpdir,
            delete_unencrypted=False  # Keep for verification
        )
        
        print(f"Encrypted file created: {encrypted_path}")
        
        # Verify encrypted file properties
        encrypted_file = Path(encrypted_path)
        assert encrypted_file.exists(), "Encrypted file not created"
        
        file_size = encrypted_file.stat().st_size
        print(f"Encrypted file size: {file_size} bytes")
        assert file_size > 100, "Encrypted file suspiciously small"
        
        # Verify file is actually encrypted (not plain ZIP)
        with open(encrypted_path, 'rb') as f:
            header = f.read(4)
            # Should NOT be a standard ZIP header (PK\x03\x04)
            assert header != b'PK\x03\x04', "File appears to be unencrypted ZIP"
            print("✓ File is properly encrypted (not plain ZIP)")
        
        # Test decryption with real decryption algorithm
        print("Testing decryption with real algorithm...")
        try:
            decrypted_path = decrypt_session_file(encrypted_path, password)
            print(f"SUCCESS: Decrypted to {decrypted_path}")
            
            # Verify decrypted file is valid ZIP
            import zipfile
            with zipfile.ZipFile(decrypted_path, 'r') as zf:
                files = zf.namelist()
                print(f"Decrypted ZIP contains {len(files)} files: {files}")
                
                # Verify session JSON is in the ZIP
                session_json_name = f"{session.session_id}.json"
                assert session_json_name in files, f"Session JSON {session_json_name} not found in ZIP"
                
                # Extract and validate session content
                with zf.open(session_json_name) as json_file:
                    decrypted_session_data = json.load(json_file)
                    assert decrypted_session_data["session_id"] == session.session_id
                    assert len(decrypted_session_data["events"]) == 6
                    print("✓ Decrypted session data validation passed")
            
            # Test wrong password fails appropriately
            print("Testing wrong password rejection...")
            try:
                wrong_password = password + "_wrong"
                decrypt_session_file(encrypted_path, wrong_password)
                print("ERROR: Wrong password was accepted!")
                return False
            except Exception as e:
                print(f"✓ Wrong password correctly rejected: {type(e).__name__}")
            
            # Test file integrity with checksum
            original_checksum = _calculate_file_checksum(session_json_path)
            
            # Extract decrypted session for comparison
            with zipfile.ZipFile(decrypted_path, 'r') as zf:
                with zf.open(session_json_name) as json_file:
                    decrypted_content = json_file.read()
                    
            temp_extracted = tmpdir_path / "extracted_session.json"
            with open(temp_extracted, 'wb') as f:
                f.write(decrypted_content)
                
            decrypted_checksum = _calculate_file_checksum(temp_extracted)
            
            if original_checksum == decrypted_checksum:
                print("✓ File integrity verified: checksums match")
            else:
                print(f"WARNING: Checksum mismatch - original: {original_checksum}, decrypted: {decrypted_checksum}")
            
            return True
            
        except Exception as e:
            print(f"FAILED: Decryption error: {e}")
            import traceback
            traceback.print_exc()
            return False


def _calculate_file_checksum(file_path: Path) -> str:
    """Calculate SHA-256 checksum of a file for integrity verification."""
    sha256_hash = hashlib.sha256()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            sha256_hash.update(chunk)
    return sha256_hash.hexdigest()


def test_encryption_performance():
    """Test encryption performance with realistic data sizes."""
    print("\n=== Performance Test: Encryption Speed ===")
    
    password = os.getenv("ENCRYPTION_PASSWORD", "test_password_123")
    
    # Import real modules
    from raw_session import RawSession
    from storage_encrypted import create_session_zip_encrypted
    
    # Create larger session for performance testing
    session = RawSession.create(
        user_id="perf_test_user",
        platform="linux", 
        hostname="perf_test",
        screen_resolution=[2560, 1440]
    )
    
    # Add many events to test with realistic data size
    base_time = datetime.now()
    for i in range(100):  # 100 events
        session.add_mouse_move_event([100 + i, 200 + i], f"App {i}", "test_app", None)
    
    print(f"Created performance test session with {len(session.events)} events")
    
    with tempfile.TemporaryDirectory() as tmpdir:
        session.save_json(tmpdir)
        
        # Measure encryption time
        import time
        start_time = time.time()
        
        encrypted_path = create_session_zip_encrypted(
            session,
            password,
            tmpdir,
            delete_unencrypted=True
        )
        
        encryption_time = time.time() - start_time
        file_size = Path(encrypted_path).stat().st_size
        
        print(f"Encryption completed in {encryption_time:.3f} seconds")
        print(f"Encrypted file size: {file_size} bytes")
        print(f"Encryption speed: {file_size / encryption_time / 1024:.1f} KB/s")
        
        # Performance should be reasonable (>100 KB/s)
        speed_kbps = file_size / encryption_time / 1024
        if speed_kbps > 100:
            print("✓ Encryption performance acceptable")
            return True
        else:
            print(f"WARNING: Encryption performance slow: {speed_kbps:.1f} KB/s")
            return False

if __name__ == "__main__":
    try:
        # Run main encryption test
        success1 = test_fresh_encryption()
        
        # Run performance test
        success2 = test_encryption_performance()
        
        overall_success = success1 and success2
        print(f"\nOverall test result: {'SUCCESS' if overall_success else 'FAILED'}")
        print(f"- Encryption cycle: {'PASS' if success1 else 'FAIL'}")
        print(f"- Performance test: {'PASS' if success2 else 'FAIL'}")
        
        sys.exit(0 if overall_success else 1)
    except Exception as e:
        print(f"Test exception: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)