/* osgEarth
 * Copyright 2025 Pelican Mapping
 * MIT License
 */
#ifndef OSGEARTH_COVERAGE_LAYER_H
#define OSGEARTH_COVERAGE_LAYER_H

#include <osgEarth/Common>
#include <osgEarth/Layer>
#include <osgEarth/ImageLayer>
#include <osgEarth/LayerReference>
#include <osgEarth/Metrics>
#include <osgEarth/Progress>
#include <osgEarth/Coverage>
#include <osgEarth/Cache>

namespace osgEarth
{
    /**
     * CoverageLayer is a data-only layer that creates gridded coverages.
     * A "coverage" is a 2-dimensional data structure, like and image,
     * but in which each cell contains a user-defined data structure.
     *
     * Internally, a Coverage is an Image consisting of a discrete
     * integer value at each pixel. A mapping tables provides information
     * on how to interpret specific values. The user of a CoverageLayer
     * must know the format of the data structure in use, and must
     * use that structure as the template parameter when initializing a CoverageCreator
     * and calling createMappings() and createCoverage().
     *
     * This is a data-only layer. The terrain engine will not render it.
     */
    class OSGEARTH_EXPORT CoverageLayer : public TileLayer
    {
    public:
        struct SourceLayerOptions : public ConfigOptions
        {
            OE_OPTION_LAYER(ImageLayer, source);
            OE_OPTION(Config, mappings);
            Config getConfig() const;
            void fromConfig(const Config& conf);
        };

        struct OSGEARTH_EXPORT Options : public TileLayer::Options
        {
            META_LayerOptions(osgEarth, Options, TileLayer::Options);
            OE_OPTION(Config, presets);
            OE_OPTION_VECTOR(SourceLayerOptions, layers);
            virtual Config getConfig() const;
        private:
            void fromConfig(const Config& conf);
        };

    public:
        META_Layer(osgEarth, CoverageLayer, Options, TileLayer, Coverage);

        //! Table that holds the mappings from coverage values
        //! to user data structures.
        template<typename T>
        struct Mappings : public std::unordered_map<unsigned, T> {
            T _nodata;
            inline const T& get(unsigned value) const;
        };

        //! Table of preset sample values.
        struct Presets : public std::unordered_map<std::string, Config> {
            Presets() : _nodata() { }
            const Config _nodata;
            inline const Config& get(const std::string& name) const;
        };

        //! Create the value mapping table. You can then pass
        //! this table to the createCoverage() method.
        template<typename T>
        void createMappings(const SourceLayerOptions&, Mappings<T>& output) const;

        //! Find one of the source layers in this coverage.
        //! @return The source layer, or nullptr if we cannot find it
        ImageLayer* getSourceLayerByName(const std::string& name) const;

        /**
        * Factory is a utility class that initializes the mappings between the int values and 
        * the user-defined datastructure.  Initialize one of these first and call createCoverage on it.
        */
        template<typename T>
        class Factory
        {
        public:
            using Ptr = std::unique_ptr<Factory<T>>;

            static Ptr create(CoverageLayer* layer) {
                return Ptr(new Factory<T>(layer));
            }

            Factory(CoverageLayer* layer) :
                _layer(layer)
            {
                // Preload the mappings for each layer
                for (auto& layer : _layer->options().layers())
                {
                    ImageLayer* imageLayer = layer.source().getLayer();

                    if (imageLayer == nullptr || !imageLayer->isOpen())
                        continue;

                    Mappings<T>& mappings = _layerToMappings[imageLayer];
                    _layer->createMappings(layer, mappings);
                }
            }

