from __future__ import annotations

import binascii
from datetime import datetime
import secrets
import stat
import struct
import zipfile
from pathlib import Path
from typing import Iterable
import zlib

try:
    import pyzipper
except ImportError:  # pragma: no cover
    pyzipper = None

ArchiveMember = tuple[Path, str]
ZIP_METHOD_AES = "aes"
ZIP_METHOD_STANDARD = "zip_standard"


def _normalized_members(files: Iterable[Path | ArchiveMember]) -> list[ArchiveMember]:
    members: list[ArchiveMember] = []
    used_names: set[str] = set()
    for item in files:
        path, requested_name = item if isinstance(item, tuple) else (item, item.name)
        requested = Path(requested_name).name or path.name
        stem = Path(requested).stem
        suffix = Path(requested).suffix
        candidate = requested
        counter = 2
        while candidate.casefold() in used_names:
            candidate = f"{stem} ({counter}){suffix}"
            counter += 1
        used_names.add(candidate.casefold())
        members.append((path, candidate))
    return members


def create_zip_archive(
    output_path: Path,
    files: Iterable[Path | ArchiveMember],
    password: str = "",  # nosec B107 - empty means an unencrypted archive.
    method: str = ZIP_METHOD_AES,
) -> Path:
    """Create a ZIP archive, optionally using AES or legacy ZipCrypto encryption."""

    output_path.parent.mkdir(parents=True, exist_ok=True)
    members = _normalized_members(files)
    if password:
        if method == ZIP_METHOD_STANDARD:
            _create_zipcrypto_archive(output_path, members, password)
            return output_path
        _create_aes_archive(output_path, members, password)
        return output_path

    with zipfile.ZipFile(output_path, "w", compression=zipfile.ZIP_DEFLATED) as zip_file:
        for file_path, archive_name in members:
            zip_file.write(file_path, arcname=archive_name)
    return output_path


def create_encrypted_zip(output_path: Path, files: list[Path], password: str, method: str = ZIP_METHOD_AES) -> Path:
    """Backward-compatible wrapper for the original per-rule ZIP helper."""

    return create_zip_archive(output_path, files, password, method)


def _create_aes_archive(output_path: Path, members: list[ArchiveMember], password: str) -> None:
    if pyzipper is None:
        raise RuntimeError("pyzipper is required for writing AES password-protected ZIP files")
    with pyzipper.AESZipFile(
        output_path,
        "w",
        compression=pyzipper.ZIP_DEFLATED,
        encryption=pyzipper.WZ_AES,
    ) as zip_file:
        zip_file.setpassword(password.encode("utf-8"))
        for file_path, archive_name in members:
            zip_file.write(file_path, arcname=archive_name)


def _create_zipcrypto_archive(output_path: Path, members: list[ArchiveMember], password: str) -> None:
    entries: list[_CentralDirectoryEntry] = []
    password_bytes = password.encode("utf-8")
    with output_path.open("wb") as zip_file:
        for file_path, archive_name in members:
            local_header_offset = zip_file.tell()
            file_data = file_path.read_bytes()
            crc = binascii.crc32(file_data) & 0xFFFFFFFF
            compressed_data = _deflate(file_data)
            encrypted_data = _zipcrypto_encrypt(compressed_data, password_bytes, crc)
            compressed_size = len(encrypted_data)
            uncompressed_size = len(file_data)
            modified_at = datetime.fromtimestamp(file_path.stat().st_mtime)
            dos_time, dos_date = _dos_timestamp(modified_at)
            filename_bytes, flags = _filename_bytes(archive_name, encrypted=True)

            zip_file.write(
                struct.pack(
                    "<IHHHHHIIIHH",
                    0x04034B50,
                    20,
                    flags,
                    zipfile.ZIP_DEFLATED,
                    dos_time,
                    dos_date,
                    crc,
                    compressed_size,
                    uncompressed_size,
                    len(filename_bytes),
                    0,
                )
            )
            zip_file.write(filename_bytes)
            zip_file.write(encrypted_data)
            entries.append(
                _CentralDirectoryEntry(
                    filename_bytes=filename_bytes,
                    flags=flags,
                    dos_time=dos_time,
                    dos_date=dos_date,
                    crc=crc,
                    compressed_size=compressed_size,
                    uncompressed_size=uncompressed_size,
                    external_attributes=_external_attributes(file_path),
                    local_header_offset=local_header_offset,
                )
            )

        central_directory_offset = zip_file.tell()
        for entry in entries:
            zip_file.write(
                struct.pack(
                    "<IHHHHHHIIIHHHHHII",
                    0x02014B50,
                    20,
                    20,
                    entry.flags,
                    zipfile.ZIP_DEFLATED,
                    entry.dos_time,
                    entry.dos_date,
                    entry.crc,
                    entry.compressed_size,
                    entry.uncompressed_size,
                    len(entry.filename_bytes),
                    0,
                    0,
                    0,
                    0,
                    entry.external_attributes,
                    entry.local_header_offset,
                )
            )
            zip_file.write(entry.filename_bytes)

        central_directory_size = zip_file.tell() - central_directory_offset
        zip_file.write(
            struct.pack(
                "<IHHHHIIH",
                0x06054B50,
                0,
                0,
                len(entries),
                len(entries),
                central_directory_size,
                central_directory_offset,
                0,
            )
        )


