import base64
import hashlib
import bcrypt
from secrets import choice
from pydantic import BaseModel

from src.models import EncipheredNKode
from src.nkode_interface import CustomerInterface
from src.utils import generate_random_nonrepeating_list, xor_lists, int_array_to_bytes


class UserCipherKeys(BaseModel):
    alpha_key: list[int]
    set_key: list[int]
    pass_key: list[int]
    mask_key: list[int]
    salt: bytes
    max_nkode_len: int = 10

    @staticmethod
    def new_user_encipher_keys(numb_of_keys: int, attrs_per_key: int, set_values: list[int]):
        assert len(set_values) == attrs_per_key

        set_key = generate_random_nonrepeating_list(attrs_per_key)
        set_key = xor_lists(set_key, set_values)

        return UserCipherKeys(
            alpha_key=generate_random_nonrepeating_list(attrs_per_key * numb_of_keys),
            pass_key=generate_random_nonrepeating_list(numb_of_keys),
            mask_key=generate_random_nonrepeating_list(numb_of_keys),
            set_key=set_key,
            salt=bcrypt.gensalt(),
        )

    def pad_user_mask(self, user_mask: list[int], set_vals: list[int]) -> list[int]:
        assert (len(user_mask) <= self.max_nkode_len)
        padded_user_mask = user_mask.copy()
        for _ in range(self.max_nkode_len - len(user_mask)):
            padded_user_mask.append(choice(set_vals))
        return padded_user_mask

    @staticmethod
    def pad_user_code(user_code: list[int], max_nkode_len: int) -> list[int]:
        assert (len(user_code) <= max_nkode_len)
        return user_code + [0 for _ in range(max_nkode_len - len(user_code))]

    @staticmethod
    def encode_base64_str(data: list[int]) -> str:
        return base64.b64encode(int_array_to_bytes(data)).decode("utf-8")

    @staticmethod
    def decode_base64_str(data: str) -> list[int]:
        byte_data = base64.b64decode(data)
        int_list = []

        for i in range(0, len(byte_data), 2):
            int_val = int.from_bytes(byte_data[i:i + 2], byteorder='big')
            int_list.append(int_val)
        return int_list

    def _hash_passcode(self, passcode: list[int]) -> str:
        passcode_bytes = int_array_to_bytes(passcode)
        passcode_digest = base64.b64encode(hashlib.sha256(passcode_bytes).digest())
        hashed_data = bcrypt.hashpw(passcode_digest, self.salt)
        return hashed_data.decode("utf-8")

    def encipher_nkode(
            self,
            passcode_attr_idx: list[int],
            customer_interface: CustomerInterface
    ) -> EncipheredNKode:

        passcode_attrs = [customer_interface.customer_interface[idx] for idx in passcode_attr_idx]
        passcode_sets = [customer_interface.get_attr_set_val(attr) for attr in passcode_attrs]
        code = self.encipher_salt_hash_code(passcode_attr_idx, customer_interface)
        mask = self.encipher_mask(passcode_sets, customer_interface)
        return EncipheredNKode(
            code=code,
            mask=mask
        )

    def encipher_salt_hash_code(
            self,
            passcode_attr_idx: list[int],
            customer_interface: CustomerInterface,
    ) -> str:
        passcode_len = len(passcode_attr_idx)
        passcode_attrs = [customer_interface.customer_interface[idx] for idx in passcode_attr_idx]

        passcode_cipher = self.pass_key.copy()

        for idx in range(passcode_len):
            attr_idx = passcode_attr_idx[idx]
            alpha = self.alpha_key[attr_idx]
            attr_val = passcode_attrs[idx]
            passcode_cipher[idx] ^= alpha ^ attr_val
        return self._hash_passcode(passcode_cipher)

    def encipher_mask(
            self,
            passcode_sets: list[int],
            customer_interface: CustomerInterface
    ) -> str:
        padded_passcode_sets = self.pad_user_mask(passcode_sets, customer_interface.set_vals)
        set_idx = [customer_interface.get_set_index(set_val) for set_val in padded_passcode_sets]
        sorted_set_key = [self.set_key[idx] for idx in set_idx]
        ciphered_mask = xor_lists(sorted_set_key, padded_passcode_sets)
        ciphered_mask = xor_lists(ciphered_mask, self.mask_key)
        mask = self.encode_base64_str(ciphered_mask)
        return mask

    def decipher_mask(self, mask: str, set_vals: list, passcode_len: int) -> list[int]:
        decoded_mask = self.decode_base64_str(mask)
        deciphered_mask = xor_lists(decoded_mask, self.mask_key)
        set_key_ciphers = xor_lists(set_vals, self.set_key)
        passcode_sets = []
        for set_cipher in deciphered_mask[:passcode_len]:
            set_idx = set_key_ciphers.index(set_cipher)
            passcode_sets.append(set_vals[set_idx])
        return passcode_sets