"""
Test SHA-256 mining operations on virtual CPU with detailed performance logging
"""
from virtual_cpu import (
    CPU, WORD_SIZE, OP_ADD, OP_XOR, OP_SHR, OP_SHL, OP_AND, OP_OR,
    OP_LOAD, OP_STORE, Memory, MEMORY_SIZE, THREADS_PER_CORE
)
from logic_gates import VDD, VSS, VTH
import time
import hashlib
import logging
import threading
from datetime import datetime

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('mining_performance.log'),
        logging.StreamHandler()
    ]
)

# SHA-256 Constants (first 32 bits of the fractional parts of the cube roots of the first 64 primes)
K = [
    0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5,
    0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
    # ... (truncated for simplicity)
]

def create_sha256_mining_program(data, target_difficulty):
    """Create a SHA-256 mining program for the virtual CPU"""
    program = []
    
    # Initialize hash values (first 32 bits of the fractional parts of the square roots of the first 8 primes)
    h0 = 0x6a09e667
    h1 = 0xbb67ae85
    h2 = 0x3c6ef372
    h3 = 0xa54ff53a
    h4 = 0x510e527f
    h5 = 0x9b05688c
    h6 = 0x1f83d9ab
    h7 = 0x5be0cd19
    
    # Store initial hash values in registers 0-7
    program.extend([
        (OP_ADD << 26) | (0 << 21) | (0 << 16) | h0,
        (OP_ADD << 26) | (1 << 21) | (0 << 16) | h1,
        (OP_ADD << 26) | (2 << 21) | (0 << 16) | h2,
        (OP_ADD << 26) | (3 << 21) | (0 << 16) | h3,
        (OP_ADD << 26) | (4 << 21) | (0 << 16) | h4,
        (OP_ADD << 26) | (5 << 21) | (0 << 16) | h5,
        (OP_ADD << 26) | (6 << 21) | (0 << 16) | h6,
        (OP_ADD << 26) | (7 << 21) | (0 << 16) | h7,
    ])
    
    # Compression function main loop
    for i in range(64):
        # Load message word into register 8
        program.extend([
            (OP_LOAD << 26) | (8 << 21) | (0 << 16) | i,
        ])
        
        # Ch(e,f,g) = (e & f) ^ (~e & g)
        program.extend([
            (OP_AND << 26) | (9 << 21) | (4 << 16) | (5 << 11),   # r9 = e & f
            (OP_XOR << 26) | (4 << 21) | (4 << 16) | (0xffffffff),# r4 = ~e
            (OP_AND << 26) | (10 << 21) | (4 << 16) | (6 << 11),  # r10 = ~e & g
            (OP_XOR << 26) | (9 << 21) | (9 << 16) | (10 << 11),  # ch = r9 ^ r10
        ])
        
        # Maj(a,b,c) = (a & b) ^ (a & c) ^ (b & c)
        program.extend([
            (OP_AND << 26) | (10 << 21) | (0 << 16) | (1 << 11),  # r10 = a & b
            (OP_AND << 26) | (11 << 21) | (0 << 16) | (2 << 11),  # r11 = a & c
            (OP_AND << 26) | (12 << 21) | (1 << 16) | (2 << 11),  # r12 = b & c
            (OP_XOR << 26) | (10 << 21) | (10 << 16) | (11 << 11),# r10 = (a&b) ^ (a&c)
            (OP_XOR << 26) | (10 << 21) | (10 << 16) | (12 << 11),# maj = r10 ^ (b&c)
        ])
        
        # Σ0(a) = ROTR(a,2) ^ ROTR(a,13) ^ ROTR(a,22)
        program.extend([
            (OP_SHR << 26) | (11 << 21) | (0 << 16) | 2,         # r11 = a>>2
            (OP_SHL << 26) | (12 << 21) | (0 << 16) | 30,        # r12 = a<<30
            (OP_OR << 26) | (11 << 21) | (11 << 16) | (12 << 11),# r11 = ROTR(a,2)
            
            (OP_SHR << 26) | (12 << 21) | (0 << 16) | 13,        # r12 = a>>13
            (OP_SHL << 26) | (13 << 21) | (0 << 16) | 19,        # r13 = a<<19
            (OP_OR << 26) | (12 << 21) | (12 << 16) | (13 << 11),# r12 = ROTR(a,13)
            
            (OP_SHR << 26) | (13 << 21) | (0 << 16) | 22,        # r13 = a>>22
            (OP_SHL << 26) | (14 << 21) | (0 << 16) | 10,        # r14 = a<<10
            (OP_OR << 26) | (13 << 21) | (13 << 16) | (14 << 11),# r13 = ROTR(a,22)
            
            (OP_XOR << 26) | (11 << 21) | (11 << 16) | (12 << 11),# r11 = ROTR(2) ^ ROTR(13)
            (OP_XOR << 26) | (11 << 21) | (11 << 16) | (13 << 11),# sigma0 = r11 ^ ROTR(22)
        ])
        
        # Update working variables
        program.extend([
            # t1 = h + Sigma1(e) + Ch(e,f,g) + k[i] + w[i]
            (OP_ADD << 26) | (15 << 21) | (7 << 16) | 0,          # r15 = h
            (OP_ADD << 26) | (15 << 21) | (15 << 16) | (9 << 11), # r15 += ch
            (OP_ADD << 26) | (15 << 21) | (15 << 16) | (8 << 11), # r15 += w[i]
            
            # t2 = Sigma0(a) + Maj(a,b,c)
            (OP_ADD << 26) | (14 << 21) | (11 << 16) | 0,         # r14 = sigma0
            (OP_ADD << 26) | (14 << 21) | (14 << 16) | (10 << 11),# r14 += maj
            
            # Update h=g, g=f, f=e, e=d+t1, d=c, c=b, b=a, a=t1+t2
            (OP_ADD << 26) | (7 << 21) | (6 << 16) | 0,           # h = g
            (OP_ADD << 26) | (6 << 21) | (5 << 16) | 0,           # g = f
            (OP_ADD << 26) | (5 << 21) | (4 << 16) | 0,           # f = e
            (OP_ADD << 26) | (4 << 21) | (3 << 16) | (15 << 11),  # e = d + t1
            (OP_ADD << 26) | (3 << 21) | (2 << 16) | 0,           # d = c
            (OP_ADD << 26) | (2 << 21) | (1 << 16) | 0,           # c = b
            (OP_ADD << 26) | (1 << 21) | (0 << 16) | 0,           # b = a
            (OP_ADD << 26) | (0 << 21) | (15 << 16) | (14 << 11), # a = t1 + t2
        ])
    
    return program

