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

"""
Performance regression tests for the Rialo Python CDK.

These tests measure performance of key operations to detect regressions.
"""

import gc
import time
import psutil
import pytest
import statistics
from concurrent.futures import ThreadPoolExecutor
from rialo_py_cdk import (
    # Core functions
    generate_keypair,
    keypair_from_seed,
    keypair_from_mnemonic,
    generate_mnemonic,
    public_key_from_keypair,
    sign,
    verify,
    # Wallet types
    Wallet,
    # Utility functions
    create_wallet_with_mnemonic,
    recover_wallet_from_mnemonic,
    rlo_to_kelvin,
    kelvin_to_rlo,
    # Core types
    PublicKey,
    Hash,
    Signature,
    # Transaction types
    TransactionBuilder,
    AccountMeta,
    Instruction,
    transfer_instruction,
    get_localnet_url,
)


class PerformanceMetrics:
    """Helper class to collect and analyze performance metrics."""
    
    def __init__(self, operation_name: str):
        self.operation_name = operation_name
        self.times = []
        self.memory_before = None
        self.memory_after = None
        self.process = psutil.Process()
    
    def start_measurement(self):
        """Start measuring performance."""
        gc.collect()  # Force garbage collection
        self.memory_before = self.process.memory_info().rss
        return time.perf_counter()
    
    def end_measurement(self, start_time: float):
        """End measuring performance."""
        end_time = time.perf_counter()
        duration = end_time - start_time
        self.times.append(duration)
        self.memory_after = self.process.memory_info().rss
        return duration
    
    def get_stats(self):
        """Get performance statistics."""
        if not self.times:
            return None
        
        return {
            'operation': self.operation_name,
            'count': len(self.times),
            'min_time': min(self.times),
            'max_time': max(self.times),
            'mean_time': statistics.mean(self.times),
            'median_time': statistics.median(self.times),
            'std_dev': statistics.stdev(self.times) if len(self.times) > 1 else 0,
            'total_time': sum(self.times),
            'ops_per_second': len(self.times) / sum(self.times) if sum(self.times) > 0 else 0,
            'memory_delta_mb': (self.memory_after - self.memory_before) / (1024 * 1024) if self.memory_before and self.memory_after else 0,
        }