            //! Create a GeoCoverage for a given tile key.
            //! @param key TileKey for which to create a coverage grid
            //! @param progress Progress indicator (can be nullptr)
            inline GeoCoverage<T> createCoverage(const TileKey& key, ProgressCallback* progress)
            {
                OE_PROFILING_ZONE;
                OE_PROFILING_ZONE_TEXT(_layer->getName() + " " + key.str());

                GeoCoverage<T> result;

                // got something from the cache? return it
                if (readFromCache(key, result, progress))
                    return result;
                
                // nothing in the cache, but we're cache-only, so fail silently
                if (_layer->getCacheSettings()->cachePolicy()->isCacheOnly())
                    return result;

                // populate the result (and write it to the cache if necessary)
                populate(result, key, progress);
                writeToCache(key, result, progress);

                return result;
            }

            inline bool readFromCache(const TileKey& key, GeoCoverage<T>& result, ProgressCallback* p)
            {
                if (_layer->_memCache.valid())
                {
                    char memCacheKey[64];
                    sprintf(memCacheKey, "%d/%s/%s", _layer->getRevision(), key.str().c_str(), key.getProfile()->getHorizSignature().c_str());
                    CacheBin* bin = _layer->_memCache->getOrCreateDefaultBin();
                    ReadResult r = bin->readObject(memCacheKey, nullptr);
                    if (r.succeeded())
                    {
                        result = GeoCoverage<T>(Coverage<T>::create(), key.getExtent());
                        Config conf;
                        conf.fromJSON(r.getString());
                        result.setConfig(conf);
                        return true;
                    }
                }

                CacheBin* cacheBin = _layer->getCacheBin(key.getProfile());
                const CachePolicy& policy = _layer->getCacheSettings()->cachePolicy().get();
                if (!policy.isCacheOnly() && !_layer->getProfile())
                {
                    _layer->disable("Could not establish a valid profile");
                    return false;
                }

                // First, attempt to read from the cache. Since the cached data is stored in the
                // map profile, we can try this first.
                if (cacheBin && policy.isCacheReadable())
                {
                    // the cache key combines the Key and the horizontal profile.
                    std::string cacheKey = Cache::makeCacheKey(key.str() + "-" + key.getProfile()->getHorizSignature(), "coverage");
                    ReadResult r = cacheBin->readString(cacheKey, nullptr);
                    if (r.succeeded() && !policy.isExpired(r.lastModifiedTime()))
                    {
                        result = GeoCoverage<T>(Coverage<T>::create(), key.getExtent());
                        Config conf;
                        conf.fromJSON(r.getString());
                        result.setConfig(conf);
                        return true;
                    }
                }

                // The data was not in the cache. If we are cache-only, fail sliently
                if (policy.isCacheOnly())
                {
                    // If it's cache only and we have an expired but cached image, just return it.
                    return !result.empty();
                }

                return result.valid();
            }

            inline void writeToCache(const TileKey& key, const GeoCoverage<T>& result, ProgressCallback* p)
            {
                if (result.empty())
                    return;

                osg::ref_ptr<StringObject> so;

                if (_layer->_memCache.valid())
                {
                    char memCacheKey[64];
                    sprintf(memCacheKey, "%d/%s/%s", _layer->getRevision(), key.str().c_str(), key.getProfile()->getHorizSignature().c_str());
                    CacheBin* bin = _layer->_memCache->getOrCreateDefaultBin();
                    Config conf = result.getConfig();
                    so = new StringObject(conf.toJSON(false));
                    bin->write(memCacheKey, so.get(), nullptr);
                }

                // If we got a result, the cache is valid and we are caching in the map profile,
                // write to the map cache.
                CacheBin* cacheBin = _layer->getCacheBin(key.getProfile());
                const CachePolicy& policy = _layer->getCacheSettings()->cachePolicy().get();
                if (cacheBin && policy.isCacheWriteable())
                {
                    std::string cacheKey = Cache::makeCacheKey(key.str()+"-"+key.getProfile()->getHorizSignature(), "coverage");
                    if (!so.valid()) {
                        Config conf = result.getConfig();
                        so = new StringObject(conf.toJSON(false));
                    }
                    cacheBin->write(cacheKey, so.get(), nullptr);
                }
            }

        private:

