import aiohttp
import asyncio
import numpy as np

def generate_middle_column(steps: int) -> bytes:
    """Generate the middle column of Rule 30 for n steps and return as bytes"""
    print(f"Starting Rule 30 generation for {steps:,} steps...")
    
    # Use boolean array for faster operations
    width = steps + 2  # Only need 1 cell padding on each side
    state = np.zeros(width, dtype=bool)
    state[width//2] = True
    
    # Pre-allocate array for ALL middle column bits
    bits = np.zeros(steps, dtype=bool)
    bits[0] = True  # First bit is always 1
    
    # Create views for faster access
    left = state[:-2]
    center = state[1:-1]
    right = state[2:]
    
    # Progress tracking
    update_every = max(1, steps // 1000)
    
    # Main loop - generate each step
    for i in range(1, steps):
        if i % update_every == 0:
            print(f'\rGenerating... {(i/steps)*100:3.1f}% ({i:,}/{steps:,})', end='', flush=True)
        
        # Store middle bit at position i in our pre-allocated array
        bits[i] = center[steps//2]  # This is where we save each middle column bit
        
        # Rule 30: next = left XOR (center OR right)
        # Update center in-place for next iteration
        np.logical_xor(left, np.logical_or(center, right), out=center)
        
        # Shift views for next iteration
        left[1:] = left[:-1]
        right[:-1] = right[1:]
        left[0] = False
        right[-1] = False
    
    print(f'\rGeneration complete: {steps:,} steps')
    
    # Convert all collected bits to bytes at once
    byte_array = np.packbits(bits)  # This packs our middle column bits into bytes
    return bytes(byte_array)

async def fetch_validator_bits(ip: str = "65.109.145.104", port: int = 8091, num_bytes: int = 7) -> tuple[int, bytes, int]:
    """Fetch step and most recent bytes from validator API"""
    print("Starting API fetch...")
    async with aiohttp.ClientSession() as session:
        # Get step
        step_url = f"http://{ip}:{port}/step"
        async with session.get(step_url, timeout=5) as resp:
            if resp.status != 200:
                raise Exception(f"Failed to get step: {resp.status}")
            step = await resp.json()
            print(f"Got step: {step}")
            if step <= 0:
                raise Exception(f"Invalid step value: {step}")
        
        # Get only the most recent bytes
        bits_url = f"http://{ip}:{port}/bits"
        end_byte = (step + 7) // 8
        start_byte = max(0, end_byte - num_bytes)
        headers = {'Range': f'bytes={start_byte}-{end_byte-1}'}
        print(f"Requesting bytes {start_byte}-{end_byte-1}...")
        
        async with session.get(bits_url, headers=headers, timeout=5) as resp:
            if resp.status != 206:
                raise Exception(f"Failed to get bits: {resp.status}")
            
            bytes_data = await resp.read()
            if not bytes_data:
                raise Exception("No data received")
            
            # Remove trailing zero bytes
            while bytes_data and bytes_data[-1] == 0:
                bytes_data = bytes_data[:-1]
                
            print(f"Got {len(bytes_data)} non-zero bytes")
            return step, bytes_data, start_byte

async def main():
    try:
        # Get most recent bytes from API
        step, api_bytes, start_pos = await fetch_validator_bits(num_bytes=7)
        print("\nAPI bytes details:")
        print(f"Step: {step}")
        print(f"Starting from byte position: {start_pos}")
        
        # Skip the last byte from API for comparison
        api_bytes = api_bytes[:-1]
        
        # Generate matching bytes
        print("\nGenerating Rule 30 bytes...")
        test_bytes = generate_middle_column(step)
        
        # Calculate positions
        total_bytes = (step + 7) // 8
        # Get bytes from the same positions as API bytes
        test_bytes = test_bytes[start_pos:start_pos + len(api_bytes)]
        
        # Show the bytes we're comparing
        print("\nBytes being compared:")
        print("API bytes:")
        for i, b in enumerate(api_bytes):
            abs_pos = start_pos + i
            print(f"Position {abs_pos}: {hex(b)} (binary: {format(b, '08b')})")
            
        print("\nGenerated bytes:")
        for i, b in enumerate(test_bytes):
            abs_pos = start_pos + i
            print(f"Position {abs_pos}: {hex(b)} (binary: {format(b, '08b')})")
        
        # Compare bytes
        if api_bytes == test_bytes:
            print("\n✅ Bytes match")
        else:
            print("\n❌ Bytes differ")
            for i, (a, b) in enumerate(zip(api_bytes, test_bytes)):
                abs_pos = start_pos + i
                if a != b:
                    print(f"\nFirst difference at position {abs_pos}:")
                    print(f"API:       {hex(a)} (binary: {format(a, '08b')})")
                    print(f"Generated: {hex(b)} (binary: {format(b, '08b')})")
                    break
        
    except Exception as e:
        print(f"Error: {str(e)}")

if __name__ == "__main__":
    asyncio.run(main())