import os
import hashlib
import time


class keygenerator:
    def __init__(self, key_size=32, min_rounds=12, max_rounds=20):
        self._key_size = key_size
        self._min_rounds = min_rounds
        self._max_rounds = max_rounds
    
    def generate_master_key(self, password=None, salt=None):
        if password is None:
            return os.urandom(self._key_size)
        
        if isinstance(password, str):
            password = password.encode("utf-8")
        
        if salt is None:
            salt = os.urandom(16)
        elif isinstance(salt, str):
            salt = salt.encode("utf-8")
        
        derived = self._kdf(password, salt, iterations=100000)
        return derived[:self._key_size], salt
    
    def _kdf(self, password, salt, iterations):
        result = password + salt
        for i in range(iterations):
            h = hashlib.sha512()
            h.update(result)
            h.update(i.to_bytes(8, "little"))
            h.update(salt)
            result = h.digest()
        
        return result
    
    def generate_round_keys(self, master_key, num_rounds=None):
        if num_rounds is None:
            num_rounds = self._calculate_rounds(master_key)
        
        round_keys = []
        state = master_key
        
        for r in range(num_rounds):
            round_key = self._derive_round_key(state, r)
            round_keys.append(round_key)
            state = self._evolve_state(state, round_key)
        
        return round_keys
    
    def _calculate_rounds(self, key):
        key_entropy = sum(key) % 256
        base_rounds = self._min_rounds
        additional = (key_entropy * (self._max_rounds - self._min_rounds)) // 256
        return base_rounds + additional
    
    def _derive_round_key(self, state, round_num):
        h1 = hashlib.sha256()
        h1.update(state)
        h1.update(round_num.to_bytes(4, "little"))
        h1.update(bytes([0x4f, 0x42, 0x53, 0x49, 0x44, 0x49, 0x41, 0x4e]))
        derived1 = h1.digest()
        
        h2 = hashlib.sha256()
        h2.update(derived1)
        h2.update((round_num ^ 0x5a5a5a5a).to_bytes(4, "little"))
        h2.update(state[::-1])
        derived2 = h2.digest()
        
        result = bytearray(self._key_size)
        for i in range(self._key_size):
            d1 = derived1[i % len(derived1)]
            d2 = derived2[i % len(derived2)]
            mixed = (d1 + d2 + round_num + i) & 0xff
            mixed ^= ((d1 * d2) ^ (round_num << 3)) & 0xff
            result[i] = mixed
        
        return bytes(result)
    
    def _evolve_state(self, state, round_key):
        h = hashlib.sha256()
        h.update(state)
        h.update(round_key)
        return h.digest()
    
    def get_round_count(self, master_key):
        return self._calculate_rounds(master_key)
    
    def generate_iv(self):
        return os.urandom(self._key_size)
    
    def generate_keypair(self, password=None):
        if password:
            master_key, salt = self.generate_master_key(password)
        else:
            master_key = self.generate_master_key()
            salt = None
        
        iv = self.generate_iv()
        round_keys = self.generate_round_keys(master_key)
        
        return {
            "master_key": master_key,
            "salt": salt,
            "iv": iv,
            "round_keys": round_keys,
            "num_rounds": len(round_keys),
        }
