// main.cpp

// This file provides the application code for the `shader-toy` example.
//
// Much of the logic here is identical to the simpler `hello-world` example,
// so we will not spend time commenting those parts that are identical or
// nearly identical. Readers who want detailed comments on a simpler example
// using Slang should look there.

// This example uses the Slang C/C++ API, alonmg with its optional type
// for managing COM-style reference-counted pointers.
//
#include "slang-com-ptr.h"
#include "slang.h"
using Slang::ComPtr;

// This example uses a graphics API abstraction layer that is implemented inside
// the Slang codebase for use in our sample programs and test cases. Use of
// this layer is *not* required or assumed when using the Slang language,
// compiler, and API.
//
#include "core/slang-basic.h"
#include "examples/example-base/example-base.h"
#include "platform/performance-counter.h"
#include "platform/window.h"
#include "slang-rhi.h"
#include "slang-rhi/shader-cursor.h"

#include <chrono>

static const ExampleResources resourceBase("shader-toy");

using namespace rhi;

// In order to display a shader toy effect using rasterization-based shader
// execution we need to render a full-screen triangle. We will define a
// small helper type that defines the data for such a triangle.
//
struct FullScreenTriangle
{
    struct Vertex
    {
        float position[2];
    };

    enum
    {
        kVertexCount = 3
    };

    static const Vertex kVertices[kVertexCount];
};
const FullScreenTriangle::Vertex FullScreenTriangle::kVertices[FullScreenTriangle::kVertexCount] = {
    {{-1, -1}},
    {{-1, 3}},
    {{3, -1}},
};

// The application itself will be encapsulated in a C++ `struct` type
// so that it can easily scope its state without use of global variables.
//
struct ShaderToyApp : public WindowedAppBase
{

    // The uniform data used by the shader is defined here as a simple
    // POD ("plain old data") type.
    //
    // Note: This type must match the declaration of `ShaderToyUniforms`
    // in the file `shader-toy.slang`.
    //
    // An application could instead use a shared header file to define
    // this type, or use Slang's reflection capabilities to allocate
    // and set parameters at runtime. For this simple example we did
    // the expedient thing of having distinct Slang and C++ declarations.
    //
    struct Uniforms
    {
        float iMouse[4];
        float iResolution[2];
        float iTime;
    };

