"""The cooked shader interface must preserve GPU resource kinds and element strides."""

import importlib.util
import json
import os
from pathlib import Path
import struct
import subprocess
import sys
import tempfile
import unittest


SCRIPT = Path(__file__).resolve().parents[1] / "tools" / "compile_shader.py"
SPEC = importlib.util.spec_from_file_location("faset_compile_shader", SCRIPT)
assert SPEC and SPEC.loader
shader = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(shader)


def parameter(name: str, index: int, shape: str, access: str, stride: int | None = None):
    result = {"kind": "scalar", "scalarType": "uint32", "sizes": []}
    if stride is not None:
        result["sizes"] = [{"kind": "uniform", "value": stride, "alignment": 16}]
    return {
        "name": name,
        "binding": {"kind": "descriptorTableSlot", "index": index, "space": 0},
        "type": {
            "kind": "resource",
            "baseShape": shape,
            "access": access,
            "resultType": result,
            "sizes": [{"kind": "descriptorTableSlot", "value": 1}],
        },
    }


class ReflectionTests(unittest.TestCase):
    def test_gpu_storage_resources_keep_kind_and_stride(self):
        parameters = [
            parameter("instances", 0, "structuredBuffer", "read", 224),
            parameter("visibleIds", 1, "structuredBuffer", "readWrite", 4),
            parameter("depthOutput", 2, "texture2D", "readWrite"),
            parameter("depthInput", 3, "texture2D", "read"),
        ]
        raw = {
            "parameters": parameters,
            "entryPoints": [{
                "name": "gpuCullMain", "stage": "compute",
                "bindings": [{"name": p["name"], "binding": {"used": 1}} for p in parameters],
            }],
        }
        spirv = struct.pack("<5I", 0x07230203, 0x00010600, 0, 1, 0)
        layout = shader.normalize(raw, spirv, "gpuCullMain")["layout"]
        descriptors = layout["descriptors"]
        self.assertEqual(
            [(d["type"], d.get("element_stride")) for d in descriptors],
            [("storage_buffer", 224), ("storage_buffer", 4),
             ("storage_image_2d", None), ("sampled_image_2d", None)],
        )

    def test_real_slang_compute_interface(self):
        compiler = os.environ["FASET_TEST_SLANGC"]
        with tempfile.TemporaryDirectory(prefix="faset-shader-reflection-") as directory:
            process = subprocess.run(
                [sys.executable, str(SCRIPT), "--compiler", compiler, "--source",
                 str(SCRIPT.parents[1] / "shaders" / "gpu_scene.slang"), "--entry",
                 "gpuCullMain", "--define", "FASET_GPU_CULL=1", "--output", directory],
                capture_output=True, text=True,
            )
            self.assertEqual(process.returncode, 0, process.stderr)
            metadata = json.loads((Path(directory) / "gpuCullMain.reflection.json").read_text())
            bindings = {item["binding"]: item for item in metadata["layout"]["descriptors"]}
            self.assertEqual([bindings[n]["element_stride"] for n in (0, 1, 2, 3, 4, 5, 6, 9)],
                             [224, 16, 16, 4, 16, 4, 4, 208])
            self.assertEqual([bindings[n]["type"] for n in (7, 8)],
                             ["sampled_image_2d", "sampled_image_2d"])

    def test_hzb_storage_image_has_explicit_r32f_format(self):
        compiler = os.environ["FASET_TEST_SLANGC"]
        with tempfile.TemporaryDirectory(prefix="faset-hzb-format-") as directory:
            process = subprocess.run(
                [sys.executable, str(SCRIPT), "--compiler", compiler, "--source",
                 str(SCRIPT.parents[1] / "shaders" / "gpu_scene.slang"), "--entry",
                 "gpuHzbMain", "--define", "FASET_GPU_HZB=1", "--output", directory],
                capture_output=True, text=True,
            )
            self.assertEqual(process.returncode, 0, process.stderr)
            bytecode = (Path(directory) / "gpuHzbMain.spv").read_bytes()
            words = struct.unpack(f"<{len(bytecode) // 4}I", bytecode)
            capabilities = set()
            image_formats = []
            offset = 5
            while offset < len(words):
                count, opcode = words[offset] >> 16, words[offset] & 0xffff
                self.assertGreater(count, 0)
                if opcode == 17:  # OpCapability
                    capabilities.add(words[offset + 1])
                elif opcode == 25:  # OpTypeImage; last operand is ImageFormat
                    image_formats.append((words[offset + 7], words[offset + 8]))
                offset += count
            self.assertNotIn(55, capabilities)  # StorageImageReadWithoutFormat
            self.assertNotIn(56, capabilities)  # StorageImageWriteWithoutFormat
            self.assertIn((2, 3), image_formats)  # storage image, R32f


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