//
// Copyright Contributors to the MaterialX Project
// SPDX-License-Identifier: Apache-2.0
//

#include <MaterialXTest/External/Catch/catch.hpp>
#include <MaterialXTest/MaterialXRender/RenderUtil.h>

#include <MaterialXRenderOsl/OslRenderer.h>

#include <MaterialXRender/StbImageLoader.h>
#if defined(MATERIALX_BUILD_OIIO)
#include <MaterialXRender/OiioImageLoader.h>
#endif

#include <MaterialXGenOsl/OslShaderGenerator.h>
#include <MaterialXGenOsl/OslNetworkShaderGenerator.h>

#include <MaterialXFormat/Util.h>

namespace mx = MaterialX;

namespace
{

//
// Define local overrides for the tangent frame in shader generation, aligning conventions
// between MaterialXRender and testrender.
//

class TangentOsl : public mx::ShaderNodeImpl
{
  public:
    static mx::ShaderNodeImplPtr create()
    {
        return std::make_shared<TangentOsl>();
    }

    void emitFunctionCall(const  mx::ShaderNode& node, mx::GenContext& context, mx::ShaderStage& stage) const override
    {
        const mx::ShaderGenerator& shadergen = context.getShaderGenerator();

        DEFINE_SHADER_STAGE(stage, mx::Stage::PIXEL)
        {
            shadergen.emitLineBegin(stage);
            shadergen.emitOutput(node.getOutput(), true, false, context, stage);
            shadergen.emitString(" = normalize(vector(N[2], 0, -N[0]))", stage);
            shadergen.emitLineEnd(stage);
        }
    }
};

class BitangentOsl : public mx::ShaderNodeImpl
{
  public:
    static mx::ShaderNodeImplPtr create()
    {
        return std::make_shared<BitangentOsl>();
    }

    void emitFunctionCall(const  mx::ShaderNode& node, mx::GenContext& context, mx::ShaderStage& stage) const override
    {
        const mx::ShaderGenerator& shadergen = context.getShaderGenerator();

        DEFINE_SHADER_STAGE(stage, mx::Stage::PIXEL)
        {
            shadergen.emitLineBegin(stage);
            shadergen.emitOutput(node.getOutput(), true, false, context, stage);
            shadergen.emitString(" = normalize(cross(N, vector(N[2], 0, -N[0])))", stage);
            shadergen.emitLineEnd(stage);
        }
    }
};

} // anonymous namespace

class OslShaderRenderTester : public RenderUtil::ShaderRenderTester
{
  public:
    explicit OslShaderRenderTester(mx::ShaderGeneratorPtr shaderGenerator, bool useOslCmdStr) :
        RenderUtil::ShaderRenderTester(shaderGenerator)
    {
        // Preprocess to resolve to absolute image file names
        // and all non-POSIX separators must be converted to POSIX ones (this only affects running on Windows)
        _resolveImageFilenames = true;
        _customFilenameResolver = mx::StringResolver::create();
        _customFilenameResolver->setFilenameSubstitution("\\\\", "/");
        _customFilenameResolver->setFilenameSubstitution("\\", "/");
        _useOslCmdStr = useOslCmdStr;
    }

  protected:
    void createRenderer(std::ostream& log) override;

    RenderUtil::RenderProfileResult runRenderer(
        const RenderUtil::RenderSession& session,
        const RenderUtil::RenderItem& item,
        mx::GenContext& context) override;

    void addSkipFiles() override
    {
        if (_useOslCmdStr)
        {
            _skipFiles.insert("filename_cm_test.mtlx");
            _skipFiles.insert("shader_ops.mtlx");
            _skipFiles.insert("chiang_hair_surfaceshader.mtlx");
            _skipFiles.insert("network_surfaceshader.mtlx");
            _skipFiles.insert("sheen.mtlx");
            _skipFiles.insert("toon_shade.mtlx");
            _skipFiles.insert("bsdf_graph.mtlx");
            _skipFiles.insert("varying_ior.mtlx");
            _skipFiles.insert("vertical_layering.mtlx");
            _skipFiles.insert("mix_bsdf.mtlx");
        }
    }

    bool saveImage(const mx::FilePath& filePath, mx::ConstImagePtr image, bool /*verticalFlip*/) const override
    {
        return _renderer->getImageHandler()->saveImage(filePath, image, false);
    }

    mx::ImageLoaderPtr _imageLoader;
    mx::OslRendererPtr _renderer;
    bool _useOslCmdStr;
};

