import rdtest
import struct
import renderdoc as rd


class VK_Shader_Editing(rdtest.TestCase):
    demos_test_name = 'VK_Shader_Editing'

    def check_capture(self):
        eid = self.find_action("Draw 1").nextAction.eventId
        self.set_event(eid, False)

        pipe = self.controller.GetPipelineState()

        fsrefl1 = pipe.GetShaderReflection(rd.ShaderStage.Fragment)

        eid = self.find_action("Draw 2").nextAction.eventId
        self.set_event(eid, False)

        pipe = self.controller.GetPipelineState()

        fsrefl2 = pipe.GetShaderReflection(rd.ShaderStage.Fragment)
        vsrefl = pipe.GetShaderReflection(rd.ShaderStage.Vertex)

        tex = pipe.GetOutputTargets()[0].resource

        # Both triangles should be green
        self.check_pixel_value(tex, 0.25, 0.5, [0.0, 1.0, 0.0, 1.0])
        self.check_pixel_value(tex, 0.75, 0.5, [0.0, 1.0, 0.0, 1.0])

        rdtest.log.success("Values are as expected initially")

        source = fsrefl1.debugInfo.files[0].contents.replace('#if 1', '#if 0')

        newShader = self.controller.BuildTargetShader(
            fsrefl1.entryPoint,
            rd.ShaderEncoding.GLSL,
            bytes(source, "UTF-8"),
            rd.ShaderCompileFlags(),
            rd.ShaderStage.Fragment,
        )

        if len(newShader[1]) != 0:
            raise rdtest.TestFailureException(f"Failed to compile edited shader: {newShader[1]}")

        FS1 = newShader[0]

        source = fsrefl2.debugInfo.files[0].contents.replace('#if 1', '#if 0')

        newShader = self.controller.BuildTargetShader(
            fsrefl2.entryPoint,
            rd.ShaderEncoding.GLSL,
            bytes(source, "UTF-8"),
            rd.ShaderCompileFlags(),
            rd.ShaderStage.Fragment,
        )

        if len(newShader[1]) != 0:
            raise rdtest.TestFailureException(f"Failed to compile edited shader: {newShader[1]}")

        FS2 = newShader[0]

        source = vsrefl.debugInfo.files[0].contents.replace('Position.xyz', 'Position.xyz+vec3(1.0)')

        newShader = self.controller.BuildTargetShader(
            vsrefl.entryPoint,
            rd.ShaderEncoding.GLSL,
            bytes(source, "UTF-8"),
            rd.ShaderCompileFlags(),
            rd.ShaderStage.Vertex,
        )

        if len(newShader[1]) != 0:
            raise rdtest.TestFailureException(f"Failed to compile edited shader: {newShader[1]}")

        offsetVS = newShader[0]

        source_bytes = vsrefl.rawBytes

        assert vsrefl.entryPoint == "main"

        # we search-replace in the SPIR-V expecting that 'main' won't appear anywhere other than in the OpEntryPoint
        patched_entry_source = source_bytes.replace(b'main', b't_st')

        newShader = self.controller.BuildTargetShader(
            "t_st",
            vsrefl.encoding,
            patched_entry_source,
            rd.ShaderCompileFlags(),
            rd.ShaderStage.Vertex,
        )

        if len(newShader[1]) != 0:
            raise rdtest.TestFailureException(f"Failed to compile edited shader: {newShader[1]}")

        nochangeVS = newShader[0]

        # Edit both fragment shaders
        self.replace_resource(fsrefl1.resourceId, FS1)
        self.replace_resource(fsrefl2.resourceId, FS2)

        # Refresh the replay if it didn't happen already
        self.set_event(eid, True)

        # Triangles have green and blue channel
        self.check_pixel_value(tex, 0.25, 0.5, [0.0, 1.0, 1.0, 1.0])
        self.check_pixel_value(tex, 0.75, 0.5, [0.0, 1.0, 1.0, 1.0])

        rdtest.log.success("Values are as expected after fragment editing")

        # Now "edit" the VS but don't change it. We should still get the same values
        self.replace_resource(vsrefl.resourceId, nochangeVS)
        self.set_event(eid, True)

        # Triangles have green and blue channel
        self.check_pixel_value(tex, 0.25, 0.5, [0.0, 1.0, 1.0, 1.0])
        self.check_pixel_value(tex, 0.75, 0.5, [0.0, 1.0, 1.0, 1.0])

        rdtest.log.success("Values are as expected after no-op vertex editing")

        # Change the VS to one that has offset the triangles off-centre
        self.replace_resource(vsrefl.resourceId, offsetVS)
        self.set_event(eid, True)

        # Original sample positions are now the clear color
        self.check_pixel_value(tex, 0.25, 0.5, [0.2, 0.2, 0.2, 1.0])
        self.check_pixel_value(tex, 0.75, 0.5, [0.2, 0.2, 0.2, 1.0])

        # Triangles have green and blue channel
        self.check_pixel_value(tex, 0.45, 0.95, [0.0, 1.0, 1.0, 1.0])
        self.check_pixel_value(tex, 0.95, 0.95, [0.0, 1.0, 1.0, 1.0])

        rdtest.log.success("Values are as expected after offset vertex editing")

        # Now undo the first FS edit
        self.remove_replacement(fsrefl1.resourceId)
        self.set_event(eid, True)

        # Original sample positions are still the clear color
        self.check_pixel_value(tex, 0.25, 0.5, [0.2, 0.2, 0.2, 1.0])
        self.check_pixel_value(tex, 0.75, 0.5, [0.2, 0.2, 0.2, 1.0])

        # The right triangle is the edited colour, the other two have reverted to green channel only
        self.check_pixel_value(tex, 0.45, 0.95, [0.0, 1.0, 0.0, 1.0])
        self.check_pixel_value(tex, 0.95, 0.95, [0.0, 1.0, 1.0, 1.0])

        rdtest.log.success("Values are as expected after removing first fragment edit")

        # Now undo the first VS edit
        self.remove_replacement(vsrefl.resourceId)
        self.set_event(eid, True)

        # The right triangle is the edited colour, but they are back in the original positions
        self.check_pixel_value(tex, 0.25, 0.5, [0.0, 1.0, 0.0, 1.0])
        self.check_pixel_value(tex, 0.75, 0.5, [0.0, 1.0, 1.0, 1.0])

        rdtest.log.success("Values are as expected after removing vertex edit")

        # finally undo the second FS edit
        self.remove_replacement(fsrefl2.resourceId)
        self.set_event(eid, True)

        # We should be back to where we started
        self.check_pixel_value(tex, 0.25, 0.5, [0.0, 1.0, 0.0, 1.0])
        self.check_pixel_value(tex, 0.75, 0.5, [0.0, 1.0, 0.0, 1.0])

        rdtest.log.success("Values are as expected after removing all edits")

        self.controller.FreeTargetResource(nochangeVS)
        self.controller.FreeTargetResource(offsetVS)
        self.controller.FreeTargetResource(FS1)
        self.controller.FreeTargetResource(FS2)

        bufout = self.get_resource_by_name("bufout").resourceId

        self.set_event(self.find_action("Pre-Dispatch").eventId, False)
        pipe = self.controller.GetPipelineState()
        csrefl = pipe.GetShaderReflection(rd.ShaderStage.Compute)

        uints = struct.unpack_from('=4L', self.controller.GetBufferData(bufout, 0, 0), 0)
        if not rdtest.value_compare(uints, [222, 222, 222, 222]):
            raise rdtest.TestFailureException(
                f'bufout data is incorrect before dispatch: {uints}')

        eid = self.find_action("Post-Dispatch").eventId
        self.set_event(eid, False)

        uints = struct.unpack_from('=4L', self.controller.GetBufferData(bufout, 0, 0), 0)
        if not rdtest.value_compare(uints, [777, 888, 999, 1110]):
            raise rdtest.TestFailureException(
                f'bufout data is incorrect after dispatch: {uints}')

        assert csrefl.debugInfo.encoding == rd.ShaderEncoding.HLSL
        assert csrefl.entryPoint == "hlsl_main"
        rdtest.log.success("Values are as expected before compute shader edits")

        raw_source = csrefl.rawBytes
        # search-replace in the SPIR-V for 'hlsl_main' replace with a string of the same length
        patched_entry_source = raw_source.replace(b'hlsl_main', b'main_hlsl')
        newShader = self.controller.BuildTargetShader(
            "main_hlsl",
            csrefl.encoding,
            patched_entry_source,
            rd.ShaderCompileFlags(),
            rd.ShaderStage.Compute,
        )

        if len(newShader[1]) != 0:
            raise rdtest.TestFailureException(f"Failed to compile edited compute shader: {newShader[1]}")

        nochangeCS = newShader[0]
        self.replace_resource(csrefl.resourceId, nochangeCS)
        self.set_event(eid, False)
        uints = struct.unpack_from('=4L', self.controller.GetBufferData(bufout, 0, 0), 0)
        if not rdtest.value_compare(uints, [777, 888, 999, 1110]):
            raise rdtest.TestFailureException(
                f'bufout data is incorrect after dispatch: {uints}')

        rdtest.log.success("Values are as expected after compute shader entry point change")

        glsl_source = b"""
#version 450 core
layout(push_constant) uniform PushData
{
  uvec4 data;
} push;

layout(binding = 0, std430) buffer inbuftype {
  uvec4 data[];
} inbuf;

layout(binding = 1, std430) buffer outbuftype {
  uvec4 data[];
} outbuf;

layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;

void main()
{
  outbuf.data[0].x += inbuf.data[0].x * push.data.w;
  outbuf.data[0].y += inbuf.data[0].y * push.data.z;
  outbuf.data[0].z += inbuf.data[0].z * push.data.y;
  outbuf.data[0].w += inbuf.data[0].w * push.data.x;
}
"""
        newShader = self.controller.BuildTargetShader(
            "main",
            rd.ShaderEncoding.GLSL,
            glsl_source,
            rd.ShaderCompileFlags(),
            rd.ShaderStage.Compute,
        )

        if len(newShader[1]) != 0:
            raise rdtest.TestFailureException(f"Failed to compile edited compute shader: {newShader[1]}")

        CS1 = newShader[0]
        self.replace_resource(csrefl.resourceId, CS1)
        self.set_event(eid, False)
        uints = struct.unpack_from('=4L', self.controller.GetBufferData(bufout, 0, 0), 0)
        if not rdtest.value_compare(uints, [1110, 999, 888, 777]):
            raise rdtest.TestFailureException(
                f'bufout data is incorrect after dispatch: {uints}')

        rdtest.log.success("Values are as expected after compute shader edit")

        self.controller.FreeTargetResource(CS1)
        self.controller.FreeTargetResource(nochangeCS)