#!/usr/bin/env python3
"""
Test the complete encryption/decryption cycle with agent v0 functions.
"""

import os
import sys
import tempfile
import json
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_encryption_cycle():
    """Test creating, encrypting, and decrypting a session."""
    
    print("=== Testing Agent V0 Encryption Cycle ===")
    print("Starting test...")
    
    # Load environment
    env_local_path = Path(".env.local")
    print(f"Looking for env file: {env_local_path}")
    if env_local_path.exists():
        print("Loading environment variables...")
        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()
    else:
        print("No .env.local file found")
    
    password = os.getenv("ENCRYPTION_PASSWORD")
    print(f"Using password: {password[:16] if password else 'None'}...")
    
    print("Attempting imports...")
    try:
        from raw_session import RawSession
        print("RawSession imported successfully")
        from storage_encrypted import create_session_zip_encrypted, decrypt_session_file
        print("Storage functions imported successfully")
        
        # Create a test session
        test_session = RawSession.create(
            user_id="test_encryption_cycle",
            platform="linux",
            hostname="test_host",
            screen_resolution=[1920, 1080]
        )
        
        # Add some test data
        test_session.add_event("click", {"x": 100, "y": 200}, datetime.now())
        test_session.add_event("key", {"key": "Enter"}, datetime.now())
        
        print(f"Created test session: {test_session.session_id}")
        
        with tempfile.TemporaryDirectory() as tmpdir:
            # Save the session
            test_session.save_json(tmpdir)
            print(f"Session saved to: {tmpdir}")
            
            # Encrypt the session
            encrypted_path = create_session_zip_encrypted(
                test_session,
                password,
                tmpdir,
                delete_unencrypted=False  # Keep original for comparison
            )
            print(f"Encryption successful: {encrypted_path}")
            print(f"   Encrypted file size: {os.path.getsize(encrypted_path)} bytes")
            
            # Try to decrypt it
            decrypted_path = decrypt_session_file(encrypted_path, password)
            print(f"Decryption successful: {decrypted_path}")
            
            # Verify it's a valid ZIP
            import zipfile
            with zipfile.ZipFile(decrypted_path, 'r') as zf:
                files = zf.namelist()
                print(f"   ZIP contains {len(files)} files: {files}")
            
            print("Complete encryption/decryption cycle successful!")
            return True
            
    except Exception as e:
        print(f"Test failed: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    success = test_encryption_cycle()
    sys.exit(0 if success else 1)