from __future__ import annotations

import binascii
from datetime import datetime
import hashlib
from importlib import metadata
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 not in {ZIP_METHOD_AES, ZIP_METHOD_STANDARD}:
            raise ValueError(f"Unsupported password-encryption method: {method}")
        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 zip_archive_evidence(
    output_path: Path,
    members: Iterable[Path | ArchiveMember],
    *,
    password_protected: bool,
    method: str,
) -> dict[str, object]:
    """Return password-free, content-addressed evidence for one built archive."""

    normalized = _normalized_members(members)
    archive_bytes = output_path.read_bytes()
    if password_protected and method == ZIP_METHOD_AES:
        try:
            implementation_version = metadata.version("pyzipper")
        except metadata.PackageNotFoundError:  # pragma: no cover - guarded by writer
            implementation_version = "unknown"
        implementation = "pyzipper"
        archive_format = "WinZip AES"
    elif password_protected and method == ZIP_METHOD_STANDARD:
        implementation = "govoplan-campaign.zipcrypto"
        implementation_version = "1"
        archive_format = "Legacy ZipCrypto"
    else:
        implementation = "python.zipfile"
        implementation_version = "stdlib"
        archive_format = "ZIP (unencrypted)"
    return {
        "format": archive_format,
        "method": method if password_protected else "none",
        "password_protected": password_protected,
        "implementation": implementation,
        "implementation_version": implementation_version,
        "archive_sha256": hashlib.sha256(archive_bytes).hexdigest(),
        "archive_size_bytes": len(archive_bytes),
        "members": [
            {
                "name": archive_name,
                "sha256": hashlib.sha256(path.read_bytes()).hexdigest(),
                "size_bytes": path.stat().st_size,
            }
            for path, archive_name in normalized
        ],
    }


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