            //! Create a GeoCoverage for a given tile key.
            //! @param key TileKey for which to create a coverage grid
            //! @param mappings Value mapping table returned by createMappings()
            //! @param progress Progress indicator (can be nullptr)
            inline GeoCoverage<T> createCoverage(
                const TileKey& key,
                ImageLayer* source,
                ProgressCallback* progress)
            {
                typename Coverage<T>::Ptr result;

                if (source &&
                    source->isOpen())
                {
                    auto mappingsItr = _layerToMappings.find(source);
                    if (mappingsItr == _layerToMappings.end())
                    {
                        return GeoCoverage<T>();
                    }

                    auto& mappings = mappingsItr->second;


                    GeoImage image = source->createImage(key, progress);
                    if (image.valid())
                    {
                        osg::Vec4 pixel;
                        GeoImagePixelReader read(image);
                        read.setBilinear(false); // unnecessary
                        //bool normalized = image.getImage()->getDataType() == GL_UNSIGNED_BYTE;

                        if (progress && progress->isCanceled())
                            return GeoCoverage<T>();

                        OE_PROFILING_ZONE;
                        OE_PROFILING_ZONE_TEXT("decode");

                        result = Coverage<T>::create();
                        result->allocate(read.s(), read.t());

                        const T* sample_ptr = nullptr;
                        unsigned value_prev;
                        int k = -1;
                        GeoImageIterator iter(image);
                        iter.forEachPixel([&]()
                            {
                                read(pixel, iter.s(), iter.t());
                                unsigned value;
                                //if (normalized)
                                if (pixel.r() < 1.0f)
                                    value = (unsigned)(pixel.r() * 255.0f);
                                else
                                    value = (unsigned)pixel.r();

                                if (!sample_ptr || value != value_prev)
                                {
                                    sample_ptr = &mappings.get(value);
                                    k = -1;
                                }

                                if (sample_ptr->valid())
                                    k = result->write(*sample_ptr, iter.s(), iter.t(), k);

                                value_prev = value;
                            });
                    }
                }
                return GeoCoverage<T>(result, key.getExtent());
            }