    // The main interesting part of the host application code is where we
    // load, compile, inspect, and compose the Slang shader code.
    //
    Result loadShaderProgram(IDevice* device, ComPtr<IShaderProgram>& outShaderProgram)
    {
        // We need to obatin a compilation session (`slang::ISession`) that will provide
        // a scope to all the compilation and loading of code we do.
        //
        // Our example application uses the `slang-rhi` graphics API abstraction layer, which
        // already creates a Slang compilation session for us, so we just grab and use it here.
        ComPtr<slang::ISession> slangSession;
        slangSession = device->getSlangSession();

        // Once the session has been obtained, we can start loading code into it.
        //
        // The simplest way to load code is by calling `loadModule` with the name of a Slang
        // module. A call to `loadModule("MyStuff")` will behave more or less as if you
        // wrote:
        //
        //      import MyStuff;
        //
        // In a Slang shader file. The compiler will use its search paths to try to locate
        // `MyModule.slang`, then compile and load that file. If a matching module had
        // already been loaded previously, that would be used directly.
        //
        // Note: The only interesting wrinkle here is that our file is named `shader-toy` with
        // a hyphen in it, so the name is not directly usable as an identifier in Slang code.
        // Instead, when trying to import this module in the context of Slang code, a user
        // needs to replace the hyphens with underscores:
        //
        //      import shader_toy;
        //
        ComPtr<slang::IBlob> diagnosticsBlob;
        Slang::String shaderToyPath = resourceBase.resolveResource("shader-toy.slang");
        slang::IModule* module =
            slangSession->loadModule(shaderToyPath.getBuffer(), diagnosticsBlob.writeRef());
        diagnoseIfNeeded(diagnosticsBlob);
        if (!module)
            return SLANG_FAIL;

        // Loading the `shader-toy` module will compile and check all the shader code in it,
        // including the shader entry points we want to use. Now that the module is loaded
        // we can look up those entry points by name.
        //
        // Note: If you are using this `loadModule` approach to load your shader code it is
        // important to tag your entry point functions with the `[shader("...")]` attribute
        // (e.g., `[shader("vertex")] void vertexMain(...)`). Without that information there
        // is no umambiguous way for the compiler to know which functions represent entry
        // points when it parses your code via `loadModule()`.
        //
        char const* vertexEntryPointName = "vertexMain";
        char const* fragmentEntryPointName = "fragmentMain";
        //
        ComPtr<slang::IEntryPoint> vertexEntryPoint;
        SLANG_RETURN_ON_FAIL(
            module->findEntryPointByName(vertexEntryPointName, vertexEntryPoint.writeRef()));
        //
        ComPtr<slang::IEntryPoint> fragmentEntryPoint;
        SLANG_RETURN_ON_FAIL(
            module->findEntryPointByName(fragmentEntryPointName, fragmentEntryPoint.writeRef()));

        // At this point we have a few different Slang API objects that represent
        // pieces of our code: `module`, `vertexEntryPoint`, and `fragmentEntryPoint`.
        //
        // A single Slang module could contain many different entry points (e.g.,
        // four vertex entry points, three fragment entry points, and two compute
        // shaders), and before we try to generate output code for our target API
        // we need to identify which entry points we plan to use together.
        //
        // Modules and entry points are both examples of *component types* in the
        // Slang API. The API also provides a way to build a *composite* out of
        // other pieces, and that is what we are going to do with our module
        // and entry points.
        //
        Slang::List<slang::IComponentType*> componentTypes;
        componentTypes.add(module);

        // Later on when we go to extract compiled kernel code for our vertex
        // and fragment shaders, we will need to make use of their order within
        // the composition, so we will record the relative ordering of the entry
        // points here as we add them.
        int entryPointCount = 0;
        int vertexEntryPointIndex = entryPointCount++;
        componentTypes.add(vertexEntryPoint);

        int fragmentEntryPointIndex = entryPointCount++;
        componentTypes.add(fragmentEntryPoint);

        // Actually creating the composite component type is a single operation
        // on the Slang session, but the operation could potentially fail if
        // something about the composite was invalid (e.g., you are trying to
        // combine multiple copies of the same module), so we need to deal
        // with the possibility of diagnostic output.
        //
        ComPtr<slang::IComponentType> composedProgram;
        SlangResult result = slangSession->createCompositeComponentType(
            componentTypes.getBuffer(),
            componentTypes.getCount(),
            composedProgram.writeRef(),
            diagnosticsBlob.writeRef());
        diagnoseIfNeeded(diagnosticsBlob);
        SLANG_RETURN_ON_FAIL(result);

        // At this point, `composedProgram` represents the shader program
        // we want to run, and the vertex and fragment shader there have
        // been checked.
        //
        // We could use the Slang reflection API on `composedProgram` at this
        // point to query things like the locations and offsets of the
        // various uniform parameters, textures, etc.
        //
        // What *cannot* be done yet at this point is actually generating
        // kernel code, because `composedProgram` includes a generic type
        // parameter as part of the `fragmentMain` entry point:
        //
        //      void fragmentMain<T : IShaderToyImageShader>(...)
        //
        // Our next task is to load code for a type we'd like to plug in
        // for `T` there.
        //
        // Because Slang supports modular programming, there is no requirement
        // that a type we want to plug in for `T` has to come from the
        // same module, and to demonstrate that we will load a different
        // module to provide the effect type we will plug in.
        //
        const char* effectTypeName = "ExampleEffect";
        Slang::String effectModulePath = resourceBase.resolveResource("example-effect.slang");
        slang::IModule* effectModule =
            slangSession->loadModule(effectModulePath.getBuffer(), diagnosticsBlob.writeRef());
        diagnoseIfNeeded(diagnosticsBlob);
        if (!module)
            return SLANG_FAIL;

        // Once we've loaded the code module that defines out effect type,
        // we can look it up by name using the reflection information on
        // the module.
        //
        // Note: A future version of the Slang API will support enumerating
        // the types declared in a module so that we do not have to hard-code
        // the name here.
        //
        auto effectType = effectModule->getLayout()->findTypeByName(effectTypeName);

        // Now that we have the `effectType` we want to plug in to our generic
        // shader, we need to specialize the shader to that type.
        //
        // Because a shader program could have zero or more specialization parameters,
        // we need to build up an array of specialization arguments.
        //
        Slang::List<slang::SpecializationArg> specializationArgs;

        {
            // In our case, we only have a single specialization argument we plan
            // to use, and it is a type argument.
            //
            slang::SpecializationArg effectTypeArg;
            effectTypeArg.kind = slang::SpecializationArg::Kind::Type;
            effectTypeArg.type = effectType;
            specializationArgs.add(effectTypeArg);
        }

        // Specialization of a component type is a single Slang API call, but
        // we need to deal with the possibility of diagnostic output on failure.
        // For example, if we tried to specialize the shader program to a
        // type like `int` that doesn't support the `IShaderToyImageShader` interface,
        // this is the step where we'd get an error message saying so.
        //
        ComPtr<slang::IComponentType> specializedProgram;
        result = composedProgram->specialize(
            specializationArgs.getBuffer(),
            specializationArgs.getCount(),
            specializedProgram.writeRef(),
            diagnosticsBlob.writeRef());
        diagnoseIfNeeded(diagnosticsBlob);
        SLANG_RETURN_ON_FAIL(result);

        // At this point we have a specialized shader program that represents our
        // intention to run the `vertexMain` and `fragmentMain` entry points,
        // specialized to the `ExampleEffect` type we loaded.
        //
        // We can now *link* the program, which ensures that all of the code that
        // it transitively depends on has been pulled together into a single
        // component type.
        //
        ComPtr<slang::IComponentType> linkedProgram;
        result = specializedProgram->link(linkedProgram.writeRef(), diagnosticsBlob.writeRef());
        diagnoseIfNeeded(diagnosticsBlob);
        SLANG_RETURN_ON_FAIL(result);

        ShaderProgramDesc programDesc = {};
        programDesc.slangGlobalScope = linkedProgram.get();
        auto shaderProgram = device->createShaderProgram(programDesc);
        outShaderProgram = shaderProgram;
        return SLANG_OK;
    }