class TestCryptographicPerformance:
    """Performance tests for cryptographic operations."""
    
    def test_keypair_generation_performance(self):
        """Test keypair generation performance."""
        metrics = PerformanceMetrics("keypair_generation")
        num_operations = 100
        
        for _ in range(num_operations):
            start = metrics.start_measurement()
            keypair = generate_keypair()
            metrics.end_measurement(start)
            
            assert len(keypair) == 64
        
        stats = metrics.get_stats()
        print(f"\n{stats['operation']}: {stats['ops_per_second']:.2f} ops/sec, "
              f"avg: {stats['mean_time']*1000:.2f}ms")
        
        # Performance assertions (adjust thresholds as needed)
        assert stats['mean_time'] < 0.01  # Should complete in < 10ms on average
        assert stats['ops_per_second'] > 50  # Should handle at least 50 ops/sec
    
    def test_keypair_from_seed_performance(self):
        """Test deterministic keypair generation performance."""
        metrics = PerformanceMetrics("keypair_from_seed")
        num_operations = 100
        seed = b"a" * 32
        
        for _ in range(num_operations):
            start = metrics.start_measurement()
            keypair = keypair_from_seed(seed)
            metrics.end_measurement(start)
            
            assert len(keypair) == 64
        
        stats = metrics.get_stats()
        print(f"\n{stats['operation']}: {stats['ops_per_second']:.2f} ops/sec, "
              f"avg: {stats['mean_time']*1000:.2f}ms")
        
        # Should be faster than random generation
        assert stats['mean_time'] < 0.005  # Should complete in < 5ms on average
        assert stats['ops_per_second'] > 100  # Should handle at least 100 ops/sec
    
    def test_signing_performance(self):
        """Test message signing performance."""
        keypair = generate_keypair()
        message = b"test message for performance testing"
        
        metrics = PerformanceMetrics("message_signing")
        num_operations = 100
        
        for _ in range(num_operations):
            start = metrics.start_measurement()
            signature = sign(keypair, message)
            metrics.end_measurement(start)
            
            assert len(signature.to_bytes()) == 64
        
        stats = metrics.get_stats()
        print(f"\n{stats['operation']}: {stats['ops_per_second']:.2f} ops/sec, "
              f"avg: {stats['mean_time']*1000:.2f}ms")
        
        # Signing should be reasonably fast
        assert stats['mean_time'] < 0.01  # Should complete in < 10ms on average
        assert stats['ops_per_second'] > 50  # Should handle at least 50 ops/sec
    
    def test_verification_performance(self):
        """Test signature verification performance."""
        keypair = generate_keypair()
        pubkey = public_key_from_keypair(keypair)
        message = b"test message for performance testing"
        signature = sign(keypair, message)
        
        metrics = PerformanceMetrics("signature_verification")
        num_operations = 100
        
        for _ in range(num_operations):
            start = metrics.start_measurement()
            result = verify(pubkey, message, signature)
            metrics.end_measurement(start)
            
            assert result is True
        
        stats = metrics.get_stats()
        print(f"\n{stats['operation']}: {stats['ops_per_second']:.2f} ops/sec, "
              f"avg: {stats['mean_time']*1000:.2f}ms")
        
        # Verification should be reasonably fast
        assert stats['mean_time'] < 0.01  # Should complete in < 10ms on average
        assert stats['ops_per_second'] > 50  # Should handle at least 50 ops/sec
    
    def test_mnemonic_operations_performance(self):
        """Test mnemonic generation and keypair derivation performance."""
        # Test mnemonic generation
        gen_metrics = PerformanceMetrics("mnemonic_generation")
        num_operations = 50  # Fewer operations as these are more expensive
        
        mnemonics = []
        for _ in range(num_operations):
            start = gen_metrics.start_measurement()
            mnemonic = generate_mnemonic()
            gen_metrics.end_measurement(start)
            
            mnemonics.append(mnemonic)
            assert len(mnemonic.split()) == 12
        
        gen_stats = gen_metrics.get_stats()
        print(f"\n{gen_stats['operation']}: {gen_stats['ops_per_second']:.2f} ops/sec, "
              f"avg: {gen_stats['mean_time']*1000:.2f}ms")
        
        # Test keypair from mnemonic
        derive_metrics = PerformanceMetrics("keypair_from_mnemonic")
        test_mnemonic = mnemonics[0]
        
        for _ in range(num_operations):
            start = derive_metrics.start_measurement()
            keypair = keypair_from_mnemonic(test_mnemonic)
            derive_metrics.end_measurement(start)
            
            assert len(keypair) == 64
        
        derive_stats = derive_metrics.get_stats()
        print(f"\n{derive_stats['operation']}: {derive_stats['ops_per_second']:.2f} ops/sec, "
              f"avg: {derive_stats['mean_time']*1000:.2f}ms")
        
        # Mnemonic operations are more expensive but should still be reasonable
        assert gen_stats['mean_time'] < 0.1  # Should complete in < 100ms on average
        assert derive_stats['mean_time'] < 0.1  # Should complete in < 100ms on average (includes HD derivation)


class TestWalletPerformance:
    """Performance tests for wallet operations."""
    
    def test_wallet_creation_performance(self):
        """Test wallet creation performance."""
        metrics = PerformanceMetrics("wallet_creation")
        num_operations = 100
        
        wallets = []
        for i in range(num_operations):
            keypair = generate_keypair()
            
            start = metrics.start_measurement()
            wallet = Wallet(f"test_wallet_{i}", keypair)
            metrics.end_measurement(start)
            
            wallets.append(wallet)
            assert wallet.name == f"test_wallet_{i}"
        
        stats = metrics.get_stats()
        print(f"\n{stats['operation']}: {stats['ops_per_second']:.2f} ops/sec, "
              f"avg: {stats['mean_time']*1000:.2f}ms")
        
        # Wallet creation should be very fast
        assert stats['mean_time'] < 0.001  # Should complete in < 1ms on average
        assert stats['ops_per_second'] > 500  # Should handle at least 500 ops/sec
    
    def test_wallet_with_mnemonic_performance(self):
        """Test wallet creation with mnemonic performance."""
        metrics = PerformanceMetrics("wallet_with_mnemonic_creation")
        num_operations = 50  # Fewer due to mnemonic overhead
        
        for i in range(num_operations):
            start = metrics.start_measurement()
            wallet, mnemonic = create_wallet_with_mnemonic(f"mnemonic_wallet_{i}")
            metrics.end_measurement(start)
            
            assert wallet.name == f"mnemonic_wallet_{i}"
            assert len(mnemonic.split()) == 12
        
        stats = metrics.get_stats()
        print(f"\n{stats['operation']}: {stats['ops_per_second']:.2f} ops/sec, "
              f"avg: {stats['mean_time']*1000:.2f}ms")
        
        # Should be reasonable even with mnemonic generation
        assert stats['mean_time'] < 0.2  # Should complete in < 200ms on average
        assert stats['ops_per_second'] > 10  # Should handle at least 10 ops/sec
    


