import contextlib
import hashlib
import importlib.util
import io
import json
import pathlib
import shutil
import struct
import subprocess
import tempfile
import unittest
from unittest import mock


MODULE_PATH = pathlib.Path(__file__).with_name("w3cs-pack-build.py")
SPEC = importlib.util.spec_from_file_location("w3cs_pack_build", MODULE_PATH)
MODULE = importlib.util.module_from_spec(SPEC)
assert SPEC.loader
SPEC.loader.exec_module(MODULE)


class PackBuildTests(unittest.TestCase):
    def write_bundle(self, spool, body):
        digest = hashlib.sha256(body).hexdigest()
        (spool / f"{digest}.bundle").write_bytes(body)
        return digest

    def write_webp_bundle(self, spool, encoded_bytes=2048):
        encoded = (b"RIFF" + struct.pack("<I", encoded_bytes - 8) +
                   b"WEBP" + b"x" * (encoded_bytes - 12))
        payload = struct.pack("<IIQQ", 123, len(encoded), 0, 0) + encoded
        body = struct.pack("<BBHI", 7, 0, 0, len(payload)) + payload
        return self.write_bundle(spool, body)

    def write_manifest(self, spool, profile, replay, *digests):
        manifests = spool / "manifests"
        manifests.mkdir(exist_ok=True)
        (manifests / f"{profile}__{replay}.digests").write_text(
            "".join(digest + "\n" for digest in digests), encoding="ascii")

    @staticmethod
    def fake_zstd(args, check):
        assert check
        shutil.copyfile(args[-3], args[-1])
        return subprocess.CompletedProcess(args, 0)

    def fixture(self, root, catalog, aliases=()):
        spool = root / "spool"
        spool.mkdir(exist_ok=True)
        replay_root = root / "replays"
        replay_root.mkdir(exist_ok=True)
        replays_json = root / "replays.json"
        replays_json.write_text(json.dumps(list(aliases)), encoding="utf-8")
        catalog_json = root / "catalog.json"
        catalog_json.write_text(json.dumps(
            {"schema": 3, "replays": catalog}), encoding="utf-8")
        return spool, replay_root

    def build(self, root, spool, replay_root):
        with mock.patch.object(
                MODULE.subprocess, "run",
                side_effect=self.fake_zstd) as zstd_run, \
                contextlib.redirect_stdout(io.StringIO()), \
                contextlib.redirect_stderr(io.StringIO()):
            MODULE.main(spool=str(spool),
                        replays_json=str(root / "replays.json"),
                        catalog_json=str(root / "catalog.json"),
                        replay_root=str(replay_root))
        return zstd_run

    @staticmethod
    def pointer(spool, replay):
        path = spool / "pack-out" / f"replay-{replay}.json"
        return json.loads(path.read_text(encoding="utf-8"))

    def test_catalog_only_on_demand_replay_gets_a_manifest(self):
        with tempfile.TemporaryDirectory() as directory:
            root = pathlib.Path(directory)
            replay = "8" * 40
            map_sha1 = "a" * 40
            catalog = {replay: {
                "profile": "native-121b",
                "mapPath": "Maps\\SW\\(2)TerenasStand.w3x",
                "mapSha1": map_sha1,
            }}
            spool, replay_root = self.fixture(root, catalog)
            digest = self.write_bundle(spool, b"a" * 4096)
            self.write_manifest(spool, "native-121b", replay, digest)
            self.build(root, spool, replay_root)

            manifest = self.pointer(spool, replay)
            self.assertEqual(
                f"native-121b__map-{map_sha1}", manifest["key"])
            self.assertEqual(1, manifest["packs"][0]["entries"])

    def test_webp_quality_variant_is_excluded_from_generic_pack(self):
        with tempfile.TemporaryDirectory() as directory:
            root = pathlib.Path(directory)
            replay = "0" * 40
            map_sha1 = "9" * 40
            catalog = {replay: {
                "profile": "native-1285",
                "mapPath": "Maps\\Download\\WebpPolicy.w3x",
                "mapSha1": map_sha1,
            }}
            spool, replay_root = self.fixture(root, catalog)
            raw_digest = self.write_bundle(spool, b"r" * 4096)
            webp_digest = self.write_webp_bundle(spool)
            self.write_manifest(
                spool, "native-1285", replay, raw_digest, webp_digest)

            self.build(root, spool, replay_root)

            manifest = self.pointer(spool, replay)
            self.assertEqual(1, manifest["packs"][0]["entries"])
            state_path = next((spool / "pack-state").glob("*.built"))
            state = json.loads(state_path.read_text())
            self.assertEqual([raw_digest], state["digests"])

    def test_same_basename_with_different_content_has_distinct_packs(self):
        with tempfile.TemporaryDirectory() as directory:
            root = pathlib.Path(directory)
            first, second = "1" * 40, "2" * 40
            first_map, second_map = "a" * 40, "b" * 40
            catalog = {
                first: {"profile": "native-121b",
                        "mapPath": "Maps\\SW\\(2)TerenasStand.w3x",
                        "mapSha1": first_map},
                second: {"profile": "native-121b",
                         "mapPath": "Maps\\FrozenThrone\\(2)TerenasStand.w3x",
                         "mapSha1": second_map},
            }
            spool, replay_root = self.fixture(root, catalog)
            first_digest = self.write_bundle(spool, b"a" * 4096)
            second_digest = self.write_bundle(spool, b"b" * 4096)
            self.write_manifest(spool, "native-121b", first, first_digest)
            self.write_manifest(spool, "native-121b", second, second_digest)
            self.build(root, spool, replay_root)

            first_pointer = self.pointer(spool, first)
            second_pointer = self.pointer(spool, second)
            self.assertEqual(
                f"native-121b__map-{first_map}", first_pointer["key"])
            self.assertEqual(
                f"native-121b__map-{second_map}", second_pointer["key"])
            self.assertNotEqual(first_pointer["packs"], second_pointer["packs"])

    def test_materialized_descriptor_resolves_non_catalog_alias(self):
        with tempfile.TemporaryDirectory() as directory:
            root = pathlib.Path(directory)
            spool, replay_root = self.fixture(root, {})
            replay = "patch-proof"
            map_sha1 = "c" * 40
            (replay_root / f"{replay}.json").write_text(json.dumps({
                "engineProfile": "native-1285",
                "mapPath": "Maps\\Download\\Arena.w3x",
                "mapContentSha1": map_sha1,
            }), encoding="utf-8")
            digest = self.write_bundle(spool, b"c" * 4096)
            self.write_manifest(spool, "native-1285", replay, digest)
            self.build(root, spool, replay_root)
            pointer = self.pointer(spool, replay)
            self.assertEqual(
                f"native-1285__map-{map_sha1}", pointer["key"])

    def test_manifest_records_replay_specific_entries_and_prunes_stale(self):
        with tempfile.TemporaryDirectory() as directory:
            root = pathlib.Path(directory)
            first, second = "5" * 40, "6" * 40
            map_sha1 = "f" * 40
            metadata = {"profile": "native-121b",
                        "mapPath": "Maps\\Download\\Shared.w3x",
                        "mapSha1": map_sha1}
            spool, replay_root = self.fixture(
                root, {first: metadata, second: metadata})
            shared = self.write_bundle(spool, b"s" * 4096)
            first_only = self.write_bundle(spool, b"f" * 4096)
            second_one = self.write_bundle(spool, b"1" * 4096)
            second_two = self.write_bundle(spool, b"2" * 4096)
            self.write_manifest(
                spool, "native-121b", first, shared, first_only)
            self.write_manifest(
                spool, "native-121b", second, shared, second_one, second_two)
            out_dir = spool / "pack-out"
            out_dir.mkdir()
            stale = out_dir / "replay-retired.json"
            stale.write_text("{}", encoding="utf-8")

            self.build(root, spool, replay_root)

            self.assertEqual(2, self.pointer(spool, first)["replayEntries"])
            self.assertEqual(3, self.pointer(spool, second)["replayEntries"])
            self.assertFalse(stale.exists())

    def test_missing_prior_zpack_forces_hysteresis_rebuild(self):
        with tempfile.TemporaryDirectory() as directory:
            root = pathlib.Path(directory)
            replay = "3" * 40
            catalog = {replay: {
                "profile": "native-121b",
                "mapPath": "Maps\\Download\\MissingPack.w3x",
                "mapSha1": "d" * 40,
            }}
            spool, replay_root = self.fixture(root, catalog)
            digests = [self.write_bundle(
                spool, bytes([index]) * 4096) for index in range(20)]
            self.write_manifest(spool, "native-121b", replay, *digests)
            self.build(root, spool, replay_root)
            old_pack = self.pointer(spool, replay)["packs"][0]
            old_path = spool / "pack-out" / old_pack["object"]
            old_path.unlink()

            digests.append(self.write_bundle(spool, b"new" * 1366))
            self.write_manifest(spool, "native-121b", replay, *digests)
            zstd_run = self.build(root, spool, replay_root)

            self.assertEqual(1, zstd_run.call_count)
            new_pack = self.pointer(spool, replay)["packs"][0]
            self.assertNotEqual(old_pack["object"], new_pack["object"])
            new_path = spool / "pack-out" / new_pack["object"]
            self.assertEqual(new_pack["packedBytes"], new_path.stat().st_size)

    def test_wrong_size_prior_zpack_forces_same_fingerprint_rebuild(self):
        with tempfile.TemporaryDirectory() as directory:
            root = pathlib.Path(directory)
            replay = "4" * 40
            catalog = {replay: {
                "profile": "native-121b",
                "mapPath": "Maps\\Download\\WrongSizePack.w3x",
                "mapSha1": "e" * 40,
            }}
            spool, replay_root = self.fixture(root, catalog)
            digest = self.write_bundle(spool, b"e" * 4096)
            self.write_manifest(spool, "native-121b", replay, digest)
            self.build(root, spool, replay_root)
            old_pack = self.pointer(spool, replay)["packs"][0]
            old_path = spool / "pack-out" / old_pack["object"]
            old_path.write_bytes(b"truncated")

            zstd_run = self.build(root, spool, replay_root)

            self.assertEqual(1, zstd_run.call_count)
            new_pack = self.pointer(spool, replay)["packs"][0]
            self.assertEqual(old_pack["object"], new_pack["object"])
            self.assertEqual(new_pack["packedBytes"], old_path.stat().st_size)

    def test_hysteresis_reports_only_digests_in_the_reused_pack(self):
        with tempfile.TemporaryDirectory() as directory:
            root = pathlib.Path(directory)
            first, second = "7" * 40, "8" * 40
            metadata = {"profile": "native-121b",
                        "mapPath": "Maps\\Download\\Hysteresis.w3x",
                        "mapSha1": "1" * 40}
            spool, replay_root = self.fixture(
                root, {first: metadata, second: metadata})
            original = [self.write_bundle(
                spool, bytes([index]) * 4096) for index in range(20)]
            self.write_manifest(spool, "native-121b", first, *original)
            self.build(root, spool, replay_root)

            new_digest = self.write_bundle(spool, b"new" * 1366)
            self.write_manifest(
                spool, "native-121b", second, original[0], new_digest)
            zstd_run = self.build(root, spool, replay_root)

            self.assertEqual(0, zstd_run.call_count)
            self.assertEqual(1, self.pointer(spool, second)["replayEntries"])

    def test_removed_digest_forces_pack_rebuild_despite_hysteresis(self):
        with tempfile.TemporaryDirectory() as directory:
            root = pathlib.Path(directory)
            replay = "a" * 40
            catalog = {replay: {
                "profile": "native-1285",
                "mapPath": "Maps\\Download\\Shrink.w3x",
                "mapSha1": "3" * 40,
            }}
            spool, replay_root = self.fixture(root, catalog)
            keep = self.write_bundle(spool, b"k" * 4096)
            retire = self.write_bundle(spool, b"x" * 4096)
            self.write_manifest(
                spool, "native-1285", replay, keep, retire)
            self.build(root, spool, replay_root)
            old_object = self.pointer(spool, replay)["packs"][0]["object"]

            webp = self.write_webp_bundle(spool)
            self.write_manifest(spool, "native-1285", replay, keep, webp)
            zstd_run = self.build(root, spool, replay_root)

            self.assertEqual(1, zstd_run.call_count)
            pointer = self.pointer(spool, replay)
            self.assertEqual(1, pointer["packs"][0]["entries"])
            self.assertNotEqual(old_object, pointer["packs"][0]["object"])

    def test_matching_legacy_state_records_exact_built_digests(self):
        with tempfile.TemporaryDirectory() as directory:
            root = pathlib.Path(directory)
            replay = "9" * 40
            catalog = {replay: {
                "profile": "native-121b",
                "mapPath": "Maps\\Download\\LegacyState.w3x",
                "mapSha1": "2" * 40,
            }}
            spool, replay_root = self.fixture(root, catalog)
            digest = self.write_bundle(spool, b"l" * 4096)
            self.write_manifest(spool, "native-121b", replay, digest)
            self.build(root, spool, replay_root)
            state_path = next((spool / "pack-state").glob("*.built"))
            state = json.loads(state_path.read_text())
            state.pop("digests")
            state_path.write_text(json.dumps(state), encoding="utf-8")

            zstd_run = self.build(root, spool, replay_root)

            self.assertEqual(0, zstd_run.call_count)
            upgraded = json.loads(state_path.read_text())
            self.assertEqual([digest], upgraded["digests"])


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