# Copyright (c) Subzero Labs, Inc.
# SPDX-License-Identifier: Apache-2.0

"""
Memory usage monitoring and leak detection tests for the Rialo Python CDK.

These tests monitor memory usage patterns and detect potential memory leaks
in long-running applications.
"""

import gc
import time
import psutil
import pytest
import threading
import weakref
from typing import List, Dict, Optional
from rialo_py_cdk import (
    # Core functions
    generate_keypair,
    keypair_from_seed,
    keypair_from_mnemonic,
    generate_mnemonic,
    public_key_from_keypair,
    sign,
    verify,
    # Wallet and provider types
    Wallet,
    # Core types
    PublicKey,
    Hash,
    Signature,
    # Transaction types
    TransactionBuilder,
    AccountMeta,
    Instruction,
    transfer_instruction,
    # Utility functions
    create_wallet_with_mnemonic,
    recover_wallet_from_mnemonic,
    # Constants
    get_localnet_url,
)


class MemoryTracker:
    """Helper class to track memory usage over time."""
    
    def __init__(self, name: str):
        self.name = name
        self.process = psutil.Process()
        self.snapshots = []
        self.start_time = None
        self.peak_memory = 0
        self.baseline_memory = 0
    
    def start_tracking(self):
        """Start memory tracking."""
        gc.collect()  # Force garbage collection
        self.start_time = time.time()
        self.baseline_memory = self.process.memory_info().rss
        self.peak_memory = self.baseline_memory
        self.snapshots = [(0, self.baseline_memory)]
        return self
    
    def take_snapshot(self, label: str = ""):
        """Take a memory snapshot."""
        if self.start_time is None:
            return
        
        current_time = time.time() - self.start_time
        current_memory = self.process.memory_info().rss
        self.peak_memory = max(self.peak_memory, current_memory)
        self.snapshots.append((current_time, current_memory, label))
    
    def stop_tracking(self):
        """Stop tracking and return summary."""
        if not self.snapshots:
            return None
        
        gc.collect()  # Force garbage collection
        final_memory = self.process.memory_info().rss
        final_time = time.time() - self.start_time if self.start_time else 0
        
        return {
            'name': self.name,
            'baseline_mb': self.baseline_memory / (1024 * 1024),
            'peak_mb': self.peak_memory / (1024 * 1024),
            'final_mb': final_memory / (1024 * 1024),
            'max_increase_mb': (self.peak_memory - self.baseline_memory) / (1024 * 1024),
            'net_increase_mb': (final_memory - self.baseline_memory) / (1024 * 1024),
            'duration_seconds': final_time,
            'snapshots': len(self.snapshots),
        }
    
    def print_summary(self):
        """Print memory usage summary."""
        summary = self.stop_tracking()
        if summary:
            print(f"\nMemory Usage Summary - {summary['name']}:")
            print(f"  Baseline: {summary['baseline_mb']:.2f} MB")
            print(f"  Peak: {summary['peak_mb']:.2f} MB")
            print(f"  Final: {summary['final_mb']:.2f} MB")
            print(f"  Max increase: {summary['max_increase_mb']:.2f} MB")
            print(f"  Net increase: {summary['net_increase_mb']:.2f} MB")
            print(f"  Duration: {summary['duration_seconds']:.2f} seconds")
            return summary
        return None