    ComPtr<IShaderProgram> gShaderProgram;
    ComPtr<IPipeline> gPipeline;
    ComPtr<IBuffer> gVertexBuffer;
    const Format format = Format::RG32Float;

    Result initialize()
    {
        SLANG_RETURN_ON_FAIL(initializeBase("Shader Toy", 1024, 768, getDeviceType()));

        // We may not have a window if we're running in test mode
        SLANG_ASSERT(isTestMode() || gWindow);
        if (gWindow)
        {
            gWindow->events.mouseMove = [this](const platform::MouseEventArgs& e)
            { handleEvent(e); };
            gWindow->events.mouseUp = [this](const platform::MouseEventArgs& e) { handleEvent(e); };
            gWindow->events.mouseDown = [this](const platform::MouseEventArgs& e)
            { handleEvent(e); };
        }

        InputElementDesc inputElements[] = {
            {"POSITION", 0, format, offsetof(FullScreenTriangle::Vertex, position)},
        };
        auto inputLayout = gDevice->createInputLayout(
            sizeof(FullScreenTriangle::Vertex),
            &inputElements[0],
            SLANG_COUNT_OF(inputElements));
        if (!inputLayout)
            return SLANG_FAIL;

        BufferDesc vertexBufferDesc;
        vertexBufferDesc.size =
            FullScreenTriangle::kVertexCount * sizeof(FullScreenTriangle::Vertex);
        vertexBufferDesc.elementSize = sizeof(FullScreenTriangle::Vertex);
        vertexBufferDesc.usage = BufferUsage::VertexBuffer;
        gVertexBuffer = gDevice->createBuffer(vertexBufferDesc, &FullScreenTriangle::kVertices[0]);
        if (!gVertexBuffer)
            return SLANG_FAIL;

        SLANG_RETURN_ON_FAIL(loadShaderProgram(gDevice, gShaderProgram));

        // Create pipeline.
        ColorTargetDesc colorTarget;
        colorTarget.format = Format::RGBA8Unorm;
        RenderPipelineDesc desc;
        desc.inputLayout = inputLayout;
        desc.program = gShaderProgram;
        desc.targetCount = 1;
        desc.targets = &colorTarget;
        desc.depthStencil.depthTestEnable = false;
        desc.depthStencil.depthWriteEnable = false;
        desc.primitiveTopology = PrimitiveTopology::TriangleList;
        gPipeline = gDevice->createRenderPipeline(desc);
        if (!gPipeline)
            return SLANG_FAIL;

        return SLANG_OK;
    }

