#!/usr/bin/env python3
# Python implementation of a Linear Congruential Generator for IP Sharding - Developed by acidvegas in Python (https://git.acid.vegas/pylcg)
# pylcg.py

import argparse
import asyncio
import ipaddress
from math import ceil


class LCG:
    '''Linear Congruential Generator for deterministic random number generation'''

    def __init__(self, seed: int, m: int = 2**32):
        self.m       = m
        self.a       = 1597
        self.c       = 51749
        self.seed    = seed
        self.current = seed


    def get_nth(self, n: int) -> int:
        '''
        Get the nth number in the sequence without generating previous numbers.

        :param n: The index of the number to get
        '''

        # For large n, use the standard next() method to avoid modular arithmetic issues
        if n > 1000:
            self.current = self.seed
            for _ in range(n):
                self.next()
            return self.current
        
        # For smaller n, use direct calculation
        result = self.seed
        for _ in range(n):
            result = (self.a * result + self.c) % self.m
        return result


    def next(self) -> int:
        '''Generate next random number'''

        self.current = (self.a * self.current + self.c) % self.m

        return self.current



class IPRange:
    '''Memory-efficient IP range iterator'''

    def __init__(self, cidr: str):
        network    = ipaddress.ip_network(cidr)
        self.start = int(network.network_address)
        self.end   = int(network.broadcast_address)
        self.total = self.end - self.start + 1
    
    def get_ip_at_index(self, index: int) -> str:
        '''
        Get IP at specific index without generating previous IPs
        
        :param index: The index of the IP to get
        '''

        if not 0 <= index < self.total:
            raise IndexError('IP index out of range')
        
        return str(ipaddress.ip_address(self.start + index))


async def get_shard_ips(cidr: str, shard_num: int, total_shards: int, seed: int, chunk_size: int = 1000):
    '''
    Asynchronously generate IPs for the specified shard.
    
    :param cidr: The CIDR range to shard
    :param shard_num: The number of the shard to generate
    :param total_shards: The total number of shards
    :param seed: The seed for the random number generator
    :param chunk_size: The size of the chunks to process
    '''

    # Initialize the IP range and LCG
    ip_range  = IPRange(cidr)
    lcg       = LCG(seed)
    total_ips = ip_range.total
    
    # Calculate which indices belong to this shard
    shard_size = ceil(total_ips / total_shards)
    start_idx  = shard_num * shard_size
    end_idx    = min(start_idx + shard_size, total_ips)
    
    # Process in chunks to maintain memory efficiency
    for chunk_start in range(start_idx, end_idx, chunk_size):
        chunk_end     = min(chunk_start + chunk_size, end_idx)
        chunk_indices = list(range(chunk_start, chunk_end))
        
        # Generate random values for this chunk
        chunk_random_values = [(i, lcg.get_nth(i)) for i in chunk_indices]
        chunk_random_values.sort(key=lambda x: x[1])
        
        # Yield IPs in randomized order
        for idx, _ in chunk_random_values:
            yield ip_range.get_ip_at_index(idx)
            
            # Allow other tasks to run (do we need this?)
            await asyncio.sleep(0)


async def main():
    parser = argparse.ArgumentParser(description='Async IP address sharding tool')
    parser.add_argument('cidr', help='Target IP range in CIDR format')
    parser.add_argument('shard_num', type=int, help='Shard number (0-based)')
    parser.add_argument('total_shards', type=int, help='Total number of shards')
    parser.add_argument('--seed', type=int, default=12345, help='Random seed for LCG')
    parser.add_argument('--chunk-size', type=int, default=1000, help='Processing chunk size')
    
    args = parser.parse_args()
    
    if args.shard_num >= args.total_shards:
        raise ValueError('Shard number must be less than total shards')
    
    if args.shard_num < 0 or args.total_shards < 1:
        raise ValueError('Invalid shard configuration')
    
    async for ip in get_shard_ips(args.cidr, args.shard_num, args.total_shards, args.seed, args.chunk_size):
        print(ip)



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