class TestBasicMemoryUsage:
    """Test basic memory usage patterns."""
    
    def test_cryptographic_operations_memory(self):
        """Test memory usage of cryptographic operations with larger payloads."""
        tracker = MemoryTracker("Cryptographic Operations").start_tracking()

        keypairs = []
        signatures = []
        large_messages = []

        # Generate keypairs and signatures with large messages
        for i in range(1000):  # Increased from 100 to 1000
            keypair = generate_keypair()
            keypairs.append(keypair)

            # Create a larger message (10KB instead of small strings)
            large_message = f"test message {i} " + "x" * (10 * 1024)  # 10KB message
            large_message_bytes = large_message.encode()
            large_messages.append(large_message_bytes)

            signature = sign(keypair, large_message_bytes)
            signatures.append(signature)

            if i % 200 == 0:  # Adjusted snapshot frequency
                tracker.take_snapshot(f"After {i+1} operations")

        tracker.take_snapshot("After all operations")
        summary = tracker.print_summary()

        # Memory usage should be reasonable (increased limit due to larger payloads)
        assert summary['max_increase_mb'] < 50  # Should not use more than 50MB

        # Clear references
        del keypairs, signatures, large_messages
        gc.collect()

        final_tracker = MemoryTracker("After cleanup").start_tracking()
        final_summary = final_tracker.print_summary()

        # Memory should be mostly recovered
        if summary['max_increase_mb'] > 0.1:  # Only check recovery if significant memory was used
            memory_recovered = summary['max_increase_mb'] - final_summary['net_increase_mb']
            assert memory_recovered > summary['max_increase_mb'] * 0.5  # At least 50% recovered
        else:
            print(f"Memory usage too small to measure recovery: {summary['max_increase_mb']:.2f} MB")
    
    def test_wallet_creation_memory(self):
        """Test memory usage of wallet creation with larger wallet names."""
        tracker = MemoryTracker("Wallet Creation").start_tracking()

        wallets = []
        for i in range(200):  # Increased from 50 to 200
            keypair = generate_keypair()
            # Create wallet with large name to ensure memory usage
            wallet_name = f"memory_test_wallet_{i}_" + "z" * 5000  # 5KB name
            wallet = Wallet(wallet_name, keypair)
            wallets.append(wallet)

            if i % 50 == 0:
                tracker.take_snapshot(f"After {i+1} wallets")

        summary = tracker.print_summary()

        # Memory usage should scale with larger wallets (increased limit)
        assert summary['max_increase_mb'] < 100  # Should not use more than 100MB for 200 large wallets

        # Test cleanup
        del wallets
        gc.collect()
    
    def test_type_creation_memory(self):
        """Test memory usage of core type creation."""
        tracker = MemoryTracker("Type Creation").start_tracking()
        
        pubkeys = []
        hashes = []
        signatures = []
        
        # Create many objects
        for i in range(1000):
            pubkey_bytes = bytes([(i % 256)] * 32)
            hash_bytes = bytes([((i+1) % 256)] * 32)
            sig_bytes = bytes([((i+2) % 256)] * 64)
            
            pubkeys.append(PublicKey(pubkey_bytes))
            hashes.append(Hash(hash_bytes))
            signatures.append(Signature(sig_bytes))
            
            if i % 200 == 0:
                tracker.take_snapshot(f"After {i+1} objects")
        
        summary = tracker.print_summary()
        
        # Memory usage should be reasonable even for many objects
        assert summary['max_increase_mb'] < 50  # Should not use more than 50MB for 3000 objects
        
        # Test that objects are properly allocated
        assert len(pubkeys) == 1000
        assert len(hashes) == 1000
        assert len(signatures) == 1000


