#include <jni.h>
#include <android/binder_ibinder_jni.h>
#include <aidl/com/omixlab/mosis/IMosisService.h>
#include <aidl/com/omixlab/mosis/IMosisListener.h>
#include <aidl/com/omixlab/mosis/BnMosisListener.h>
#include <android/native_window_jni.h>
#include <android/asset_manager.h>
#include <android/asset_manager_jni.h>
#include "logger.h"
#include "egl_context.h"
#include "assets_manager.h"
#include "shader.h"
#include "external_texture.h"
#include "quad.h"
#include <glad/gles2.h>
#include <glad/egl.h>
#include <memory>
#include <thread>
#include <vector>
#include <future>
#include <functional>

using namespace aidl::com::omixlab::mosis;
using namespace aidl::android::hardware;

class Renderer : public BnMosisListener
{
    std::unique_ptr<egl::Context> m_egl_context;
    std::unique_ptr<ExternalTexture> m_texture;
    std::unique_ptr<Shader> m_quad_shader;
    std::unique_ptr<Quad> m_quad;
    AHardwareBuffer* m_hwbuffer = nullptr;
    std::vector<std::function<void()>> m_task_queue;
    std::mutex m_task_mutex;
    std::thread m_render_loop;
    bool m_active = false;
    void render_loop(ANativeWindow* window)
    {
        m_egl_context = std::make_unique<egl::Context>();
        if (m_egl_context->create(window))
        {
            init_resources();
            m_active = true;
            while (m_active)
            {
                process_tasks();
                render_frame();
            }
        }
        else
        {
            Logger::Log("Failed to create EGL context");
        }
    }
    void init_resources()
    {
        m_quad_shader = std::make_unique<Shader>();
        m_quad_shader->load("quad.vs.glsl", "quad.fs.glsl");
        m_quad = std::make_unique<Quad>();
        m_quad->create();
    }
    void process_tasks()
    {
        std::lock_guard _lock(m_task_mutex);
        for (auto& task : m_task_queue)
        {
            task();
        }
        m_task_queue.clear();
    }
    void render_frame()
    {
        glClearColor(0.5f, 0.5f, 0.5f, 1.0f);
        glClear(GL_COLOR_BUFFER_BIT);
        if (m_texture)
        {
            m_quad_shader->use();
            glUniform1i(m_quad_shader->getUniformLocation("uTexture"), 0);
            glActiveTexture(GL_TEXTURE0);
            m_texture->bind();
            m_quad->draw();
        }
        m_egl_context->swap();
    }
public:
    Renderer() = default;
    bool create(ANativeWindow* window)
    {
        m_render_loop = std::thread(&Renderer::render_loop, this, window);
        return true;
    }
    void destroy()
    {
        m_active = false;
        if (m_render_loop.joinable())
            m_render_loop.join();
        if (m_egl_context)
            m_egl_context->destroy();
        m_egl_context.reset();
    }
    template<typename T>
    [[nodiscard]] std::future<void> exec(T&& lambda)
    {
        std::lock_guard _lock(m_task_mutex);
        auto promise = std::make_shared<std::promise<void>>();
        auto future = promise->get_future();
        m_task_queue.emplace_back(
            [l=std::forward<T>(lambda),p=std::move(promise)]
            {
                l();
                p->set_value();
            }
        );
        return future;
    }
    template<typename T>
    void exec_async(T&& lambda)
    {
        std::lock_guard _lock(m_task_mutex);
        m_task_queue.emplace_back(std::forward<T>(lambda));
    }
    void update_surface(ANativeWindow* window)
    {
        exec_async([this, window]{
            m_egl_context->make_current(window);
        });
    }
    ndk::ScopedAStatus onServiceInitialized(bool in_success) override
    {
        return ndk::ScopedAStatus::ok();
    }
    ndk::ScopedAStatus onFrameAvailable() override
    {
        Logger::Log("onFrameAvailable");
        return ndk::ScopedAStatus::ok();
    }
    ndk::ScopedAStatus onBufferAvailable(const HardwareBuffer &in_buffer) override
    {
        m_hwbuffer = in_buffer.get();
        AHardwareBuffer_acquire(m_hwbuffer);
        exec_async([this]{
            m_texture = std::make_unique<ExternalTexture>();
            if (!m_texture->create(m_hwbuffer))
            {
                Logger::Log("Failed to create texture");
            }
        });
        return ndk::ScopedAStatus::ok();
    }
};

std::shared_ptr<IMosisService> g_service;
std::shared_ptr<Renderer> g_renderer;

extern "C"
JNIEXPORT void JNICALL
Java_com_omixlab_mosis_MainActivity_serviceConnected(JNIEnv *env, jobject thiz,
    jobject binder)
{
    AIBinder* pBinder = AIBinder_fromJavaBinder(env, binder);
    const ndk::SpAIBinder spBinder(pBinder);
    g_service = IMosisService::fromBinder(spBinder);
    Logger::Log("Service Connected");
}

extern "C"
JNIEXPORT void JNICALL
Java_com_omixlab_mosis_MainActivity_initOS(JNIEnv *env, jobject thiz)
{
    if (g_service && g_renderer)
    {
        bool result{};
        g_service->initOS(g_renderer, &result);
    }
}

extern "C"
JNIEXPORT void JNICALL
Java_com_omixlab_mosis_MainActivity_setSurface(JNIEnv *env, jobject thiz,
   jobject surface)
{
    Logger::Log("setSurface");
    if (!g_renderer)
    {
        g_renderer = ndk::SharedRefBase::make<Renderer>();
        g_renderer->create(ANativeWindow_fromSurface(env, surface));
    }
    else
    {
        g_renderer->update_surface(ANativeWindow_fromSurface(env, surface));
    }
}
extern "C"
JNIEXPORT void JNICALL
Java_com_omixlab_mosis_MainActivity_destroySurface(JNIEnv *env, jobject thiz)
{
    Logger::Log("destroySurface");
    if (g_renderer)
    {
        g_renderer->destroy();
        g_renderer.reset();
    }
}

extern "C"
JNIEXPORT void JNICALL
Java_com_omixlab_mosis_MainActivity_setAssetManager(JNIEnv *env, jobject thiz,
    jobject asset_manager)
{
    AssetsManager::Init(AAssetManager_fromJava(env, asset_manager));
}