            bool populate(
                GeoCoverage<T>& output,
                const TileKey& key,
                ProgressCallback* progress)
            {
                auto& sources = _layer->options().layers();

                // trivial rejection when no sources are present
                if (sources.empty())
                {
                    return false;
                }
                
                // trivial case of a SINGLE layer, when no compositing
                // is necessary
                if (sources.size() == 1)
                {
                    auto imageLayer = sources.front().source().getLayer();
                    if (imageLayer->isOpen())
                    {
                        output = createCoverage(key, imageLayer, progress);
                    }
                    return output.valid();
                }

                // Compositing is necessary. Collect the best available
                // coverage from each source, and then composite them based
                // on their order of appearance.
                // Iterate backwards since the last source has the highest priority.
                // If we get a coverage with all valid values, we are finished.
                std::vector<GeoCoverage<T>> inputs;
                bool fallback = false;

                for (auto i = sources.rbegin(); i != sources.rend(); ++i)
                {
                    auto imageLayer = i->source().getLayer();

                    if (imageLayer && imageLayer->isOpen())
                    {
                        GeoCoverage<T> input;
                        TileKey inputKey = key;

                        if (fallback)
                        {
                            // This path runs when we already have a partial tile from
                            // a higher priority layer, and need to fill in the empty samples
                            // from other lower priority layers.
                            // NB: layer->mayHaveData() is insufficient here. We want to be sure
                            // to fail and fall back until we get data, and mayHaveData() will return
                            // false if our LOD exceeds that of the source layer. Instead we just query
                            // the "best available" key and start the search from there if that
                            // key is valid. (It will be invalid, for example, if the extents don't intersect.)
                            inputKey = imageLayer->getBestAvailableTileKey(inputKey);

                            while (inputKey.valid() && !input.valid() && imageLayer->isKeyInLegalRange(inputKey))
                            {
                                input = createCoverage(inputKey, imageLayer, progress);

                                if (!input.valid())
                                    inputKey.makeParent();
                            }
                        }
                        else if (imageLayer->mayHaveData(inputKey))
                        {
                            // This path runs when we have no data at all yet.
                            input = createCoverage(inputKey, imageLayer, progress);
                        }

                        // if we got a valid coverage, add it to the composition list.
                        if (!input.empty())
                        {
                            inputs.emplace_back(std::move(input));

                            // if this coverage is fully populated, we're done
                            if (inputs.back().nodataCount() == 0)
                                break;
                            else
                                fallback = true;
                        }
                    }
                }

                // Got nothing? bye
                if (inputs.empty())
                {
                    return false;
                }

                // Special case: we found one valid coverage - just return that.
                if (inputs.size() == 1 && 
                    inputs.back().getExtent() == key.getExtent())
                {

                    output = std::move(inputs.back());
                    return true;
                }

                // Time to composite.
                // make a new coverage with the same dimensions as our
                // highest priority input:
                T value;
                osg::Matrix scalebias;
                unsigned width = inputs.front().s();
                unsigned height = inputs.front().t();

                typename Coverage<T>::Ptr composite = Coverage<T>::create();
                composite->allocate(width, height);

                for(int i=0; i<inputs.size(); ++i)
                {
                    auto& input = inputs[i];

                    key.getExtent().createScaleBias(input.getExtent(), scalebias);
                    double scale = scalebias(0, 0);
                    double bias_s = scalebias(3, 0) * width;
                    double bias_t = scalebias(3, 1) * height;

                    for (unsigned t = 0u; t < height; ++t)
                    {
                        unsigned input_t = (unsigned)((double)t * scale + bias_t);
                        for (unsigned s = 0u; s < width; ++s)
                        {
                            unsigned input_s = (unsigned)((double)s * scale + bias_s);

                            if (i == 0 || composite->hasDataAt(s, t) == false) // nodata, OK to overwrite:
                            {
                                if (input.read(value, input_s, input_t)) // true if has data
                                {
                                    composite->write(value, s, t);
                                }
                            }
                        }
                    }

                    if (progress && progress->isCanceled())
                        return false;
                }

                // done. wrap it and return it.
                output = GeoCoverage<T>(composite, key.getExtent());
                return true;
            }

        private:
            CoverageLayer* _layer;
            std::unordered_map<ImageLayer*, CoverageLayer::Mappings<T>> _layerToMappings;
        };

    protected: // Layer

        virtual Status openImplementation() override;

    public: // Layer

        virtual void addedToMap(const Map*) override;

        virtual void removedFromMap(const Map*) override;
    };


    // inline functions for CoverageLayer
    template<typename T>
    const T& CoverageLayer::Mappings<T>::get(unsigned value) const
    {
        auto i = this->find(value);
        return i != this->end() ? i->second : _nodata;
    }

    const Config& CoverageLayer::Presets::get(const std::string& name) const
    {
        auto i = this->find(name);
        return i != this->end() ? i->second : _nodata;
    }

    template<typename T>
    void CoverageLayer::createMappings(
        const SourceLayerOptions& layer,
        CoverageLayer::Mappings<T>& output) const
    {
        Presets presets;
        for (auto& child : options().presets()->children("preset"))
        {
            presets[child.value("name")] = child;
        }
        
        for (auto& child : layer.mappings()->children("mapping"))
        {            
            optional<unsigned> value;
            std::string preset_name;
            if (child.get("value", value))
            {
                if (value == 0u)
                {
                    OE_WARN << "[CoverageLayer] " << getName() << " : Illegal value 0 in mapping (zero is reserved for 'no data'); skipping." << std::endl;
                }
                else if (child.get("preset", preset_name))
                {
                    Config temp = presets.get(preset_name);
                    temp.merge(child);
                    output[value.get()] = T(temp);
                }
                else
                {
                    output[value.get()] = T(child);
                }
            }
        }
    }
} // namespace osgEarth

#endif // OSGEARTH_COVERAGE_LAYER_H
