import base64
import gzip
import json
import os
from typing import Any

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes, padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

from . import bfile, bfunc
from .btype import XPath

FLAG = 'bcrypto---'


def encrypt(data: bytes, password: str) -> bytes:
    compress_data = gzip.compress(data)
    if len(compress_data) < len(data):
        data = compress_data
    assert len(data) < 200 * 1024, '内容太大无法加密'
    salt = os.urandom(16)
    iv = os.urandom(16)
    key = _get_encrypt_key(password, salt)
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    encryptor = cipher.encryptor()
    padder = padding.PKCS7(128).padder()
    padded_data = padder.update(data) + padder.finalize()
    encrypted_message = encryptor.update(padded_data) + encryptor.finalize()
    result = salt + iv + encrypted_message
    result = bfunc.obfuscate(result)
    return result


def encrypt_text(data: str, password: str) -> str:
    content = encrypt(data.encode(), password)
    return FLAG + base64.b64encode(content).decode('utf-8')


def encrypt_json(data: dict[str, Any], password: str) -> str:
    content = bfunc.json_dumps_mini(data)
    return encrypt_text(content, password)


async def encrypt_file(file: XPath, password: str) -> None:
    data = await bfile.read_bytes(file)
    data = FLAG.encode() + encrypt(data, password)
    await bfile.write_bytes(file, data)


def decrypt(data: bytes, password: str) -> bytes:
    data = bfunc.deobfuscate(data)
    salt = data[:16]
    iv = data[16:32]
    data = data[32:]
    key = _get_encrypt_key(password, salt)
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    decryptor = cipher.decryptor()
    decrypted_message = decryptor.update(data) + decryptor.finalize()
    unpadder = padding.PKCS7(128).unpadder()
    result = unpadder.update(decrypted_message) + unpadder.finalize()
    try:
        result = gzip.decompress(result)
    except Exception:
        pass
    return result


def decrypt_text(data: str, password: str) -> str:
    data = data.strip().replace('\n', ' ').split(' ')[::-1][0]
    if data.startswith(FLAG):
        data = data[len(FLAG) :]
    content = decrypt(base64.b64decode(data), password)
    return content.decode()


def decrypt_json(data: str, password: str) -> dict[str, Any]:
    content = decrypt_text(data, password)
    return json.loads(content)


async def decrypt_file(file: XPath, password: str) -> None:
    data = await bfile.read_bytes(file)
    flag = FLAG.encode()
    if data.startswith(flag):
        data = data[len(flag) :]
    data = decrypt(data, password)
    await bfile.write_bytes(file, data)


# ------------------------------------------------------------------------------------


def _get_encrypt_key(password: str, salt: bytes) -> bytes:
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
        backend=default_backend(),
    )
    return kdf.derive(password.encode())
