import os
from typing import Self

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

from ..db.models import Account
from ..decrypted_account import DecryptedAccount


class Cipher:
    def __init__(self, key: bytes) -> None:
        self._chacha = ChaCha20Poly1305(key)

    @classmethod
    def generate_cipher(cls, salt: bytes, password: bytes) -> Self:
        """Generates cipher which uses key derived from a given password"""
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
            backend=default_backend(),
        )
        return cls(kdf.derive(password))

    def encrypt(self, data: bytes) -> bytes:
        nonce = os.urandom(12)
        return nonce + self._chacha.encrypt(
            nonce,
            data,
            associated_data=None,
        )

    def decrypt(self, data: bytes) -> bytes:
        return self._chacha.decrypt(
            nonce=data[:12],
            data=data[12:],
            associated_data=None,
        )


def encrypt(
    account: DecryptedAccount,
    master_pass: str,
) -> Account:
    """Encrypts account using master password and returns Account object"""
    salt = os.urandom(64)
    cipher = Cipher.generate_cipher(salt, master_pass.encode("utf-8"))

    enc_login = cipher.encrypt(account.login.encode("utf-8"))
    enc_password = cipher.encrypt(account.password.encode("utf-8"))

    return Account(
        user_id=account.user_id,
        name=account.name,
        salt=salt,
        enc_login=enc_login,
        enc_password=enc_password,
    )


def decrypt(
    account: Account,
    master_pass: str,
) -> DecryptedAccount:
    """Decrypts account using master password and returns
    DecryptedAccount object"""
    cipher = Cipher.generate_cipher(account.salt, master_pass.encode("utf-8"))

    login = cipher.decrypt(account.enc_login).decode("utf-8")
    password = cipher.decrypt(account.enc_password).decode("utf-8")

    return DecryptedAccount(
        user_id=account.user_id,
        name=account.name,
        login=login,
        password=password,
    )