class _CentralDirectoryEntry:
    def __init__(
        self,
        *,
        filename_bytes: bytes,
        flags: int,
        dos_time: int,
        dos_date: int,
        crc: int,
        compressed_size: int,
        uncompressed_size: int,
        external_attributes: int,
        local_header_offset: int,
    ) -> None:
        self.filename_bytes = filename_bytes
        self.flags = flags
        self.dos_time = dos_time
        self.dos_date = dos_date
        self.crc = crc
        self.compressed_size = compressed_size
        self.uncompressed_size = uncompressed_size
        self.external_attributes = external_attributes
        self.local_header_offset = local_header_offset


def _deflate(data: bytes) -> bytes:
    compressor = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
    return compressor.compress(data) + compressor.flush()


def _zipcrypto_encrypt(data: bytes, password: bytes, crc: int) -> bytes:
    cipher = _ZipCrypto(password)
    header = secrets.token_bytes(11) + bytes([(crc >> 24) & 0xFF])
    return cipher.encrypt(header + data)


class _ZipCrypto:
    def __init__(self, password: bytes) -> None:
        self.key0 = 0x12345678
        self.key1 = 0x23456789
        self.key2 = 0x34567890
        for value in password:
            self._update_keys(value)

    def encrypt(self, data: bytes) -> bytes:
        encrypted = bytearray()
        for value in data:
            encrypted.append(value ^ self._decrypt_byte())
            self._update_keys(value)
        return bytes(encrypted)

    def _decrypt_byte(self) -> int:
        temp = self.key2 | 2
        return ((temp * (temp ^ 1)) >> 8) & 0xFF

    def _update_keys(self, value: int) -> None:
        self.key0 = _zipcrypto_crc32_byte(value, self.key0)
        self.key1 = (self.key1 + (self.key0 & 0xFF)) & 0xFFFFFFFF
        self.key1 = (self.key1 * 134775813 + 1) & 0xFFFFFFFF
        self.key2 = _zipcrypto_crc32_byte((self.key1 >> 24) & 0xFF, self.key2)


def _build_zipcrypto_crc(value: int) -> int:
    for _ in range(8):
        if value & 1:
            value = 0xEDB88320 ^ (value >> 1)
        else:
            value >>= 1
    return value


_ZIPCRYPTO_CRC_TABLE = [_build_zipcrypto_crc(value) for value in range(256)]


def _zipcrypto_crc32_byte(value: int, crc: int) -> int:
    return ((crc >> 8) ^ _ZIPCRYPTO_CRC_TABLE[(crc ^ value) & 0xFF]) & 0xFFFFFFFF


def _filename_bytes(filename: str, *, encrypted: bool) -> tuple[bytes, int]:
    flags = 0x1 if encrypted else 0
    try:
        return filename.encode("ascii"), flags
    except UnicodeEncodeError:
        return filename.encode("utf-8"), flags | 0x800


def _dos_timestamp(value: datetime) -> tuple[int, int]:
    year = min(max(value.year, 1980), 2107)
    month = value.month if value.year == year else 1
    day = value.day if value.year == year else 1
    dos_time = value.hour << 11 | value.minute << 5 | value.second // 2
    dos_date = (year - 1980) << 9 | month << 5 | day
    return dos_time, dos_date


def _external_attributes(path: Path) -> int:
    try:
        permissions = stat.S_IMODE(path.stat().st_mode)
    except OSError:
        return 0
    return permissions << 16