class TestMemoryLeakDetection:
    """Test for memory leaks in various scenarios."""
    
    def test_repeated_keypair_generation_leak(self):
        """Test for memory leaks in repeated keypair generation."""
        tracker = MemoryTracker("Repeated Keypair Generation").start_tracking()
        
        baseline_memory = tracker.baseline_memory
        
        # Generate many keypairs without keeping references
        for i in range(500):
            keypair = generate_keypair()
            # Don't keep reference - should be garbage collected
            del keypair
            
            # Force garbage collection periodically
            if i % 100 == 0:
                gc.collect()
                tracker.take_snapshot(f"After {i+1} keypairs")
        
        # Final garbage collection
        gc.collect()
        summary = tracker.print_summary()
        
        # Memory increase should be minimal since we're not keeping references
        assert summary['net_increase_mb'] < 5  # Should not grow more than 5MB
        
        print(f"  Memory leak test - net increase: {summary['net_increase_mb']:.2f} MB")
    
    
    def test_signature_operations_leak(self):
        """Test for memory leaks in signature operations."""
        tracker = MemoryTracker("Signature Operations").start_tracking()
        
        # Create a keypair once
        keypair = generate_keypair()
        pubkey = public_key_from_keypair(keypair)
        
        # Perform many sign/verify operations
        for i in range(200):
            message = f"leak test message {i}".encode()
            
            # Sign message
            signature = sign(keypair, message)
            
            # Verify signature
            verified = verify(pubkey, message, signature)
            assert verified is True
            
            # Don't keep references
            del signature, message
            
            if i % 50 == 0:
                gc.collect()
                tracker.take_snapshot(f"After {i+1} sign/verify operations")
        
        # Final cleanup
        gc.collect()
        summary = tracker.print_summary()
        
        # Memory increase should be minimal
        assert summary['net_increase_mb'] < 5  # Should not use more than 5MB
        
        print(f"  Signature operations leak test - net increase: {summary['net_increase_mb']:.2f} MB")
    
    def test_weak_references_cleanup(self):
        """Test that objects are properly cleaned up using weak references."""
        weak_refs = []
        strong_refs = []  # Keep strong references to ensure objects stay alive initially

        # Create objects with larger payloads to ensure they're substantial
        for i in range(100):
            keypair = generate_keypair()
            # Create wallet with large name to increase memory footprint
            wallet_name = f"weak_ref_test_wallet_{i}_" + "x" * 1000  # 1KB name
            wallet = Wallet(wallet_name, keypair)

            # Keep both strong and weak references initially
            strong_refs.append(wallet)
            weak_refs.append(weakref.ref(wallet))

        # Objects should exist with strong references
        alive_count = sum(1 for ref in weak_refs if ref() is not None)
        assert alive_count == 100, f"Expected 100 objects alive with strong refs, got {alive_count}"

        # Force garbage collection - should not affect objects with strong references
        gc.collect()

        # Objects should still be alive
        alive_after_gc = sum(1 for ref in weak_refs if ref() is not None)
        print(f"  Objects alive after GC (with strong refs): {alive_after_gc}/100")
        assert alive_after_gc == 100, f"Strong references should keep objects alive, got {alive_after_gc}"

        # Clear strong references but keep weak refs
        del strong_refs
        gc.collect()

        # Now objects should be cleaned up
        final_alive = sum(1 for ref in weak_refs if ref() is not None)
        print(f"  Objects alive after clearing strong references: {final_alive}/100")

        # Most objects should be cleaned up (allow some to remain due to GC timing)
        assert final_alive < 50, f"Too many objects still alive after cleanup: {final_alive}/100"


class TestLongRunningApplications:
    """Test memory behavior in long-running application scenarios."""
    
    def test_continuous_operations_memory(self):
        """Test memory usage during continuous operations."""
        tracker = MemoryTracker("Continuous Operations").start_tracking()
        
        # Simulate long-running application with continuous operations
        iteration_count = 0
        max_iterations = 300
        
        while iteration_count < max_iterations:
            # Generate and use keypair
            keypair = generate_keypair()
            pubkey = public_key_from_keypair(keypair)
            
            # Create wallet
            wallet = Wallet(f"continuous_wallet_{iteration_count}", keypair)
            
            # Sign a message
            message = f"continuous test {iteration_count}".encode()
            signature = sign(keypair, message)
            verified = verify(pubkey, message, signature)
            assert verified is True
            
            # Don't keep references (simulate processing and discarding)
            del keypair, pubkey, wallet, signature
            
            iteration_count += 1
            
            # Take periodic snapshots
            if iteration_count % 50 == 0:
                gc.collect()  # Force GC periodically
                tracker.take_snapshot(f"After {iteration_count} iterations")
        
        # Final cleanup and summary
        gc.collect()
        summary = tracker.print_summary()
        
        # Memory should not grow significantly over time
        assert summary['net_increase_mb'] < 10  # Should not use more than 10MB net
        
        # Peak memory should be reasonable
        assert summary['max_increase_mb'] < 20  # Peak should not exceed 20MB
        
        print(f"  Continuous operations - {max_iterations} iterations")
        print(f"  Net memory increase: {summary['net_increase_mb']:.2f} MB")
        print(f"  Peak memory increase: {summary['max_increase_mb']:.2f} MB")
    
    def test_batch_processing_memory(self):
        """Test memory usage during batch processing scenarios."""
        tracker = MemoryTracker("Batch Processing").start_tracking()
        
        batch_size = 50
        num_batches = 10
        
        for batch_num in range(num_batches):
            # Process a batch of operations
            batch_data = []
            
            for i in range(batch_size):
                # Generate test data
                keypair = generate_keypair()
                message = f"batch {batch_num} item {i}".encode()
                signature = sign(keypair, message)
                
                batch_data.append({
                    'keypair': keypair,
                    'message': message,
                    'signature': signature,
                })
            
            # Process the batch
            for item in batch_data:
                pubkey = public_key_from_keypair(item['keypair'])
                verified = verify(pubkey, item['message'], item['signature'])
                assert verified is True
            
            # Clear batch data (simulate batch completion)
            del batch_data
            gc.collect()
            
            tracker.take_snapshot(f"After batch {batch_num + 1}")
        
        summary = tracker.print_summary()
        
        # Memory should be stable across batches
        assert summary['net_increase_mb'] < 15  # Should not accumulate more than 15MB
        
        print(f"  Batch processing - {num_batches} batches of {batch_size} items each")
        print(f"  Net memory increase: {summary['net_increase_mb']:.2f} MB")
    