    bool wasMouseDown = false;
    bool isMouseDown = false;
    float lastMouseX = 0.0f;
    float lastMouseY = 0.0f;
    float clickMouseX = 0.0f;
    float clickMouseY = 0.0f;

    bool firstTime = true;
    platform::TimePoint startTime;

    virtual void renderFrame(ITexture* texture) override
    {
        auto commandEncoder = gQueue->createCommandEncoder();
        if (firstTime)
        {
            startTime = platform::PerformanceCounter::now();
            firstTime = false;
        }

        // Update uniform buffer.

        Uniforms uniforms = {};
        {
            bool isMouseClick = isMouseDown && !wasMouseDown;
            wasMouseDown = isMouseDown;

            if (isMouseClick)
            {
                clickMouseX = lastMouseX;
                clickMouseY = lastMouseY;
            }

            uniforms.iMouse[0] = lastMouseX;
            uniforms.iMouse[1] = lastMouseY;
            uniforms.iMouse[2] = isMouseDown ? clickMouseX : -clickMouseX;
            uniforms.iMouse[3] = isMouseClick ? clickMouseY : -clickMouseY;
            uniforms.iTime = platform::PerformanceCounter::getElapsedTimeInSeconds(startTime);
            uniforms.iResolution[0] = float(windowWidth);
            uniforms.iResolution[1] = float(windowHeight);
        }

        // Encode render commands.
        ComPtr<ITextureView> textureView = gDevice->createTextureView(texture, {});
        RenderPassColorAttachment colorAttachment = {};
        colorAttachment.view = textureView;
        colorAttachment.loadOp = LoadOp::Clear;

        RenderPassDesc renderPass = {};
        renderPass.colorAttachments = &colorAttachment;
        renderPass.colorAttachmentCount = 1;

        auto encoder = commandEncoder->beginRenderPass(renderPass);

        RenderState renderState = {};
        renderState.viewports[0] = Viewport::fromSize(windowWidth, windowHeight);
        renderState.viewportCount = 1;
        renderState.scissorRects[0] = ScissorRect::fromSize(windowWidth, windowHeight);
        renderState.scissorRectCount = 1;

        auto rootObject = encoder->bindPipeline(static_cast<IRenderPipeline*>(gPipeline.get()));
        auto constantBuffer = rootObject->getObject(ShaderOffset());
        constantBuffer->setData(ShaderOffset(), &uniforms, sizeof(uniforms));

        renderState.vertexBuffers[0] = gVertexBuffer;
        renderState.vertexBufferCount = 1;
        encoder->setRenderState(renderState);

        DrawArguments drawArgs = {};
        drawArgs.vertexCount = 3;
        encoder->draw(drawArgs);

        encoder->end();

        gQueue->submit(commandEncoder->finish());

        if (!isTestMode())
        {
            gSurface->present();
        }
    }

    void handleEvent(const platform::MouseEventArgs& event)
    {
        isMouseDown = ((int)event.buttons & (int)platform::ButtonState::Enum::LeftButton) != 0;
        lastMouseX = (float)event.x;
        lastMouseY = (float)event.y;
    }
};

// This macro instantiates an appropriate main function to
// run the application defined above.
EXAMPLE_MAIN(innerMain<ShaderToyApp>);