// Renderer setup
void OslShaderRenderTester::createRenderer(std::ostream& log)
{
    bool initialized = false;

    _renderer = mx::OslRenderer::create();
    _imageLoader = mx::StbImageLoader::create();

    // Set up additional utilities required to run OSL testing including
    // oslc and testrender paths and OSL include path
    //
    const std::string oslcExecutable(MATERIALX_OSL_BINARY_OSLC);
    _renderer->setOslCompilerExecutable(oslcExecutable);
    const std::string testRenderExecutable(MATERIALX_OSL_BINARY_TESTRENDER);
    _renderer->setOslTestRenderExecutable(testRenderExecutable);
    mx::FilePath oslStandardIncludePath = mx::FilePath(MATERIALX_OSL_INCLUDE_PATH);
    if (!oslStandardIncludePath.isEmpty())
    {
        _renderer->setOslIncludePath(mx::FileSearchPath(oslStandardIncludePath));
    }

    try
    {
        _renderer->initialize();

        mx::StbImageLoaderPtr stbLoader = mx::StbImageLoader::create();
        mx::ImageHandlerPtr imageHandler = mx::ImageHandler::create(stbLoader);
        imageHandler->setSearchPath(mx::getDefaultDataSearchPath());
#if defined(MATERIALX_BUILD_OIIO)
        mx::OiioImageLoaderPtr oiioLoader = mx::OiioImageLoader::create();
        imageHandler->addLoader(oiioLoader);
#endif
        _renderer->setImageHandler(imageHandler);
        _renderer->setLightHandler(nullptr);
        initialized = true;

        // Pre-compile some required shaders for testrender
        if (!oslcExecutable.empty() && !testRenderExecutable.empty())
        {
            mx::FileSearchPath searchPath = mx::getDefaultDataSearchPath();
            mx::FilePath shaderPath = searchPath.find("resources/Utilities/");
            _renderer->setOslOutputFilePath(shaderPath);

            const std::string OSL_EXTENSION("osl");
            for (const mx::FilePath& filename : shaderPath.getFilesInDirectory(OSL_EXTENSION))
            {
                _renderer->compileOSL(shaderPath / filename);
            }

            // Set the search path for these compiled shaders.
            _renderer->setOslUtilityOSOPath(shaderPath);
        }
    }
    catch (mx::ExceptionRenderError& e)
    {
        for (const auto& error : e.errorLog())
        {
            log << e.what() << " " << error << std::endl;
        }
    }
    REQUIRE(initialized);
}

