from .sbox import obsidiansbox
from .permutation import obsidianpermutation
from .mixer import obsidianmixer
from .keygen import keygenerator
from .utils import pad_data, unpad_data, xor_bytes, split_blocks, join_blocks


class obsidian256cipher:
    block_size = 32
    key_size = 32
    
    def __init__(self, key=None, iv=None):
        self._keygen = keygenerator(key_size=self.key_size)
        self._master_key = key if key else self._keygen.generate_master_key()
        self._iv = iv if iv else self._keygen.generate_iv()
        self._round_keys = None
        self._sboxes = []
        self._permutations = []
        self._mixer = obsidianmixer(self.block_size)
        self._initialized = False
        self._decoy_operations = True
    
    def _initialize(self):
        if self._initialized:
            return
        
        self._round_keys = self._keygen.generate_round_keys(self._master_key)
        num_rounds = len(self._round_keys)
        
        self._sboxes = []
        self._permutations = []
        
        for r in range(num_rounds):
            seed = int.from_bytes(self._round_keys[r][:8], "little")
            sbox = obsidiansbox(seed)
            self._sboxes.append(sbox)
            
            perm = obsidianpermutation(self.block_size)
            perm.generate_permutation(seed ^ 0xdeadbeef)
            self._permutations.append(perm)
        
        self._initialized = True
    
    def _decoy_round(self, data, round_num):
        if not self._decoy_operations:
            return data
        
        temp = bytearray(data)
        for i in range(len(temp)):
            temp[i] ^= (round_num * 0x37) & 0xff
        
        for i in range(len(temp)):
            temp[i] ^= (round_num * 0x37) & 0xff
        
        return bytes(temp)
    
    def _encrypt_block(self, block):
        self._initialize()
        
        if len(block) != self.block_size:
            raise ValueError(f"block must be {self.block_size} bytes")
        
        state = block
        num_rounds = len(self._round_keys)
        
        state = xor_bytes(state, self._master_key)
        
        for r in range(num_rounds):
            state = self._decoy_round(state, r)
            
            state = self._sboxes[r].substitute(state)
            
            state = self._permutations[r].permute(state)
            
            state = self._mixer.mix_forward(state, self._round_keys[r])
            
            state = xor_bytes(state, self._round_keys[r])
            
            if r % 2 == 0:
                state = self._decoy_round(state, r + 100)
        
        state = xor_bytes(state, self._master_key)
        
        return state
    
    def _decrypt_block(self, block):
        self._initialize()
        
        if len(block) != self.block_size:
            raise ValueError(f"block must be {self.block_size} bytes")
        
        state = block
        num_rounds = len(self._round_keys)
        
        state = xor_bytes(state, self._master_key)
        
        for r in range(num_rounds - 1, -1, -1):
            if r % 2 == 0:
                state = self._decoy_round(state, r + 100)
            
            state = xor_bytes(state, self._round_keys[r])
            
            state = self._mixer.mix_inverse(state, self._round_keys[r])
            
            state = self._permutations[r].inverse_permute(state)
            
            state = self._sboxes[r].inverse_substitute(state)
            
            state = self._decoy_round(state, r)
        
        state = xor_bytes(state, self._master_key)
        
        return state
    
    def encrypt(self, plaintext):
        if isinstance(plaintext, str):
            plaintext = plaintext.encode("utf-8")
        
        padded = pad_data(plaintext, self.block_size)
        blocks = split_blocks(padded, self.block_size)
        
        encrypted_blocks = []
        prev_block = self._iv
        
        for block in blocks:
            xored = xor_bytes(block, prev_block)
            encrypted = self._encrypt_block(xored)
            encrypted_blocks.append(encrypted)
            prev_block = encrypted
        
        return join_blocks(encrypted_blocks)
    
    def decrypt(self, ciphertext):
        if len(ciphertext) % self.block_size != 0:
            raise ValueError("invalid ciphertext length")
        
        blocks = split_blocks(ciphertext, self.block_size)
        
        decrypted_blocks = []
        prev_block = self._iv
        
        for block in blocks:
            decrypted = self._decrypt_block(block)
            xored = xor_bytes(decrypted, prev_block)
            decrypted_blocks.append(xored)
            prev_block = block
        
        padded = join_blocks(decrypted_blocks)
        return unpad_data(padded)
    
    def encrypt_with_iv(self, plaintext):
        ciphertext = self.encrypt(plaintext)
        return self._iv + ciphertext
    
    def decrypt_with_iv(self, data):
        if len(data) < self.block_size:
            raise ValueError("data too short")
        iv = data[:self.block_size]
        ciphertext = data[self.block_size:]
        self._iv = iv
        return self.decrypt(ciphertext)
    
    def get_key(self):
        return self._master_key
    
    def get_iv(self):
        return self._iv
    
    def get_round_count(self):
        self._initialize()
        return len(self._round_keys)
    
    def set_decoy_operations(self, enabled):
        self._decoy_operations = enabled
    
    @classmethod
    def generate_key(cls):
        kg = keygenerator(key_size=cls.key_size)
        return kg.generate_master_key()
    
    @classmethod
    def generate_iv(cls):
        kg = keygenerator(key_size=cls.key_size)
        return kg.generate_iv()
