"""
Test 3: Per-agent isolation

Demonstrates:
- Multiple agents with isolated vaults
- Agent-specific access
- Export/import between agents
- Service-level isolation
"""

import os
import sys
import tempfile
import shutil

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

temp_dir = tempfile.mkdtemp()
os.environ["VAULT_DIR"] = temp_dir
os.environ["VAULT_AUDIT_LOG_DIR"] = temp_dir

from secure_vault import SecureVault


def test_agent_isolation():
    """Test per-agent credential isolation"""
    
    print("=== Test: Per-Agent Isolation ===\n")
    
    vault = SecureVault()
    vault.init_vault("master-password")
    
    # Create credentials for different agents
    print("1. Creating credentials for different agents...")
    
    # Sofiia's credentials
    vault.store("sofiia", "gmail", "password", "sofiia-gmail-pass")
    vault.store("sofiia", "github", "token", "sofiia-gh-token")
    
    # Helion's credentials
    vault.store("helion", "gmail", "password", "helion-gmail-pass")
    vault.store("helion", "aws", "access_key", "helion-aws-key")
    
    # Druid's credentials
    vault.store("druid", "slack", "webhook", "druid-slack-url")
    
    print("   - Sofiia: gmail, github")
    print("   - Helion: gmail, aws")
    print("   - Druid: slack")
    
    # List for each agent
    print("\n2. Listing services per agent...")
    print(f"   Sofiia: {vault.list('sofiia')}")
    print(f"   Helion: {vault.list('helion')}")
    print(f"   Druid: {vault.list('druid')}")
    
    # Verify isolation - can't access other agent's creds
    print("\n3. Verifying isolation...")
    
    # Sofiia tries to access Helion's AWS
    helion_aws = vault.get("sofiia", "helion", "access_key")
    print(f"   Sofiia accessing Helion's AWS: {helion_aws} (should be None)")
    
    # But can access own
    sofiia_gh = vault.get("sofiia", "github", "token")
    print(f"   Sofiia accessing own GitHub: {sofiia_gh}")
    
    # Helion tries to access Sofiia's Gmail
    sofiia_gmail = vault.get("helion", "gmail", "password")
    print(f"   Helion accessing Sofiia's Gmail: {sofiia_gmail} (should be None)")
    
    # Helion can access own
    helion_aws = vault.get("helion", "aws", "access_key")
    print(f"   Helion accessing own AWS: {helion_aws}")
    
    # Clean up
    shutil.rmtree(temp_dir)
    
    print("\n✅ Agent isolation test passed!")
    return True


def test_export_import():
    """Test export/import for P2P"""
    
    print("\n=== Test: Export/Import ===\n")
    
    # Source vault
    print("1. Creating source vault...")
    source_dir = tempfile.mkdtemp()
    os.environ["VAULT_DIR"] = source_dir
    
    vault = SecureVault()
    vault.init_vault("source-password")
    
    vault.store("agent1", "service1", "cred1", "value1")
    vault.store("agent1", "service1", "cred2", "value2")
    vault.store("agent1", "service2", "api_key", "secret-api-key")
    print("   - Stored 3 credentials")
    
    # Export for agent
    print("\n2. Exporting agent vault...")
    export = vault.export_for_agent("agent1")
    print(f"   Exported services: {export['services']}")
    print(f"   Export key provided: {len(export['export_key'])} chars")
    
    # Clean up source
    shutil.rmtree(source_dir)
    
    # Import into new vault
    print("\n3. Importing into new vault...")
    target_dir = tempfile.mkdtemp()
    os.environ["VAULT_DIR"] = target_dir
    
    vault2 = SecureVault()
    vault2.init_vault("target-password")
    
    result = vault2.import_for_agent(
        encrypted_vault=export["encrypted_vault"],
        export_key=export["export_key"],
        agent_id="agent1"
    )
    print(f"   Import status: {result['status']}")
    print(f"   Services imported: {result['services_imported']}")
    
    # Verify import
    print("\n4. Verifying imported credentials...")
    creds = vault2.list("agent1")
    print(f"   Services: {creds}")
    
    val1 = vault2.get("agent1", "service1", "cred1")
    print(f"   cred1: {val1}")
    
    # Clean up
    shutil.rmtree(target_dir)
    
    print("\n✅ Export/Import test passed!")
    return True


def test_service_isolation():
    """Test service-level isolation within same agent"""
    
    print("\n=== Test: Service-Level Isolation ===\n")
    
    temp_dir = tempfile.mkdtemp()
    os.environ["VAULT_DIR"] = temp_dir
    
    vault = SecureVault()
    vault.init_vault("password")
    
    # Store multiple services for same agent
    vault.store("sofiia", "gmail", "password", "gmail-pass")
    vault.store("sofiia", "github", "token", "github-token")
    vault.store("sofiia", "aws", "access_key", "aws-key")
    vault.store("sofiia", "aws", "secret_key", "aws-secret")
    
    print("1. Stored credentials for 3 services")
    
    # List all services
    services = vault.list("sofiia")
    print(f"\n2. All services: {services}")
    
    # List credentials per service
    print("\n3. Credentials per service:")
    for svc in services:
        creds = vault.list("sofiia", svc)
        print(f"   {svc}: {creds}")
    
    # Delete only AWS credentials
    print("\n4. Deleting AWS service...")
    vault.delete("sofiia", "aws", "access_key")
    vault.delete("sofiia", "aws", "secret_key")
    
    # Verify other services intact
    remaining = vault.list("sofiia")
    print(f"\n5. Remaining services: {remaining}")
    
    gmail = vault.get("sofiia", "gmail", "password")
    github = vault.get("sofiia", "github", "token")
    print(f"   Gmail: {gmail}")
    print(f"   GitHub: {github}")
    
    # Clean up
    shutil.rmtree(temp_dir)
    
    print("\n✅ Service isolation test passed!")
    return True


if __name__ == "__main__":
    test_agent_isolation()
    test_export_import()
    test_service_isolation()