import base64
import gzip
import hashlib
import logging
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

logger = logging.getLogger(__name__)


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, verbose: bool = False) -> 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)
    if verbose:
        logger.info(f"Unzipped: {gz_path} -> {orig_path}")


def zip_files_together(dirpath: str, output_path: str, verbose: bool = False) -> None:
    file_paths = find_files(dirpath, "*")
    file_paths = [f for f in file_paths if not f.name.endswith(".zip")]
    if verbose:
        logger.info(f"Zipping {len(file_paths)} files from {dirpath} to {output_path}")
        for file_path in file_paths:
            logger.debug(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))
    if verbose:
        logger.info(f"Zipped: {len(file_paths)} files -> {output_path}")


def unzip_files_together(zip_path: str, output_path: str, verbose: bool = False) -> None:
    with zipfile.ZipFile(zip_path, "r") as zipf:
        zipf.extractall(output_path)
    if verbose:
        logger.info(f"Unzipped: {zip_path} -> {output_path}")


def password_protect_file(file_path: str, password: str, verbose: bool = False) -> 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)

    if verbose:
        logger.info(f"Password protected: {file_path} -> {file_path + '.enc'}")


def decrypt_file(file_path: str, password: str, verbose: bool = False) -> 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)

        if verbose:
            logger.info(f"Decrypted: {file_path} -> {file_path + '.dec'}")

    except Exception as e:
        if verbose:
            logger.error(f"Decryption failed: {e}")
            logger.warning("This usually means the password is incorrect.")
        # Re-raise the exception so the CLI can handle it properly
        raise ValueError(f"Decryption failed: {e}. 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,
    verbose: bool = False,
) -> 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, verbose)
    if password is not None:
        # Password protect the zip
        password_protect_file(output_path, password, verbose)
        return f"{output_path}.enc"
    else:
        return output_path


def unzip_and_decrypt(
    zip_path: str,
    password: str | None = None,
    output_dir: str | None = None,
    verbose: bool = False,
) -> 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, verbose)
        # Unzip the decrypted file
        decrypted_path = f"{zip_path}.dec"
        unzip_files_together(decrypted_path, output_dir, verbose)
    else:
        # No password, unzip directly
        unzip_files_together(zip_path, output_dir, verbose)

    return output_dir
