from __future__ import annotations

import json
import struct
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path

FRONTEND_ROOT = Path(__file__).resolve().parents[1]
TOOLS_ROOT = FRONTEND_ROOT / "tools"
sys.path.insert(0, str(TOOLS_ROOT))

from _common import (  # noqa: E402
    CORE_DYLIB,
    GROUP_DYLIBS,
    ORIGINAL_DEPLOYMENT_SEED,
    ORIGINAL_REPORTING_SEED,
)
from _path_patch import PATH_TEMPLATE, patch_initial_daily_path  # noqa: E402


TEST_CHANNEL = "0123456789abcdef0123456789abcdef"


class InitialDailyPathPatchTests(unittest.TestCase):
    def test_both_group_dylibs_retarget_validated_cfstring(self) -> None:
        expected_architectures = {"A": "arm64", "B": "arm64e"}
        encoded_path = PATH_TEMPLATE.format(channel=TEST_CHANNEL).encode() + b"\x00"

        for group, path in GROUP_DYLIBS.items():
            with self.subTest(group=group):
                source = path.read_bytes()
                patched, metadata = patch_initial_daily_path(
                    source, TEST_CHANNEL, label=path.name
                )
                self.assertEqual(metadata["architecture"], expected_architectures[group])
                self.assertEqual(metadata["path"], encoded_path[:-1].decode())

                cave = int(metadata["path_file_offset"])
                cfstring = int(metadata["cfstring_file_offset"])
                reference = int(metadata["reference_file_offset"])
                self.assertEqual(patched[cave : cave + len(encoded_path)], encoded_path)
                self.assertEqual(
                    struct.unpack_from("<Q", patched, cfstring + 24)[0],
                    len(encoded_path) - 1,
                )

                changed = {
                    index
                    for index, (before, after) in enumerate(zip(source, patched))
                    if before != after
                }
                allowed = (
                    set(range(cave, cave + len(encoded_path)))
                    | set(range(reference, reference + 8))
                    | set(range(cfstring + 24, cfstring + 32))
                )
                self.assertTrue(changed)
                self.assertLessEqual(changed, allowed)

                old_raw = struct.unpack_from("<Q", source, reference)[0]
                new_raw = struct.unpack_from("<Q", patched, reference)[0]
                if group == "A":
                    self.assertEqual(new_raw, int(metadata["path_vmaddr"], 16))
                else:
                    target_mask = (1 << 43) - 1
                    self.assertEqual(old_raw & ~target_mask, new_raw & ~target_mask)
                    self.assertEqual(
                        new_raw & target_mask, int(metadata["path_vmaddr"], 16)
                    )

    def test_fails_closed_when_cfstring_reference_differs(self) -> None:
        for group, path in GROUP_DYLIBS.items():
            with self.subTest(group=group):
                source = path.read_bytes()
                _patched, metadata = patch_initial_daily_path(source, TEST_CHANNEL)
                damaged = bytearray(source)
                reference = int(metadata["reference_file_offset"])
                damaged[reference] ^= 1
                with self.assertRaisesRegex(SystemExit, "Refusing unsafe"):
                    patch_initial_daily_path(bytes(damaged), TEST_CHANNEL)

    def test_fails_closed_when_validated_cave_differs(self) -> None:
        for group, path in GROUP_DYLIBS.items():
            with self.subTest(group=group):
                source = path.read_bytes()
                _patched, metadata = patch_initial_daily_path(source, TEST_CHANNEL)
                damaged = bytearray(source)
                damaged[int(metadata["path_file_offset"])] = 0x41
                with self.assertRaisesRegex(SystemExit, "zero-filled __TEXT cave differs"):
                    patch_initial_daily_path(bytes(damaged), TEST_CHANNEL)

    def test_fails_closed_for_unknown_macho_uuid(self) -> None:
        source = bytearray(GROUP_DYLIBS["A"].read_bytes())
        # LC_UUID payload for the known group-A input.
        uuid_bytes = bytes.fromhex("73827b4262ba3a5989ada4fe6f67e266")
        uuid_offset = source.find(uuid_bytes)
        self.assertGreater(uuid_offset, 0)
        source[uuid_offset] ^= 1
        with self.assertRaisesRegex(SystemExit, "unsupported Mach-O UUID"):
            patch_initial_daily_path(bytes(source), TEST_CHANNEL)

    def test_fat_core_retargets_both_slices(self) -> None:
        source = CORE_DYLIB.read_bytes()
        self.assertEqual(source[:4], bytes.fromhex("cafebabe"))
        patched, metadata = patch_initial_daily_path(
            source, TEST_CHANNEL, label=CORE_DYLIB.name
        )
        self.assertEqual(metadata["container"], "fat")
        expected = PATH_TEMPLATE.format(channel=TEST_CHANNEL).encode() + b"\x00"
        self.assertEqual(metadata["path"], expected[:-1].decode())
        slices = metadata["slices"]
        self.assertEqual(len(slices), 2)
        self.assertEqual(
            {item["architecture"] for item in slices},
            {"arm64", "arm64e"},
        )
        for item in slices:
            slice_off = int(item["fat_slice_offset"])
            cave = slice_off + int(item["path_file_offset"])
            cfstring = slice_off + int(item["cfstring_file_offset"])
            self.assertEqual(patched[cave : cave + len(expected)], expected)
            self.assertEqual(
                struct.unpack_from("<Q", patched, cfstring + 24)[0],
                len(expected) - 1,
            )

    def test_secondary_pack_manifest_records_native_path_patch(self) -> None:
        with tempfile.TemporaryDirectory() as temp:
            command = [
                sys.executable,
                str(TOOLS_ROOT / "patch_secondary_packs.py"),
                "--deployment-seed",
                ORIGINAL_DEPLOYMENT_SEED,
                "--reporting-seed",
                ORIGINAL_REPORTING_SEED,
                "--channel-id",
                TEST_CHANNEL,
                "--out",
                temp,
            ]
            subprocess.run(command, check=True, capture_output=True, text=True)
            manifest = json.loads((Path(temp) / "MANIFEST.json").read_text())
            path_patch = manifest["initial_daily_path_patch"]
            self.assertEqual(
                path_patch["path"], PATH_TEMPLATE.format(channel=TEST_CHANNEL)
            )
            self.assertEqual(
                {group: item["architecture"] for group, item in path_patch["groups"].items()},
                {"A": "arm64", "B": "arm64e"},
            )


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