import renderdoc as rd
import rdtest


# Not a direct test, re-used by API-specific tests
class Buffer_Truncation(rdtest.TestCase):
    internal = True
    draw_action = None

    def check_capture(self):
        action = self.draw_action
        if action is None:
            action = self.find_action("Draw")

        assert action is not None

        self.set_event(action.eventId, False)

        vsin_ref: rdtest.MeshReference = {
            0: {
                'vtx': 0,
                'idx': 1,
                'POSITION': [-0.5, -0.5, 0.0],
                'COLOR': [0.0, 1.0, 0.0, 1.0],
            },
            1: {
                'vtx': 1,
                'idx': 2,
                'POSITION': [0.0, 0.5, 0.0],
                'COLOR': [0.0, 1.0, 0.0, 1.0],
            },
            2: {
                'vtx': 2,
                'idx': 3,
                'POSITION': [0.5, -0.5, 0.0],
                'COLOR': [0.0, 1.0, 0.0, 1.0],
            },
            3: {
                'vtx': 3,
                'idx': 4,
                'POSITION': [8.8, 0.0, 0.0],
                'COLOR': [0.0, 0.0, 0.0, 1.0],
            },
            4: {
                'vtx': 4,
                'idx': 5,
                'POSITION': None,
                'COLOR': None,
            },
            5: {
                'vtx': 5,
                'idx': None,
                'POSITION': None,
                'COLOR': None,
            },
        }

        self.check_mesh_data(vsin_ref, self.get_vsin(action))

        postvs_data = self.get_postvs(action, rd.MeshDataStage.VSOut, 0, action.numIndices)

        postvs_ref: rdtest.MeshReference = {
            0: {
                'vtx': 0,
                'idx': 1,
                'OUTPOSITION': [-0.5, -0.5, 0.0, 1.0],
                'OUTCOLOR': [0.0, 1.0, 0.0, 1.0],
            },
            1: {
                'vtx': 1,
                'idx': 2,
                'OUTPOSITION': [0.0, 0.5, 0.0, 1.0],
                'OUTCOLOR': [0.0, 1.0, 0.0, 1.0],
            },
            2: {
                'vtx': 2,
                'idx': 3,
                'OUTPOSITION': [0.5, -0.5, 0.0, 1.0],
                'OUTCOLOR': [0.0, 1.0, 0.0, 1.0],
            },
            3: {
                'vtx': 3,
                'idx': 4,
                'OUTPOSITION': [8.8, 0.0, 0.0, 1.0],
                'OUTCOLOR': [0.0, 0.0, 0.0, 1.0],
            },
            4: {
                'vtx': 4,
                'idx': 5,
                # don't rely on any particular OOB behaviour for postvs, as this may come from the driver/API
            },
            5: {
                'vtx': 5,
                'idx': None,
                # don't rely on any particular OOB behaviour for postvs, as this may come from the driver/API
            },
        }

        self.check_mesh_data(postvs_ref, postvs_data)

        rdtest.log.success("vertex/index buffers were truncated as expected")

        pipe = self.controller.GetPipelineState()

        stage = rd.ShaderStage.Pixel

        cbuf = pipe.GetConstantBlock(stage, 0, 0).descriptor

        if self.find_action('NoCBufferRange') == None:
            assert cbuf.byteSize == 256

        variables = self.controller.GetCBufferVariableContents(pipe.GetGraphicsPipelineObject(),
                                                               pipe.GetShader(stage), stage,
                                                               pipe.GetShaderEntryPoint(stage), 0,
                                                               cbuf.resource, cbuf.byteOffset, cbuf.byteSize)

        outcol = variables[1]

        assert outcol.name == "outcol"
        if not rdtest.value_compare(outcol.value.f32v[0:4], [0.0, 0.0, 0.0, 0.0]):
            raise rdtest.TestFailureException(f"expected outcol to be 0s, but got {outcol.value.f32v[0:4]}")

        x, y = self.get_view_centre()

        # Debug the shader
        with self.debug_pixel(x, y, rd.DebugPixelInputs()) as debug:
            cycles, variables = self.process_trace(debug.trace)

            cbuf_sourceVars = [s for s in debug.trace.sourceVars if s.variables[0].type == rd.DebugVariableType.Constant and s.rows > 0]

            # Vulkan style, one source var for the cbuffer
            if len(cbuf_sourceVars) == 1:
                debugged_cb = debug.trace.constantBlocks[0]

                assert debugged_cb.members[0].name == 'padding'
                assert debugged_cb.members[1].name == 'outcol'

                if not rdtest.value_compare(debugged_cb.members[1].value.f32v[0:4], [0.0, 0.0, 0.0, 0.0]):
                    raise rdtest.TestFailureException(f"expected outcol to be 0s, but got {debugged_cb.members[1].value.f32v[0:4]}")
            # D3D style, one source var for each member mapping to a register
            elif len(cbuf_sourceVars) == 17:
                debugged_cb = debug.trace.constantBlocks[0].members[16]

                assert all(['consts.padding[' in c.name for c in cbuf_sourceVars[0:16]])
                assert cbuf_sourceVars[16].name == 'consts.outcol'

                assert cbuf_sourceVars[16].variables[0].name == 'cb0[16]' or cbuf_sourceVars[16].variables[0].name == 'consts[16]'

                if not rdtest.value_compare(debugged_cb.value.f32v[0:4], [0.0, 0.0, 0.0, 0.0]):
                    raise rdtest.TestFailureException(f"expected outcol to be 0s, but got {debugged_cb.value.f32v[0:4]}")
            else:
                raise rdtest.TestFailureException(f"Unexpected number of constant buffer source vars {len(cbuf_sourceVars)}")

            rdtest.log.success("CBuffer value was truncated as expected")