#include <jni.h>
#include <aidl/com/omixlab/mosis/IMosisListener.h>
#include <aidl/com/omixlab/mosis/BnMosisListener.h>
#include <logger.h>
#include <aidl/com/omixlab/mosis/IMosisService.h>
#include <android/binder_ibinder_jni.h>
#include <format>
#include <glad/gles2.h>
#include <glad/egl.h>
#include <IUnityGraphics.h>
#include <external_texture.h>
#include <render_target.h>

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

std::shared_ptr<class IMosisService> g_service;
std::shared_ptr<class ServiceContext> g_context;

class TextureBlitter
{
    GLuint source_texture = 0;
    GLuint source_fb = 0;
    GLuint dest_texture = 0;
    GLuint dest_fb = 0;
    GLint width;
    GLint height;
public:
    bool create(AHardwareBuffer* hardwareBuffer)
    {
        AHardwareBuffer_Desc desc{};
        AHardwareBuffer_describe(hardwareBuffer, &desc);
        width = desc.width;
        height = desc.height;
        EGLClientBuffer clientBuffer = eglGetNativeClientBufferANDROID(hardwareBuffer);
        EGLImageKHR eglImage = eglCreateImageKHR(eglGetCurrentDisplay(), EGL_NO_CONTEXT,
            EGL_NATIVE_BUFFER_ANDROID, clientBuffer, nullptr);
        if (eglImage == EGL_NO_IMAGE_KHR)
        {
            Logger::Log("Failed to create EGL image");
            return false;
        }
        glGenTextures(1, &source_texture);
        glBindTexture(GL_TEXTURE_EXTERNAL_OES, source_texture);
        glTexParameteri(GL_TEXTURE_EXTERNAL_OES, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
        glTexParameteri(GL_TEXTURE_EXTERNAL_OES, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
        glTexParameteri(GL_TEXTURE_EXTERNAL_OES, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
        glTexParameteri(GL_TEXTURE_EXTERNAL_OES, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
        glEGLImageTargetTexture2DOES(GL_TEXTURE_EXTERNAL_OES, eglImage);
        eglDestroyImageKHR(eglGetCurrentDisplay(), eglImage);

        glGenFramebuffers(1, &source_fb);
        glBindFramebuffer(GL_FRAMEBUFFER, source_fb);
        glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0,
            GL_TEXTURE_EXTERNAL_OES, source_texture, 0);
        if (glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE)
            return false;
        glGenTextures(1, &dest_texture);
        glBindTexture(GL_TEXTURE_2D, dest_texture);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, desc.width, desc.height,
            0, GL_RGB, GL_UNSIGNED_BYTE, nullptr);
        glGenFramebuffers(1, &dest_fb);
        glBindFramebuffer(GL_FRAMEBUFFER, dest_fb);
        glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0,
            GL_TEXTURE_2D, dest_texture, 0);
        if (glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE)
            return false;
        blit();
        return true;
    }
    [[nodiscard]] GLuint dest_texture_id() const
    {
        return dest_texture;
    }
    void blit()
    {
        glBindFramebuffer(GL_READ_FRAMEBUFFER, source_fb);
        glBindFramebuffer(GL_DRAW_FRAMEBUFFER, dest_fb);
        glBlitFramebuffer(0, 0, width, height, 0, 0, width, height, GL_COLOR_BUFFER_BIT, GL_NEAREST);
    }
};

typedef void (*OnMessageCallback)(const char*);
typedef void (*OnServiceInitializedCallback)(bool success);
typedef void (*OnFrameAvailableCallback)();
typedef void (*OnBufferReadyCallback)();
typedef void (*OnTextureReadyCallback)(GLuint gl_texture);
struct NativeCallbacks
{
    OnMessageCallback OnMessage;
    OnServiceInitializedCallback OnServiceInitialized;
    OnFrameAvailableCallback OnFrameAvailable;
    OnBufferReadyCallback OnBufferReady;
    OnTextureReadyCallback OnTextureReady;
};
NativeCallbacks g_callbacks{};

class ServiceContext : public BnMosisListener
{
    AHardwareBuffer* m_hwbuffer = nullptr;
    std::unique_ptr<TextureBlitter> m_texture;
public:
    ndk::ScopedAStatus onServiceInitialized(bool in_success) override
    {
        Logger::Log("onServiceInitialized");
        g_callbacks.OnMessage("onServiceInitialized");
        g_callbacks.OnServiceInitialized(in_success);
        return ndk::ScopedAStatus::ok();
    }
    ndk::ScopedAStatus onFrameAvailable() override
    {
        Logger::Log("onFrameAvailable");
        g_callbacks.OnMessage("onFrameAvailable");
        g_callbacks.OnFrameAvailable();
        return ndk::ScopedAStatus::ok();
    }
    ndk::ScopedAStatus onBufferAvailable(const HardwareBuffer &in_buffer) override
    {
        Logger::Log("onBufferAvailable");
        g_callbacks.OnMessage("onBufferAvailable");
        m_hwbuffer = in_buffer.get();
        AHardwareBuffer_acquire(m_hwbuffer);
        g_callbacks.OnBufferReady();
        return ndk::ScopedAStatus::ok();
    }
    [[nodiscard]] GLuint create_texture()
    {
        m_texture = std::make_unique<TextureBlitter>();
        m_texture->create(m_hwbuffer);
        return m_texture->dest_texture_id();
    }
    void update_texture()
    {
        if (m_texture)
            m_texture->blit();
    }
};

extern "C" void SetNativeCallbacks(const NativeCallbacks& ptr)
{
    g_callbacks = ptr;
}

extern "C" UnityRenderingEvent InitGLAD()
{
    return [](int eventId){
        gladLoaderLoadEGL(EGL_NO_DISPLAY);
        int egl_version = gladLoadEGL(eglGetCurrentDisplay(), eglGetProcAddress);
        if (egl_version == 0)
        {
            Logger::Log("Failed to load EGL");
            return;
        }
        int gl_version = gladLoaderLoadGLES2();
        if (gl_version == 0)
        {
            Logger::Log("Failed to load GL");
            return;
        }
        if (g_context)
        {
            GLuint texture = g_context->create_texture();
            g_callbacks.OnTextureReady(texture);
        }
    };
}

extern "C" UnityRenderingEvent UpdateTexture()
{
    return [](int eventId){
        if (g_context)
        {
            g_context->update_texture();
        }
    };
}

extern "C"
JNIEXPORT void JNICALL
Java_com_omixlab_mosis_unity_MyKotlinPlugin_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");
    g_context = ndk::SharedRefBase::make<ServiceContext>();
    bool result{};
    g_service->initOS(g_context, &result);
    Logger::Log(std::format("InitOS returned {}", result));
}