import base64
import gzip
import hashlib
import os
import shutil
import zipfile
from pathlib import Path

from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC


def find_files(root_dir: str, pattern: str) -> list[Path]:
    return list(Path(root_dir).rglob(pattern))


def gzip_file(file_path: Path) -> Path:
    gz_path = file_path.with_suffix(file_path.suffix + ".gz")
    with open(file_path, "rb") as f_in, gzip.open(gz_path, "wb") as f_out:
        shutil.copyfileobj(f_in, f_out)
    return gz_path


def gunzip_file(gz_path: Path) -> None:
    orig_path = gz_path.with_suffix("")
    with gzip.open(gz_path, "rb") as f_in, open(orig_path, "wb") as f_out:
        shutil.copyfileobj(f_in, f_out)
    print(f"Unzipped: {gz_path} -> {orig_path}")


def zip_files_together(dirpath: str, output_path: str) -> None:
    file_paths = find_files(dirpath, "*")
    file_paths = [f for f in file_paths if not f.name.endswith(".zip")]
    print(f"Zipping {len(file_paths)} files from {dirpath} to {output_path}")
    for file_path in file_paths:
        print(f"Zipping {file_path}")
    with zipfile.ZipFile(output_path, "w") as zipf:
        for file_path in file_paths:
            zipf.write(file_path, file_path.relative_to(dirpath))
    print(f"Zipped: {file_paths} -> {output_path}")


def unzip_files_together(zip_path: str, output_path: str) -> None:
    with zipfile.ZipFile(zip_path, "r") as zipf:
        zipf.extractall(output_path)
    print(f"Unzipped: {zip_path} -> {output_path}")


def password_protect_file(file_path: str, password: str) -> None:
    """
    Actually encrypt a file with a password using Fernet (AES-128).
    Works on Linux, Mac, and Windows.
    """
    # Generate a key from the password
    salt = os.urandom(16)
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
    )
    key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
    fernet = Fernet(key)

    # Read and encrypt the file
    with open(file_path, "rb") as f_in:
        data = f_in.read()

    encrypted_data = fernet.encrypt(data)

    # Write encrypted data with salt
    with open(file_path + ".enc", "wb") as f_out:
        f_out.write(salt + encrypted_data)

    print(f"Password protected: {file_path} -> {file_path + '.enc'}")


def decrypt_file(file_path: str, password: str) -> None:
    """
    Decrypt a file with a password.
    Works on Linux, Mac, and Windows.
    """
    # Read the encrypted file
    with open(file_path, "rb") as f_in:
        data = f_in.read()

    # Extract salt and encrypted data
    salt = data[:16]
    encrypted_data = data[16:]

    # Generate the key from password and salt
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
    )
    key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
    fernet = Fernet(key)

    # Decrypt the data
    try:
        decrypted_data = fernet.decrypt(encrypted_data)

        # Write decrypted data
        with open(file_path + ".dec", "wb") as f_out:
            f_out.write(decrypted_data)

        print(f"Decrypted: {file_path} -> {file_path + '.dec'}")

    except Exception as e:
        print(f"Decryption failed: {e}")
        print("This usually means the password is incorrect.")


def get_hash_of_zip(file_path: str) -> str:
    with open(file_path, "rb") as f:
        return hashlib.sha256(f.read()).hexdigest()


def zip_and_password_protect(dir_path: str, password: str | None = None, output_path: str | None = None) -> str:
    """
    Zip a directory and password protect it.
    Works on Linux, Mac, and Windows.
    """
    if output_path is None:
        output_path = f"{dir_path}.zip"

    # Zip the directory
    zip_files_together(dir_path, output_path)
    if password is not None:
        # Password protect the zip
        password_protect_file(output_path, password)

    return f"{output_path}.enc"


def unzip_and_decrypt(zip_path: str, password: str | None = None, output_dir: str | None = None) -> str:
    """
    Decrypt and unzip a password-protected zip file.
    Works on Linux, Mac, and Windows.
    """
    if output_dir is None:
        output_dir = f"{zip_path}.extracted"
    if password is not None:
        # Decrypt the zip
        decrypt_file(zip_path, password)

    # Unzip the decrypted file
    decrypted_path = f"{zip_path}.dec"
    unzip_files_together(decrypted_path, output_dir)

    return output_dir
