import hashlib
import json
from pathlib import Path
import struct
import tempfile
import unittest
from unittest import mock
import zipfile

from package_windows import package, CRT_LIBRARIES
from embeddinggemma2 import runtime_provenance as provenance
import stage_extraction_worker as extraction


def executable(path, machine=0x8664):
    path.parent.mkdir(parents=True, exist_ok=True)
    data = bytearray(70)
    data[:2] = b"MZ"
    struct.pack_into("<I", data, 60, 64)
    data[64:68] = b"PE\0\0"
    struct.pack_into("<H", data, 68, machine)
    path.write_bytes(data)


def inventory(resources):
    worker = resources / "searchi-gemma-worker.exe"
    records = []
    for name in ["ggml-cpu-x64.dll", "ggml-vulkan.dll", "vulkan-1.dll", "ggml.dll", "ggml-base.dll", "llama.dll", "mtmd.dll", *sorted(CRT_LIBRARIES)]:
        data = (resources / name).read_bytes()
        records.append({"path": name, "bytes": len(data), "sha256": hashlib.sha256(data).hexdigest()})
    (resources / "runtime-libraries.json").write_text(json.dumps({"worker_sha256": hashlib.sha256(worker.read_bytes()).hexdigest(), "build_provenance": provenance.fingerprint(resources.parents[3]), "libraries": records}))


