from __future__ import annotations

import io
import tarfile
import unittest
import zipfile
from types import SimpleNamespace
from unittest.mock import patch

import pyzipper

from govoplan_files.backend.storage.archives import (
    ArchivePasswordError,
    extract_archive_upload,
    inspect_archive,
)
from govoplan_files.backend.storage.common import FileStorageError


def _zip_bytes(entries: dict[str, bytes]) -> bytes:
    target = io.BytesIO()
    with zipfile.ZipFile(target, "w", compression=zipfile.ZIP_DEFLATED) as archive:
        for path, data in entries.items():
            archive.writestr(path, data)
    return target.getvalue()


def _tar_bytes(entries: dict[str, bytes], mode: str = "w:gz") -> bytes:
    target = io.BytesIO()
    with tarfile.open(fileobj=target, mode=mode) as archive:
        for path, data in entries.items():
            info = tarfile.TarInfo(path)
            info.size = len(data)
            archive.addfile(info, io.BytesIO(data))
    return target.getvalue()


def _encrypted_zip_bytes(password: str) -> bytes:
    target = io.BytesIO()
    with pyzipper.AESZipFile(
        target,
        "w",
        compression=zipfile.ZIP_DEFLATED,
        encryption=pyzipper.WZ_AES,
    ) as archive:
        archive.setpassword(password.encode("utf-8"))
        archive.writestr("secure/report.txt", b"classified")
    return target.getvalue()


class ArchiveInspectionTests(unittest.TestCase):
    def test_zip_preview_derives_folders_and_sizes(self) -> None:
        inspection = inspect_archive(
            _zip_bytes(
                {
                    "department/one.txt": b"one",
                    "department/nested/two.csv": b"two",
                }
            ),
            filename="monthly.zip",
        )

        self.assertEqual("zip", inspection.archive_format)
        self.assertEqual(2, inspection.file_count)
        self.assertEqual(2, inspection.directory_count)
        self.assertEqual(6, inspection.expanded_size_bytes)
        self.assertEqual(
            [
                "department",
                "department/nested",
                "department/nested/two.csv",
                "department/one.txt",
            ],
            [entry.path for entry in inspection.entries],
        )

    def test_tar_compression_variants_are_inspected(self) -> None:
        for filename, mode in (
            ("monthly.tar", "w"),
            ("monthly.tar.gz", "w:gz"),
            ("monthly.tar.bz2", "w:bz2"),
            ("monthly.tar.xz", "w:xz"),
        ):
            with self.subTest(filename=filename):
                inspection = inspect_archive(
                    _tar_bytes({"report.txt": b"report"}, mode=mode),
                    filename=filename,
                )
                self.assertEqual(1, inspection.file_count)
                self.assertEqual(filename.removeprefix("monthly."), inspection.archive_format)

    def test_unsafe_member_path_is_rejected(self) -> None:
        with self.assertRaisesRegex(FileStorageError, "Unsafe archive member"):
            inspect_archive(
                _zip_bytes({"../escape.txt": b"no"}),
                filename="unsafe.zip",
            )

    def test_archive_bomb_ratio_is_rejected(self) -> None:
        with self.assertRaisesRegex(FileStorageError, "expansion ratio"):
            inspect_archive(
                _zip_bytes({"zeros.bin": b"\0" * 50_000}),
                filename="bomb.zip",
                max_expansion_ratio=2,
            )

    def test_encrypted_zip_reports_and_verifies_password(self) -> None:
        archive = _encrypted_zip_bytes("correct horse")

        missing = inspect_archive(archive, filename="secure.zip")
        self.assertTrue(missing.requires_password)
        self.assertFalse(missing.password_verified)

        verified = inspect_archive(
            archive,
            filename="secure.zip",
            password="correct horse",
        )
        self.assertTrue(verified.requires_password)
        self.assertTrue(verified.password_verified)

        with self.assertRaisesRegex(ArchivePasswordError, "incorrect"):
            inspect_archive(
                archive,
                filename="secure.zip",
                password="wrong",
            )

    def test_folder_selection_extracts_only_descendants(self) -> None:
        archive = _zip_bytes(
            {
                "selected/one.txt": b"one",
                "selected/two.txt": b"two",
                "other/three.txt": b"three",
            }
        )
        stored = SimpleNamespace(
            asset=SimpleNamespace(id="asset"),
            version=SimpleNamespace(id="version"),
            blob=SimpleNamespace(id="blob"),
        )
        with patch(
            "govoplan_files.backend.storage.archives.create_file_asset",
            return_value=stored,
        ) as create_asset:
            result = extract_archive_upload(
                object(),  # type: ignore[arg-type]
                tenant_id="tenant-1",
                owner_type="user",
                owner_id="user-1",
                user_id="user-1",
                archive_data=archive,
                filename="selection.zip",
                folder="imports",
                campaign_id=None,
                selected_paths=["selected"],
            )

        self.assertEqual(2, len(result))
        self.assertEqual(
            {"imports/selected/one.txt", "imports/selected/two.txt"},
            {
                call.kwargs["display_path"]
                for call in create_asset.call_args_list
            },
        )


if __name__ == "__main__":
    unittest.main()