class TestTransactionPerformance:
    """Performance tests for transaction operations."""
    
    def test_instruction_creation_performance(self):
        """Test instruction creation performance."""
        program_id = PublicKey(b"a" * 32)
        account_pubkey = PublicKey(b"b" * 32)
        instruction_data = b"test_instruction_data"
        
        metrics = PerformanceMetrics("instruction_creation")
        num_operations = 1000  # Many operations since this should be very fast
        
        for _ in range(num_operations):
            account_meta = AccountMeta(account_pubkey, True, True)
            
            start = metrics.start_measurement()
            instruction = Instruction(program_id, [account_meta], instruction_data)
            metrics.end_measurement(start)
            
            assert instruction is not None
        
        stats = metrics.get_stats()
        print(f"\n{stats['operation']}: {stats['ops_per_second']:.2f} ops/sec, "
              f"avg: {stats['mean_time']*1000:.2f}ms")
        
        # Instruction creation should be extremely fast
        assert stats['mean_time'] < 0.001  # Should complete in < 1ms on average
        assert stats['ops_per_second'] > 1000  # Should handle at least 1000 ops/sec
    
    def test_transfer_instruction_performance(self):
        """Test transfer instruction creation performance."""
        from_pubkey = PublicKey(b"a" * 32)
        to_pubkey = PublicKey(b"b" * 32)
        amount = 1000000  # 1M kelvin
        
        metrics = PerformanceMetrics("transfer_instruction_creation")
        num_operations = 1000
        
        for _ in range(num_operations):
            start = metrics.start_measurement()
            instruction = transfer_instruction(from_pubkey, to_pubkey, amount)
            metrics.end_measurement(start)
            
            assert instruction is not None
        
        stats = metrics.get_stats()
        print(f"\n{stats['operation']}: {stats['ops_per_second']:.2f} ops/sec, "
              f"avg: {stats['mean_time']*1000:.2f}ms")
        
        # Transfer instruction creation should be very fast
        assert stats['mean_time'] < 0.001  # Should complete in < 1ms on average
        assert stats['ops_per_second'] > 1000  # Should handle at least 1000 ops/sec
    
    def test_transaction_builder_performance(self):
        """Test transaction builder performance."""
        payer = PublicKey(b"a" * 32)
        
        metrics = PerformanceMetrics("transaction_building")
        num_operations = 100
        
        for _ in range(num_operations):
            start = metrics.start_measurement()
            
            builder = TransactionBuilder(payer, 0, 0)
            builder.add_transfer_instruction(payer, PublicKey(b"c" * 32), 1000000)
            transaction_bytes = builder.build()
            
            metrics.end_measurement(start)
            
            assert transaction_bytes is not None
            assert len(transaction_bytes) > 0
        
        stats = metrics.get_stats()
        print(f"\n{stats['operation']}: {stats['ops_per_second']:.2f} ops/sec, "
              f"avg: {stats['mean_time']*1000:.2f}ms")
        
        # Transaction building should be reasonably fast
        assert stats['mean_time'] < 0.01  # Should complete in < 10ms on average
        assert stats['ops_per_second'] > 50  # Should handle at least 50 ops/sec