def test_mining_performance():
    """Test SHA-256 mining performance on virtual CPU"""
    logging.info("Initializing Virtual CPU for SHA-256 mining test...")
    cpu = CPU(num_cores=4, clock_freq=2000)  # 4 cores at 2KHz
    
    # Test data setup
    test_data = b"Hello, Mining World!"
    target_difficulty = 0x00000fffffffffffffffffffffffffffffffffffffffffffffffffffffffffff
    
    # Create mining program
    mining_program = create_sha256_mining_program(test_data, target_difficulty)
    start_time = time.time()
    nonce = 0
    hashes_completed = 0
    
    logging.info(f"Starting mining test with data: {test_data}")
    logging.info(f"Target difficulty: {target_difficulty:064x}")
    
    try:
        while nonce < 1000:  # Test with 1000 nonces
            # Update nonce in memory
            cpu.memory.write(0, nonce)
            
            # Execute mining program
            for instruction in mining_program:
                clk = cpu.clock.tick()
                if clk > VTH:
                    pc = cpu.control_unit.program_counter.read()
                    cpu.memory.write(pc, instruction)
                    cpu.execute_instruction(clk)
            
            # Read final hash from registers
            hash_result = 0
            for i in range(8):
                hash_part = cpu.registers[i].read()
                hash_result = (hash_result << 32) | hash_part
            
            # Log every 100 hashes
            if hashes_completed % 100 == 0:
                elapsed = time.time() - start_time
                hash_rate = hashes_completed / elapsed if elapsed > 0 else 0
                logging.info(f"Completed {hashes_completed} hashes at {hash_rate:.2f} H/s")
                logging.info(f"Latest hash: {hash_result:064x}")
                logging.info(f"CPU Stats - Clock: {cpu.clock.current_frequency} Hz, "
                           f"Active Cores: {len([c for c in cpu.cores if any(ctx['active'] for ctx in c.thread_contexts)])}")
            
            nonce += 1
            hashes_completed += 1
            
            # Check if hash meets target
            if hash_result < target_difficulty:
                logging.info(f"Found solution! Nonce: {nonce-1}")
                logging.info(f"Hash: {hash_result:064x}")
                break
                
    except KeyboardInterrupt:
        logging.info("\nMining test interrupted by user")
    finally:
        elapsed = time.time() - start_time
        hash_rate = hashes_completed / elapsed if elapsed > 0 else 0
        
        logging.info("\nMining Test Results:")
        logging.info(f"Total hashes completed: {hashes_completed}")
        logging.info(f"Total time: {elapsed:.2f} seconds")
        logging.info(f"Average hash rate: {hash_rate:.2f} H/s")
        logging.info(f"Instructions executed: {cpu.performance_counters['instructions_executed']}")
        logging.info(f"Cache hit rate: {cpu.get_cache_hit_rate():.2f}%")
        logging.info(f"Active cores: {len([c for c in cpu.cores if any(ctx['active'] for ctx in c.thread_contexts)])} of {len(cpu.cores)}")
        logging.info(f"Threads per core: {THREADS_PER_CORE}")

if __name__ == "__main__":
    test_mining_performance()