/* osgEarth
 * Copyright 2025 Pelican Mapping
 * MIT License
 */
#pragma once

#include <osgEarth/Common>
#include <osgDB/Options>
#include <map>
#include <set>

namespace osgEarth
{
    class VirtualProgram;
}

namespace osgEarth { namespace Util
{

    /**
     * Functions to help load shader code.
     */
    class OSGEARTH_EXPORT ShaderPackage
    {
    public:
        /**
         * Adds a function from this package to the VirtualProgram.
         */
        bool load(
            VirtualProgram*       vp,
            const std::string&    filename,
            const osgDB::Options* dbOptions =0L ) const;

        /**
         * Removes a function loaded by load from the VirtualProgram.
         */
        bool unload(
            VirtualProgram*       vp,
            const std::string&    filename,
            const osgDB::Options* dbOptions =0L ) const;

        /**
         * Adds all the functions in this package to the VirtualProgram.
         */
        bool loadAll(
            VirtualProgram*       vp,
            const osgDB::Options* dbOptions =0L ) const;

        /**
         * Removes all the functions in this package from the VirtualProgram.
         */
        bool unloadAll(
            VirtualProgram*       vp,
            const osgDB::Options* dbOptions =0L ) const;

        /**
         * Defs or undefs a GLSL #define proprocessor macro.
         * Don't include the "#" in the defineName.
         */
        void define(
            const std::string& defineName,
            bool               defOrUndef);

        /**
         * Replaces the specified string with another string in the loaded
         * shader source. Nested replacements are not supported.
         */
        void replace(
            const std::string& pattern,
            const std::string& value);

        /**
         * Add a file/source to the package. If you only have inline source, just
         * set the filename to a unique description string.
         */
        void add(const std::string& filename, const std::string& inlineSource)
        {
            _sources[filename] = inlineSource;
        }

    public:
        using SourceMap = std::map<std::string, std::string>;
        using ReplaceMap = std::map<std::string, std::string>;
        using DefineMap = std::map<std::string, bool>;

        const SourceMap& context() const { return _sources; }


    protected:        
        SourceMap _sources;
        DefineMap _defines;
        ReplaceMap _replaces;
        friend class ShaderLoader;
    };

    /**
     * Base class for local shader file/source pairs.
     */
    class OSGEARTH_EXPORT ShaderLoader
    {
    public:
        /**
         * Loads shader source from the specified filename or string, and calls 
         * setFunction() on the virtual program to install the shader.
         * The shader much include #pragma definitions for both its
         * entry point and its location, e.g.:
         *
         *   #pragma vp_entryPoint oe_my_shader
         *   #pargma vp_location   VERTEX_VIEW
         */
        static bool load(
            VirtualProgram* vp,
            const std::string& filename_or_string,
            const ShaderPackage& package,
            const osgDB::Options* dbOptions = 0L);

        static bool load(
            VirtualProgram* vp,
            const std::string& source);

        //! Loads a shader into the virtual program on a node.
        //! (convenience function. Same as calling
        //! load(VirtualProgram::getOrCreate(node->getOrCreateStateSet()), source)
        static bool load(osg::Node* node, const std::string& source);

        /**
         * Removes a function loaded by load from the VirtualProgram.
         */
        static bool unload(
            VirtualProgram* vp,
            const std::string& filename_or_string,
            const ShaderPackage& package,
            const osgDB::Options* dbOptions =0L );

        static bool unload(
            VirtualProgram* vp,
            const std::string& source);

        //! Loads shader source from the specified filename. If the
        //! file can't be found in the OSG file path, use the source
        //! provided in backupSource.
        static std::string load(
            const std::string& filename_or_string,
            const ShaderPackage& package,
            const osgDB::Options* dbOptions = nullptr);

        /**
         * Utility function that returns the value of #pragma; e.g.,
         * #pragma vp_something You are here
         * returns "You are here" (without the quotes).
         */
        static std::string getPragmaValue(
            const std::string& source,
            const std::string& key);

        //! Find the #pragma called "key" in the string "input".
        //! Store the entire line in "line_out".
        //! Tokenize the value into tokens_out. The value may or may not
        //! be enclosed in parntheses, and its components may or may not be
        //! delimited by commas.
        static bool getPragmaValueAsTokens(
            const std::string& input,
            const std::string& key,
            std::string& line_out,
            std::vector<std::string>& tokens_out);

        static void getAllPragmaValues(
            const std::string&     source,
            const std::string&     key,
            std::set<std::string>& output);

        static void split(
            const std::string& multisource,
            std::vector<std::string>& out_sources);

        static void configureHeader(
            std::string& in_out_source);

        //! Apply the GLSL header and sort the components for use.
        //! This should be the last thing you do.
        static void finalize(
            std::string& in_out_source);

    private:

        // internal function to load a file but apply no processing to it
        static std::string load_raw_source(
            const std::string& filename_or_string,
            const ShaderPackage& package,
            const osgDB::Options* dbOptions);

        static void sort_components(
            std::string& in_out_source);

    };
} } // namespace osgEarth