// Renderer execution
RenderUtil::RenderProfileResult OslShaderRenderTester::runRenderer(
    const RenderUtil::RenderSession& session,
    const RenderUtil::RenderItem& item,
    mx::GenContext& context)
{
    RenderUtil::RenderProfileResult result;
    const std::string& shaderName = item.shaderName;
    mx::DocumentPtr doc = item.document;
    mx::TypedElementPtr element = item.element;
    const GenShaderUtil::TestSuiteOptions& testOptions = session.testOptions;
    std::ostream& log = session.log;

    std::cout << "Validating OSL rendering for: " << doc->getSourceUri() << std::endl;

    mx::ScopedTimer totalOSLTime(&result.languageTimes.totalTime);

    mx::ShaderGenerator& shadergen = context.getShaderGenerator();

    std::vector<mx::GenOptions> optionsList;
    getGenerationOptions(testOptions, context.getOptions(), optionsList);

    if (element && doc)
    {
        log << "------------ Run OSL validation with element: " << element->getNamePath() << "-------------------" << std::endl;

        for (const auto& options : optionsList)
        {
            result.elementsTested++;

            mx::ShaderPtr shader;
            try
            {
                mx::ScopedTimer genTimer(&result.languageTimes.generationTime);
                mx::GenOptions& contextOptions = context.getOptions();
                contextOptions = options;
                contextOptions.targetColorSpaceOverride = "lin_rec709";
                contextOptions.oslConnectCiWrapper = true;

                // Apply local overrides for shader generation.
                shadergen.registerImplementation("IM_tangent_vector3_" + mx::OslShaderGenerator::TARGET, TangentOsl::create);
                shadergen.registerImplementation("IM_bitangent_vector3_" + mx::OslShaderGenerator::TARGET, BitangentOsl::create);

                shader = shadergen.generate(shaderName, element, context);
            }
            catch (mx::Exception& e)
            {
                log << ">> " << e.what() << "\n";
                shader = nullptr;
            }
            CHECK(shader != nullptr);
            if (shader == nullptr)
            {
                log << ">> Failed to generate shader\n";
                result.success = false;
                return result;
            }
            CHECK(shader->getSourceCode().length() > 0);

            std::string oslCmdStr = shader->getSourceCode();

            std::string shaderPath;
            mx::FilePath outputFilePath = item.outputPath;
            
            // Note: mkdir will fail if the directory already exists which is ok.
            {
                mx::ScopedTimer ioDir(&result.languageTimes.ioTime);
                outputFilePath.createDirectory(true);
                
                // Use separate directory for reduced output
                if (options.shaderInterfaceType == mx::SHADER_INTERFACE_REDUCED)
                {
                    outputFilePath = outputFilePath / mx::FilePath("reduced");
                    outputFilePath.createDirectory();
                }
            }

            shaderPath = mx::FilePath(outputFilePath) / mx::FilePath(shaderName);

            // Write out osl file
            if (testOptions.dumpGeneratedCode)
            {
                mx::ScopedTimer ioTimer(&result.languageTimes.ioTime);
                std::ofstream file;
                file.open(shaderPath + (_useOslCmdStr ? ".oslcmd" : ".osl"));
                file << shader->getSourceCode();
                file.close();
            }

            // Validate
            bool validated = false;
            try
            {
                // Set renderer properties.
                _renderer->setOslOutputFilePath(outputFilePath);
                _renderer->setOslShaderName(shaderName);
                // The antialiasing flag in testrender traces N*N samples per pixel,
                // so these values are effectively squared.
                _renderer->setAaLit(testOptions.enableReferenceQuality ? 16 : 2);
                _renderer->setAaUnlit(testOptions.enableReferenceQuality ? 2 : 1);

                // Validate compilation
                if (!_useOslCmdStr)
                {
                    mx::ScopedTimer compileTimer(&result.languageTimes.compileTime);
                    _renderer->createProgram(shader);
                }

                _renderer->setSize(static_cast<unsigned int>(testOptions.renderSize[0]), static_cast<unsigned int>(testOptions.renderSize[1]));

                const mx::ShaderStage& stage = shader->getStage(mx::Stage::PIXEL);

                // Bind IBL image name overrides.
                mx::StringVec envOverrides;
                std::string envmap_filename("string envmap_filename \"");
                envmap_filename += testOptions.radianceIBLPath;
                envmap_filename += "\";\n";
                envOverrides.push_back(envmap_filename);

                _renderer->setEnvShaderParameterOverrides(envOverrides);

                const mx::VariableBlock& outputs = stage.getOutputBlock(mx::OSL::OUTPUTS);
                if (outputs.size() > 0)
                {
                    const mx::ShaderPort* output = outputs[0];
                    const mx::TypeSyntax& typeSyntax = shadergen.getSyntax().getTypeSyntax(output->getType());

                    const std::string& outputName = output->getVariable();
                    const std::string& outputType = typeSyntax.getTypeAlias().empty() ? typeSyntax.getName() : typeSyntax.getTypeAlias();
                    const std::string& sceneTemplateFile = (_useOslCmdStr ? "scene_template_oslcmd.xml" : "scene_template.xml");

                    // Set shader output name and type to use
                    _renderer->setOslShaderOutput(outputName, outputType);

                    // Get oslcmdstr
                    _renderer->setOSLCmdStr(oslCmdStr);
                    _renderer->useOslCommandString(_useOslCmdStr);
                    // Osl Network Shaders record the oso path in an attribute
                    auto osoPathAttr = shader->getAttribute("osoPath");
                    if (osoPathAttr)
                    {
                        _renderer->setDataLibraryOSOPath(osoPathAttr->getValueString());
                    }

                    // Set scene template file. For now we only have the constant color scene file
                    mx::FileSearchPath searchPath = mx::getDefaultDataSearchPath();
                    mx::FilePath sceneTemplatePath = searchPath.find("resources/Utilities/");
                    sceneTemplatePath = sceneTemplatePath / sceneTemplateFile;
                    _renderer->setOslTestRenderSceneTemplateFile(sceneTemplatePath.asString());

                    // Validate rendering
                    {
                        mx::ScopedTimer renderTimer(&result.languageTimes.renderTime);
                        _renderer->render();
                        result.images.push_back(_renderer->captureImage());
                    }
                }
                else
                {
                    CHECK(false);
                    log << ">> Shader has no output to render from\n";
                }

                validated = true;
            }
            catch (mx::ExceptionRenderError& e)
            {
                // Always dump shader on error
                std::ofstream file;
                file.open(shaderPath + ".osl");
                file << shader->getSourceCode();
                file.close();

                for (const auto& error : e.errorLog())
                {
                    log << e.what() << " " << error << std::endl;
                }
                log << ">> Refer to shader code in dump file: " << shaderPath << ".osl file" << std::endl;
            }
            catch (mx::Exception& e)
            {
                log << e.what();
            }
            CHECK(validated);
            if (!validated)
                result.success = false;
        }
    }

    return result;
}

TEST_CASE("Render: OSL TestSuite", "[renderosl]")
{
    if (std::string(MATERIALX_OSL_BINARY_OSLC).empty() &&
        std::string(MATERIALX_OSL_BINARY_TESTRENDER).empty())
    {
        INFO("Skipping the OSL test suite as its executable locations haven't been set.");
        return;
    }

    mx::FileSearchPath searchPath = mx::getDefaultDataSearchPath();
    mx::FilePath optionsFilePath = searchPath.find("resources/Materials/TestSuite/_options.mtlx");

    OslShaderRenderTester renderTester(mx::OslShaderGenerator::create(), false);
    renderTester.validate(optionsFilePath);
}

#ifdef MATERIALX_BUILD_OSOS
TEST_CASE("Render: OSL Network TestSuite", "[renderoslnetwork]")
{
    if (std::string(MATERIALX_OSL_BINARY_OSLC).empty() &&
        std::string(MATERIALX_OSL_BINARY_TESTRENDER).empty())
    {
        INFO("Skipping the OSL test suite as its executable locations haven't been set.");
        return;
    }

    mx::FileSearchPath searchPath = mx::getDefaultDataSearchPath();
    mx::FilePath optionsFilePath = searchPath.find("resources/Materials/TestSuite/_options.mtlx");

    OslShaderRenderTester renderTester(mx::OslNetworkShaderGenerator::create(), true);
    renderTester.validate(optionsFilePath);
}
#endif