class TestTypeConversionPerformance:
    """Performance tests for type conversions and utility operations."""
    
    def test_type_creation_performance(self):
        """Test core type creation performance."""
        # Test PublicKey creation
        pubkey_metrics = PerformanceMetrics("pubkey_creation")
        pubkey_bytes = b"a" * 32
        
        for _ in range(1000):
            start = pubkey_metrics.start_measurement()
            pubkey = PublicKey(pubkey_bytes)
            pubkey_metrics.end_measurement(start)
            
            assert pubkey.to_bytes() == pubkey_bytes
        
        pubkey_stats = pubkey_metrics.get_stats()
        print(f"\n{pubkey_stats['operation']}: {pubkey_stats['ops_per_second']:.2f} ops/sec")
        
        # Test Hash creation
        hash_metrics = PerformanceMetrics("hash_creation")
        hash_bytes = b"b" * 32
        
        for _ in range(1000):
            start = hash_metrics.start_measurement()
            hash_obj = Hash(hash_bytes)
            hash_metrics.end_measurement(start)
            
            assert hash_obj.to_bytes() == hash_bytes
        
        hash_stats = hash_metrics.get_stats()
        print(f"\n{hash_stats['operation']}: {hash_stats['ops_per_second']:.2f} ops/sec")
        
        # Test Signature creation
        sig_metrics = PerformanceMetrics("signature_creation")
        sig_bytes = b"c" * 64
        
        for _ in range(1000):
            start = sig_metrics.start_measurement()
            signature = Signature(sig_bytes)
            sig_metrics.end_measurement(start)
            
            assert signature.to_bytes() == sig_bytes
        
        sig_stats = sig_metrics.get_stats()
        print(f"\n{sig_stats['operation']}: {sig_stats['ops_per_second']:.2f} ops/sec")
        
        # All type creations should be very fast
        assert pubkey_stats['ops_per_second'] > 5000
        assert hash_stats['ops_per_second'] > 5000
        assert sig_stats['ops_per_second'] > 5000
    
    def test_unit_conversion_performance(self):
        """Test unit conversion performance."""
        # Test RLO to kelvin conversion
        rlo_metrics = PerformanceMetrics("rlo_to_kelvin_conversion")
        
        for i in range(10000):
            rlo_amount = float(i) / 1000.0  # 0.0 to 10.0 RLO
            start = rlo_metrics.start_measurement()
            kelvin_amount = rlo_to_kelvin(rlo_amount)
            rlo_metrics.end_measurement(start)
            
            assert kelvin_amount >= 0
        
        rlo_stats = rlo_metrics.get_stats()
        print(f"\n{rlo_stats['operation']}: {rlo_stats['ops_per_second']:.2f} ops/sec")
        
        # Test kelvin to RLO conversion
        kelvin_metrics = PerformanceMetrics("kelvin_to_rlo_conversion")
        
        for i in range(10000):
            kelvin_amount = i * 1000000  # Various kelvin amounts
            start = kelvin_metrics.start_measurement()
            rlo_amount = kelvin_to_rlo(kelvin_amount)
            kelvin_metrics.end_measurement(start)
            
            assert rlo_amount >= 0.0
        
        kelvin_stats = kelvin_metrics.get_stats()
        print(f"\n{kelvin_stats['operation']}: {kelvin_stats['ops_per_second']:.2f} ops/sec")
        
        # Unit conversions should be extremely fast (threshold accounts for gc.collect() overhead in measurement)
        assert rlo_stats['ops_per_second'] > 30000
        assert kelvin_stats['ops_per_second'] > 30000


    
    def test_threaded_cryptographic_operations(self):
        """Test performance of cryptographic operations in multiple threads."""
        num_threads = 4
        operations_per_thread = 25
        
        def crypto_operations():
            """Perform cryptographic operations in a thread."""
            times = []
            for _ in range(operations_per_thread):
                start = time.perf_counter()
                
                # Generate keypair
                keypair = generate_keypair()
                
                # Sign message
                message = b"threaded test message"
                signature = sign(keypair, message)
                
                # Verify signature
                pubkey = public_key_from_keypair(keypair)
                verified = verify(pubkey, message, signature)
                
                end = time.perf_counter()
                times.append(end - start)
                
                assert verified is True
            
            return times
        
        # Run operations in multiple threads
        start_total = time.perf_counter()
        with ThreadPoolExecutor(max_workers=num_threads) as executor:
            futures = [executor.submit(crypto_operations) for _ in range(num_threads)]
            results = [future.result() for future in futures]
        total_time = time.perf_counter() - start_total
        
        # Analyze results
        all_times = [time for thread_times in results for time in thread_times]
        total_operations = num_threads * operations_per_thread
        
        print(f"\nThreaded Cryptographic Operations:")
        print(f"  Threads: {num_threads}")
        print(f"  Operations per thread: {operations_per_thread}")
        print(f"  Total operations: {total_operations}")
        print(f"  Total time: {total_time:.3f}s")
        print(f"  Avg operation time: {statistics.mean(all_times)*1000:.2f}ms")
        print(f"  Operations per second: {total_operations / total_time:.2f}")
        
        # Threaded operations should scale reasonably
        assert total_time < 30.0  # Should complete within 30 seconds
        assert statistics.mean(all_times) < 0.1  # Avg operation < 100ms


