import base64
import hashlib
from dataclasses import dataclass
import bcrypt
import numpy as np
from src.models import EncipheredNKode, KeypadSize
from src.customer_cipher import CustomerCipher


@dataclass
class UserCipher:
    property_key: np.ndarray
    combined_position_key: np.ndarray
    pass_key: np.ndarray
    mask_key: np.ndarray
    max_nkode_len: int

    @classmethod
    def create(cls, keypad_size: KeypadSize, customer_pos_key: np.ndarray, max_nkode_len: int) -> 'UserCipher':
        if len(customer_pos_key) != keypad_size.props_per_key:
            raise ValueError("Invalid position values")
        user_pos_key = np.random.choice(2**16,size=keypad_size.props_per_key, replace=False)
        return UserCipher(
            property_key=np.random.choice(2 ** 16, size=keypad_size.total_props, replace=False),
            pass_key=np.random.choice(2 ** 16, size=max_nkode_len, replace=False),
            mask_key=np.random.choice(2**16, size=max_nkode_len, replace=False),
            combined_position_key=user_pos_key ^ customer_pos_key,
            max_nkode_len=max_nkode_len
        )

    def pad_user_mask(self, user_mask: np.ndarray, pos_vals: np.ndarray) -> np.ndarray:
        # TODO: replace with new method
        if len(user_mask) >= self.max_nkode_len:
            raise ValueError("User encoded_mask is too long")
        padding_size = self.max_nkode_len - len(user_mask)
        padding = np.random.choice(pos_vals, size=padding_size, replace=True).astype(np.uint16)
        return np.concatenate([user_mask, padding])

    @staticmethod
    def encode_base64_str(data: np.ndarray) -> str:
        return base64.b64encode( b"".join([int(num).to_bytes(2, byteorder='big') for num in data])).decode("utf-8")

    @staticmethod
    def decode_base64_str(data: str) -> np.ndarray:
        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 np.array(int_list, dtype=np.uint16)

    def encipher_nkode(
            self,
            passcode_prop_idx: list[int],
            customer_cipher: CustomerCipher
    ) -> EncipheredNKode:
        mask = self.encipher_mask(passcode_prop_idx, customer_cipher)
        code = self.hash_nkode(passcode_prop_idx, customer_cipher)
        return EncipheredNKode(
            code=code,
            mask=mask
        )

    def hash_nkode(
            self,
            passcode_prop_idx: list[int],
            customer_prop: CustomerCipher,
    ) -> str:
        salt = bcrypt.gensalt(rounds=12)
        passcode_bytes = self.prehash_passcode(passcode_prop_idx, customer_prop)
        passcode_digest = base64.b64encode(hashlib.sha256(passcode_bytes).digest())
        hashed_data = bcrypt.hashpw(passcode_digest, salt)
        return hashed_data.decode("utf-8")

    def compare_nkode(
            self,
            passcode_prop_idx: list[int],
            customer_prop: CustomerCipher,
            hashed_passcode: str
    ) -> bool:
        passcode_bytes = self.prehash_passcode(passcode_prop_idx, customer_prop)
        passcode_digest = base64.b64encode(hashlib.sha256(passcode_bytes).digest())
        return bcrypt.checkpw(passcode_digest, hashed_passcode.encode('utf-8'))

    def prehash_passcode(
            self,
            passcode_prop_idx: list[int],
            customer_prop: CustomerCipher
    ) -> bytes:
        passcode_len = len(passcode_prop_idx)
        passcode_cipher = self.pass_key.copy()
        passcode_cipher[:passcode_len] = (
            passcode_cipher[:passcode_len] ^
            self.property_key[passcode_prop_idx] ^
            customer_prop.property_key[passcode_prop_idx]
        )
        return passcode_cipher.astype(np.uint16).tobytes()

    def encipher_mask(
            self,
            passcode_prop_idx: list[int],
            customer_cipher: CustomerCipher
    ) -> str:
        pos_idxs = customer_cipher.get_passcode_position_indices_padded(passcode_prop_idx, len(self.mask_key))
        ordered_pos_key = self.combined_position_key[pos_idxs]
        ordered_customer_pos_key = customer_cipher.position_key[pos_idxs]
        mask = ordered_pos_key ^ ordered_customer_pos_key ^ self.mask_key
        encoded_mask = self.encode_base64_str(mask)
        return encoded_mask

    def decipher_mask(self, encoded_mask: str, customer_pos_key: np.ndarray, passcode_len: int) -> list[int]:
        mask = self.decode_base64_str(encoded_mask)
        # user_pos_key ordered by the user's nkode and padded to be length max_nkode_len
        ordered_user_pos_key = mask  ^ self.mask_key
        user_pos_key = customer_pos_key ^ self.combined_position_key
        passcode_position = []
        for position_val in ordered_user_pos_key[:passcode_len]:
            position_idx = np.where(user_pos_key == position_val)[0][0]
            passcode_position.append(int(customer_pos_key[position_idx]))
        return passcode_position