class PackageTests(unittest.TestCase):
    def test_extraction_missing_tampered_stale_preserves_existing_package(self):
        for mutation in ("missing", "tampered", "stale"):
            with self.subTest(mutation=mutation), tempfile.TemporaryDirectory() as directory:
                root, resources, app, output = self.fixture(directory)
                output.write_bytes(b"previous package")
                helper = resources / "searchi-extract-worker.exe"
                if mutation == "missing":
                    helper.unlink()
                elif mutation == "tampered":
                    helper.write_bytes(helper.read_bytes() + b"changed")
                else:
                    source = root / "crates/searchi-extract-worker/src/main.rs"
                    source.write_bytes(source.read_bytes() + b"\n// changed\n")
                with self.assertRaises(ValueError):
                    package(root, app, output)
                self.assertEqual(output.read_bytes(), b"previous package")
                self.assertEqual(list(output.parent.glob("*.partial")), [])

    def test_optional_cli_manifest_and_architecture(self):
        with tempfile.TemporaryDirectory() as directory:
            root, _, app, output = self.fixture(directory)
            cli = root / "build/cli.exe"
            executable(cli)
            package(root, app, output, cli)
            with zipfile.ZipFile(output) as archive:
                self.assertEqual(archive.read("searchi-cli.exe"), cli.read_bytes())
                manifest = json.loads(archive.read("PACKAGE.json"))
                record = next(row for row in manifest["files"] if row["path"] == "searchi-cli.exe")
                self.assertEqual(record["sha256"], hashlib.sha256(cli.read_bytes()).hexdigest())
                self.assertIn("--db", archive.read("AGENT-SETUP.txt").decode())
            previous = output.read_bytes()
            executable(cli, 0xAA64)
            with self.assertRaisesRegex(ValueError, "architecture differs"):
                package(root, app, output, cli)
            self.assertEqual(output.read_bytes(), previous)

    def test_old_inventory_rejected_without_replacing_archive(self):
        with tempfile.TemporaryDirectory() as directory:
            root, resources, app, output = self.fixture(directory)
            output.write_bytes(b"previous archive")
            inventory = resources / "runtime-libraries.json"
            manifest = json.loads(inventory.read_text())
            del manifest["build_provenance"]
            inventory.write_text(json.dumps(manifest))
            with self.assertRaisesRegex(ValueError, "provenance"):
                package(root, app, output)
            self.assertEqual(output.read_bytes(), b"previous archive")

    def fixture(self, directory):
        root = Path(directory)
        for name in provenance.INPUTS:
            path = root / name
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_bytes((provenance.ROOT / name).read_bytes())
        for name in extraction.fingerprint()["inputs"]:
            path = root / name
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_bytes((extraction.ROOT / name).read_bytes())
        resources = root / "apps/desktop/src-tauri/resources"
        app = root / "build/app.exe"
        executable(app)
        (root / "LICENSE").write_text("application notice")
        for name in ["searchi-gemma-worker.exe", "vulkan-1.dll", "ggml-cpu-x64.dll", "ggml-vulkan.dll", "ggml.dll", "ggml-base.dll", "llama.dll", "mtmd.dll", "media-runtime/bin/ffmpeg.exe", "media-runtime/bin/ffprobe.exe"]:
            executable(resources / name)
        for name in sorted(CRT_LIBRARIES):
            executable(resources / name)
        for name in ["LLAMA-LICENSE.txt", "MSVC-REDISTRIBUTION.txt", "VULKAN-LICENSE.txt", "licenses/Inter-OFL.txt", "licenses/Fluent-MIT.txt", "media-runtime/licenses/LICENSE.txt", "media-runtime/manifest.json"]:
            path = resources / name
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_text("notice")
        (resources / "settings.json").write_text("private settings")
        media = root / ".searchi/media-runtime"
        media.mkdir(parents=True)
        upstream = media / "pinned-windows.zip"
        with zipfile.ZipFile(upstream, "w") as archive:
            for staged, member in [("bin/ffmpeg.exe", "bin/ffmpeg.exe"), ("bin/ffprobe.exe", "bin/ffprobe.exe"), ("licenses/LICENSE.txt", "LICENSE.txt")]:
                archive.writestr(member, (resources / "media-runtime" / staged).read_bytes())
        manifest = root / "tools/media-runtime/manifest.json"
        manifest.parent.mkdir(parents=True)
        manifest.write_text(json.dumps({"platforms": {"windows-x86_64": {
            "archive": upstream.name, "bytes": upstream.stat().st_size,
            "sha256": hashlib.sha256(upstream.read_bytes()).hexdigest()}}}))
        (resources / "media-runtime/manifest.json").write_bytes(manifest.read_bytes())
        inventory(resources)
        helper = root / "build/searchi-extract-worker.exe"
        executable(helper)
        extraction.stage(helper, extraction.fingerprint(root), resources, root)
        return root, resources, app, root / "archive.zip"

    def test_valid_pe_media_mutation_and_notice_corruption_preserve_previous_archive(self):
        for name in ["bin/ffmpeg.exe", "bin/ffprobe.exe", "licenses/LICENSE.txt", "manifest.json"]:
            with self.subTest(name=name), tempfile.TemporaryDirectory() as directory:
                root, resources, app, output = self.fixture(directory)
                output.write_bytes(b"previous package")
                path = resources / "media-runtime" / name
                path.write_bytes(path.read_bytes() + b"tampered")
                with self.assertRaisesRegex(ValueError, "differs from pinned"):
                    package(root, app, output)
                self.assertEqual(output.read_bytes(), b"previous package")
                self.assertEqual(list(output.parent.glob("*.partial")), [])

    def test_pinned_archive_corruption_preserves_previous_package(self):
        with tempfile.TemporaryDirectory() as directory:
            root, _, app, output = self.fixture(directory)
            output.write_bytes(b"previous package")
            source = root / ".searchi/media-runtime/pinned-windows.zip"
            source.write_bytes(source.read_bytes() + b"changed")
            with self.assertRaisesRegex(ValueError, "provenance mismatch"):
                package(root, app, output)
            self.assertEqual(output.read_bytes(), b"previous package")

    def test_media_change_during_archive_write_preserves_previous_package(self):
        with tempfile.TemporaryDirectory() as directory:
            root, resources, app, output = self.fixture(directory)
            output.write_bytes(b"previous package")
            original = zipfile.ZipFile.open
            def mutate(archive, name, mode="r", *args, **kwargs):
                if name == "media-runtime/bin/ffmpeg.exe" and mode == "w":
                    path = resources / name
                    path.write_bytes(path.read_bytes() + b"changed")
                return original(archive, name, mode, *args, **kwargs)
            with mock.patch.object(zipfile.ZipFile, "open", mutate), self.assertRaisesRegex(ValueError, "changed while packaging"):
                package(root, app, output)
            self.assertEqual(output.read_bytes(), b"previous package")

    def test_duplicate_helper_and_link_notice_are_rejected_before_output(self):
        for kind in ["duplicate", "link"]:
            with self.subTest(kind=kind), tempfile.TemporaryDirectory() as directory:
                root, resources, app, output = self.fixture(directory)
                output.write_bytes(b"previous package")
                upstream = root / ".searchi/media-runtime/pinned-windows.zip"
                with zipfile.ZipFile(upstream, "a") as archive:
                    if kind == "duplicate":
                        archive.writestr("other/ffmpeg.exe", (resources / "media-runtime/bin/ffmpeg.exe").read_bytes())
                    else:
                        entry = zipfile.ZipInfo("COPYING.txt")
                        entry.create_system = 3
                        entry.external_attr = 0o120777 << 16
                        archive.writestr(entry, "somewhere")
                manifest = root / "tools/media-runtime/manifest.json"
                data = json.loads(manifest.read_text())
                data["platforms"]["windows-x86_64"].update(bytes=upstream.stat().st_size, sha256=hashlib.sha256(upstream.read_bytes()).hexdigest())
                manifest.write_text(json.dumps(data))
                (resources / "media-runtime/manifest.json").write_bytes(manifest.read_bytes())
                with self.assertRaisesRegex(ValueError, "Invalid or duplicate"):
                    package(root, app, output)
                self.assertEqual(output.read_bytes(), b"previous package")

    def test_complete_archive_excludes_settings_and_verifies_recorded_bytes(self):
        with tempfile.TemporaryDirectory() as directory:
            root, _, app, output = self.fixture(directory)
            package(root, app, output)
            with zipfile.ZipFile(output) as archive:
                self.assertNotIn("settings.json", archive.namelist())
                self.assertIn("media-runtime/bin/ffprobe.exe", archive.namelist())
                for record in json.loads(archive.read("PACKAGE.json"))["files"]:
                    self.assertEqual(hashlib.sha256(archive.read(record["path"])).hexdigest(), record["sha256"])
                    self.assertEqual(len(archive.read(record["path"])), record["bytes"])
                self.assertEqual(len(archive.namelist()), len(set(archive.namelist())))

    def test_mixed_architecture_preserves_previous_archive(self):
        with tempfile.TemporaryDirectory() as directory:
            root, resources, app, output = self.fixture(directory)
            output.write_bytes(b"previous package")
            executable(resources / "searchi-gemma-worker.exe", 0xAA64)
            inventory(resources)
            with self.assertRaisesRegex(ValueError, "architecture differs"):
                package(root, app, output)
            self.assertEqual(output.read_bytes(), b"previous package")

    def test_changed_source_records_exact_archived_bytes(self):
        with tempfile.TemporaryDirectory() as directory:
            root, _, app, output = self.fixture(directory)
            original_open = zipfile.ZipFile.open

            def mutate_source(archive, name, mode="r", *args, **kwargs):
                if name == "searchi.exe" and mode == "w":
                    # Change bytes after preflight and after the input stream
                    # opens, at the point ZIP copying starts.
                    app.write_bytes(app.read_bytes() + b"changed while packaging")
                return original_open(archive, name, mode, *args, **kwargs)

            with mock.patch.object(zipfile.ZipFile, "open", mutate_source):
                package(root, app, output)
            with zipfile.ZipFile(output) as archive:
                content = archive.read("searchi.exe")
                record = json.loads(archive.read("PACKAGE.json"))["files"][0]
                self.assertIn(b"changed while packaging", content)
                self.assertEqual(record["bytes"], len(content))
                self.assertEqual(record["sha256"], hashlib.sha256(content).hexdigest())

    def test_mid_write_failure_preserves_previous_archive_and_cleans_temporary(self):
        with tempfile.TemporaryDirectory() as directory:
            root, _, app, output = self.fixture(directory)
            output.write_bytes(b"previous package")
            original_open = zipfile.ZipFile.open

            def fail_second_entry(archive, name, mode="r", *args, **kwargs):
                if name == "LICENSE.txt" and mode == "w":
                    raise OSError("simulated storage failure")
                return original_open(archive, name, mode, *args, **kwargs)

            with mock.patch.object(zipfile.ZipFile, "open", fail_second_entry):
                with self.assertRaisesRegex(OSError, "storage failure"):
                    package(root, app, output)
            self.assertEqual(output.read_bytes(), b"previous package")
            self.assertEqual(list(output.parent.glob("*.partial")), [])

    def test_missing_named_notice_preserves_previous_archive(self):
        with tempfile.TemporaryDirectory() as directory:
            root, resources, app, output = self.fixture(directory)
            output.write_bytes(b"previous package")
            (resources / "licenses/Inter-OFL.txt").unlink()
            with self.assertRaisesRegex(ValueError, "Inter-OFL"):
                package(root, app, output)
            self.assertEqual(output.read_bytes(), b"previous package")

    def test_mixed_loader_architecture_preserves_previous_archive(self):
        with tempfile.TemporaryDirectory() as directory:
            root, resources, app, output = self.fixture(directory)
            output.write_bytes(b"previous package")
            executable(resources / "vulkan-1.dll", 0xAA64)
            inventory(resources)
            with self.assertRaisesRegex(ValueError, "architecture differs: vulkan-1.dll"):
                package(root, app, output)
            self.assertEqual(output.read_bytes(), b"previous package")

    def test_inventory_includes_dispatch_libraries_but_not_unlisted_dll(self):
        with tempfile.TemporaryDirectory() as directory:
            root, resources, app, output = self.fixture(directory)
            executable(resources / "private-plugin.dll")
            package(root, app, output)
            with zipfile.ZipFile(output) as archive:
                self.assertIn("ggml-cpu-x64.dll", archive.namelist())
                self.assertIn("ggml-vulkan.dll", archive.namelist())
                self.assertIn("runtime-libraries.json", archive.namelist())
                self.assertIn("msvcp140.dll", archive.namelist())
                self.assertIn("MSVC-REDISTRIBUTION.txt", archive.namelist())
                self.assertNotIn("private-plugin.dll", archive.namelist())

    def test_missing_inventory_or_corrupt_library_preserves_archive(self):
        with tempfile.TemporaryDirectory() as directory:
            root, resources, app, output = self.fixture(directory)
            output.write_bytes(b"previous package")
            (resources / "ggml-cpu-x64.dll").write_bytes(b"corrupt")
            with self.assertRaisesRegex(ValueError, "checksum mismatch"):
                package(root, app, output)
            self.assertEqual(output.read_bytes(), b"previous package")
            (resources / "runtime-libraries.json").unlink()
            with self.assertRaisesRegex(ValueError, "inventory"):
                package(root, app, output)
            self.assertEqual(output.read_bytes(), b"previous package")

    def test_dependency_change_during_copy_preserves_archive(self):
        with tempfile.TemporaryDirectory() as directory:
            root, resources, app, output = self.fixture(directory)
            output.write_bytes(b"previous package")
            original_open = zipfile.ZipFile.open

            def mutate_library(archive, name, mode="r", *args, **kwargs):
                if name == "ggml-cpu-x64.dll" and mode == "w":
                    library = resources / name
                    library.write_bytes(library.read_bytes() + b"changed")
                return original_open(archive, name, mode, *args, **kwargs)

            with mock.patch.object(zipfile.ZipFile, "open", mutate_library):
                with self.assertRaisesRegex(ValueError, "changed while packaging"):
                    package(root, app, output)
            self.assertEqual(output.read_bytes(), b"previous package")
            self.assertEqual(list(output.parent.glob("*.partial")), [])

    def test_missing_crt_or_notice_preserves_previous_archive(self):
        with tempfile.TemporaryDirectory() as directory:
            root, resources, app, output = self.fixture(directory)
            output.write_bytes(b"previous package")
            (resources / "msvcp140.dll").unlink()
            with self.assertRaisesRegex(ValueError, "msvcp140"):
                package(root, app, output)
            self.assertEqual(output.read_bytes(), b"previous package")
            executable(resources / "msvcp140.dll")
            (resources / "MSVC-REDISTRIBUTION.txt").unlink()
            with self.assertRaisesRegex(ValueError, "MSVC-REDISTRIBUTION"):
                package(root, app, output)
            self.assertEqual(output.read_bytes(), b"previous package")

    def test_inventory_change_during_copy_preserves_previous_archive(self):
        with tempfile.TemporaryDirectory() as directory:
            root, resources, app, output = self.fixture(directory)
            output.write_bytes(b"previous package")
            original_open = zipfile.ZipFile.open

            def mutate_inventory(archive, name, mode="r", *args, **kwargs):
                if name == "runtime-libraries.json" and mode == "w":
                    manifest = resources / name
                    manifest.write_bytes(manifest.read_bytes() + b" ")
                return original_open(archive, name, mode, *args, **kwargs)

            with mock.patch.object(zipfile.ZipFile, "open", mutate_inventory):
                with self.assertRaisesRegex(ValueError, "changed while packaging"):
                    package(root, app, output)
            self.assertEqual(output.read_bytes(), b"previous package")


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