import base64
import gzip
import hashlib
import io
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 zip_files_to_memory(dirpath: str, verbose: bool = False) -> bytes:
    """Zip files to memory and return the zip data as bytes."""
    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 memory")
        for file_path in file_paths:
            logger.debug(f"Zipping {file_path}")

    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, "w") as zipf:
        for file_path in file_paths:
            zipf.write(file_path, file_path.relative_to(dirpath))

    zip_data = zip_buffer.getvalue()
    zip_buffer.close()

    if verbose:
        logger.info(f"Zipped: {len(file_paths)} files to memory ({len(zip_data)} bytes)")
    return zip_data


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 encrypt_data_in_memory(data: bytes, password: str, verbose: bool = False) -> bytes:
    """
    Encrypt data in memory with a password using Fernet (AES-128).
    Returns the encrypted data as bytes.
    """
    # 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)

    # Encrypt the data
    encrypted_data = fernet.encrypt(data)

    # Return salt + encrypted data
    return salt + encrypted_data


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"

    if password is not None:
        # Zip to memory and encrypt in one step
        zip_data = zip_files_to_memory(dir_path, verbose)
        encrypted_data = encrypt_data_in_memory(zip_data, password, verbose)

        # Write encrypted data directly to .zip.enc file
        encrypted_path = f"{output_path}.enc"
        with open(encrypted_path, "wb") as f_out:
            f_out.write(encrypted_data)

        if verbose:
            logger.info(f"Zipped and password protected: {dir_path} -> {encrypted_path}")
        return encrypted_path
    else:
        # No password, create normal zip file
        zip_files_together(dir_path, output_path, verbose)

        if verbose:
            logger.info(f"Zipped: {dir_path} -> {output_path}")
        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