class TestMemoryUsagePerformance:
    """Performance tests focused on memory usage patterns."""
    
    def test_memory_usage_during_operations(self):
        """Test memory usage patterns during various operations with larger payloads."""
        process = psutil.Process()

        # Baseline memory
        gc.collect()
        baseline_memory = process.memory_info().rss

        # Perform memory-intensive operations with larger payloads
        wallets = []
        keypairs = []
        signatures = []
        large_messages = []

        print(f"\nMemory Usage Test:")
        print(f"  Baseline memory: {baseline_memory / (1024*1024):.2f} MB")

        # Create many objects with larger memory footprints
        for i in range(500):  # Increased from 100 to 500
            # Generate keypair
            keypair = generate_keypair()
            keypairs.append(keypair)

            # Create wallet with large name
            wallet_name = f"memory_test_wallet_{i}_" + "x" * 2000  # 2KB name
            wallet = Wallet(wallet_name, keypair)
            wallets.append(wallet)

            # Sign large message
            large_message = f"memory test message {i} " + "y" * (20 * 1024)  # 20KB message
            large_message_bytes = large_message.encode()
            large_messages.append(large_message_bytes)

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

        peak_memory = process.memory_info().rss
        print(f"  Peak memory: {peak_memory / (1024*1024):.2f} MB")
        print(f"  Memory increase: {(peak_memory - baseline_memory) / (1024*1024):.2f} MB")

        # Clear references and force garbage collection
        del wallets, keypairs, signatures, large_messages
        gc.collect()

        final_memory = process.memory_info().rss
        print(f"  Final memory: {final_memory / (1024*1024):.2f} MB")
        print(f"  Memory recovered: {(peak_memory - final_memory) / (1024*1024):.2f} MB")

        # Memory should be reasonable (increased limit due to larger payloads)
        memory_increase_mb = (peak_memory - baseline_memory) / (1024*1024)
        assert memory_increase_mb < 100  # Should not use more than 100MB for 500 large objects

        # Test memory behavior with realistic expectations for Rust/PyO3 backend
        if peak_memory > baseline_memory:
            recovery_ratio = (peak_memory - final_memory) / (peak_memory - baseline_memory)
            memory_recovered_mb = (peak_memory - final_memory) / (1024*1024)

            print(f"  Memory recovery ratio: {recovery_ratio:.1%}")

            if recovery_ratio > 0.1:
                print(f"  ✅ Memory recovery detected: {memory_recovered_mb:.2f} MB freed")
            else:
                print(f"  ℹ️  Memory retained by efficient Rust backend (normal for PyO3)")
                print(f"  ℹ️  This indicates memory pooling/caching for performance optimization")

                # Test that memory usage was reasonable and bounded
                assert memory_increase_mb > 5, f"Should use measurable memory for large operations, got {memory_increase_mb:.2f} MB"
                assert memory_increase_mb < 100, f"Memory usage should be bounded, got {memory_increase_mb:.2f} MB"

                print(f"  ✅ Memory usage is reasonable and bounded: {memory_increase_mb:.2f} MB")
        else:
            print("  No measurable memory increase detected - highly efficient implementation")


if __name__ == "__main__":
    pytest.main([__file__, "-v", "-s"])