class TestResourceUsageBenchmarks:
    """Benchmark resource usage for different operations."""
    
    def test_cpu_and_memory_benchmark(self):
        """Benchmark CPU and memory usage for various operations."""
        import time
        
        results = {}
        
        # Benchmark keypair generation
        tracker = MemoryTracker("Keypair Generation Benchmark").start_tracking()
        start_time = time.perf_counter()
        
        for _ in range(100):
            keypair = generate_keypair()
            del keypair  # Don't keep reference
        
        end_time = time.perf_counter()
        memory_summary = tracker.print_summary()
        
        results['keypair_generation'] = {
            'operations': 100,
            'time_seconds': end_time - start_time,
            'ops_per_second': 100 / (end_time - start_time),
            'memory_mb': memory_summary['max_increase_mb'],
        }
        
        # Benchmark signature operations
        tracker = MemoryTracker("Signature Benchmark").start_tracking()
        start_time = time.perf_counter()
        
        keypair = generate_keypair()
        message = b"benchmark message"
        
        for _ in range(100):
            signature = sign(keypair, message)
            pubkey = public_key_from_keypair(keypair)
            verified = verify(pubkey, message, signature)
            assert verified is True
        
        end_time = time.perf_counter()
        memory_summary = tracker.print_summary()
        
        results['signature_operations'] = {
            'operations': 100,
            'time_seconds': end_time - start_time,
            'ops_per_second': 100 / (end_time - start_time),
            'memory_mb': memory_summary['max_increase_mb'],
        }
        
        # Benchmark mnemonic operations
        tracker = MemoryTracker("Mnemonic Benchmark").start_tracking()
        start_time = time.perf_counter()
        
        for _ in range(20):  # Fewer operations as these are more expensive
            mnemonic = generate_mnemonic()
            keypair = keypair_from_mnemonic(mnemonic)
            del mnemonic, keypair
        
        end_time = time.perf_counter()
        memory_summary = tracker.print_summary()
        
        results['mnemonic_operations'] = {
            'operations': 20,
            'time_seconds': end_time - start_time,
            'ops_per_second': 20 / (end_time - start_time),
            'memory_mb': memory_summary['max_increase_mb'],
        }
        
        # Print benchmark results
        print("\nResource Usage Benchmarks:")
        for operation, metrics in results.items():
            print(f"  {operation}:")
            print(f"    {metrics['ops_per_second']:.2f} ops/sec")
            print(f"    {metrics['memory_mb']:.2f} MB peak memory")
            print(f"    {metrics['time_seconds']:.3f}s total time")
        
        # Performance assertions
        assert results['keypair_generation']['ops_per_second'] > 10  # At least 10 ops/sec
        assert results['signature_operations']['ops_per_second'] > 10  # At least 10 ops/sec
        assert results['mnemonic_operations']['ops_per_second'] > 1   # At least 1 ops/sec
        
        # Memory assertions
        for operation, metrics in results.items():
            assert metrics['memory_mb'] < 50  # No operation should use more than 50MB


if __name__ == "__main__":
    # Add psutil to test requirements
    try:
        import psutil
    except ImportError:
        print("psutil is required for memory tests. Install with: pip install psutil")
        exit(1)
    
    pytest.main([__file__, "-v", "-s"])