// Copyright (c) 2023 Dominic Masters
// 
// This software is released under the MIT License.
// https://opensource.org/licenses/MIT

#pragma once
#include "display/shader/Shader.hpp"

namespace Dawn {
  typedef int64_t shaderlock_t;
  typedef int16_t shaderid_t;

  class ShaderManager {
    private:
      int32_t nextId;
      shaderlock_t nextLock;
      std::map<shaderid_t, std::shared_ptr<Shader>> shaders;
      std::map<shaderlock_t, shaderid_t> shaderLocks;
      std::map<shaderid_t, std::vector<shaderlock_t>> shaderLocksByShader;

      /**
       * Returns the shader id for the given shader type, or -1 if it is not
       * loaded.
       * 
       * @return The shader id for the shader, or -1 if it is not loaded.
       */
      template<class T>
      shaderid_t getShaderId() {
        auto it = shaders.begin();
        while(it != shaders.end()) {
          auto asT = std::dynamic_pointer_cast<T>(it->second);
          if(asT != nullptr) return asT->shaderId;
          ++it;
        }
        return -1;
      }

    public:
      /**
       * Creates a new shader manager.
       */
      ShaderManager();

      /**
       * Locks a shader of the given type. If the shader is not already loaded,
       * it will be loaded. If the shader is already loaded, it will be
       * returned.
       * 
       * @return The shader lock for the shader of the given type.
       */
      template<class T>
      shaderlock_t lockShader() {
        auto shaderId = this->getShaderId<T>();
        if(shaderId == -1) {
          auto shader = std::make_shared<T>();
          shader->compile();
          shader->shaderId = this->nextId++;
          this->shaders[shader->shaderId] = shader;
          shaderId = shader->shaderId;
        }

        shaderlock_t lock = this->nextLock++;
        this->shaderLocks[lock] = shaderId;
        this->shaderLocksByShader[shaderId].push_back(lock);
        return lock;
      }

      /**
       * Returns the shader for the given lock.
       * 
       * @param lock The shader lock.
       * @return The shader for the given lock.
       */
      template<class T>
      std::shared_ptr<T> getShader(shaderlock_t lock) {
        auto shaderId = this->shaderLocks[lock];
        return std::static_pointer_cast<T>(this->shaders[shaderId]);
      }

      /**
       * Releases the shader for the given lock. This will unload any shader 
       * that is no longer in use.
       * 
       * @param lock Lock to release.
       */
      template<class T>
      void releaseShader(shaderlock_t lock) {
        auto shaderId = this->shaderLocks[lock];
        this->shaderLocks.erase(lock);
        
        auto& locks = this->shaderLocksByShader[shaderId];
        auto it = std::find(locks.begin(), locks.end(), lock);
        if(it != locks.end()) locks.erase(it);

        if(locks.size() == 0) {
          this->shaderLocksByShader.erase(shaderId);
          this->shaders.erase(shaderId);
        }
      }

      /**
       * Destroys the shader manager.
       */
      ~ShaderManager();
  };
}