#include "streaming_engine.h"

#include <jni.h>
#include <android/hardware_buffer_jni.h>
#include <android/native_window_jni.h>
#include <android/log.h>

#define TAG "LckJniBridge"
#define LOGI(...) __android_log_print(ANDROID_LOG_INFO, TAG, __VA_ARGS__)
#define LOGE(...) __android_log_print(ANDROID_LOG_ERROR, TAG, __VA_ARGS__)

static JavaVM* gJavaVM = nullptr;

// Cache for callback method IDs
static jmethodID gOnStatsMethod = nullptr;
static jmethodID gOnErrorMethod = nullptr;
static jmethodID gOnBufferReleasedMethod = nullptr;

JNIEXPORT jint JNI_OnLoad(JavaVM* vm, void* reserved) {
    gJavaVM = vm;
    return JNI_VERSION_1_6;
}

extern "C" {

JNIEXPORT jlong JNICALL
Java_com_omixlab_lckcontrol_streaming_NativeStreamingEngine_nativeCreate(
        JNIEnv* env, jobject thiz,
        jint width, jint height,
        jint videoBitrate, jint audioBitrate,
        jint sampleRate, jint channels,
        jint keyframeInterval) {

    auto* engine = new StreamingEngine();
    engine->Configure(width, height, videoBitrate, audioBitrate,
                      sampleRate, channels, keyframeInterval);

    // Set up callbacks that call back into Kotlin
    jobject globalRef = env->NewGlobalRef(thiz);

    // Cache method IDs
    jclass cls = env->GetObjectClass(thiz);
    gOnStatsMethod = env->GetMethodID(cls, "onNativeStats", "(JJII)V");
    gOnErrorMethod = env->GetMethodID(cls, "onNativeError", "(ILjava/lang/String;)V");
    gOnBufferReleasedMethod = env->GetMethodID(cls, "onNativeBufferReleased", "(I)V");

    engine->SetStatsCallback([globalRef](const StreamingStats& stats) {
        JNIEnv* env;
        if (gJavaVM->GetEnv(reinterpret_cast<void**>(&env), JNI_VERSION_1_6) != JNI_OK) {
            if (gJavaVM->AttachCurrentThread(&env, nullptr) != JNI_OK) return;
        }
        if (gOnStatsMethod) {
            env->CallVoidMethod(globalRef, gOnStatsMethod,
                                (jlong)stats.videoBitrate, (jlong)stats.audioBitrate,
                                (jint)stats.fps, (jint)stats.droppedFrames);
        }
    });

    engine->SetErrorCallback([globalRef](int code, const std::string& message) {
        JNIEnv* env;
        if (gJavaVM->GetEnv(reinterpret_cast<void**>(&env), JNI_VERSION_1_6) != JNI_OK) {
            if (gJavaVM->AttachCurrentThread(&env, nullptr) != JNI_OK) return;
        }
        if (gOnErrorMethod) {
            jstring msg = env->NewStringUTF(message.c_str());
            env->CallVoidMethod(globalRef, gOnErrorMethod, (jint)code, msg);
            env->DeleteLocalRef(msg);
        }
    });

    engine->SetBufferReleasedCallback([globalRef](int bufferIndex) {
        JNIEnv* env;
        if (gJavaVM->GetEnv(reinterpret_cast<void**>(&env), JNI_VERSION_1_6) != JNI_OK) {
            if (gJavaVM->AttachCurrentThread(&env, nullptr) != JNI_OK) return;
        }
        if (gOnBufferReleasedMethod) {
            env->CallVoidMethod(globalRef, gOnBufferReleasedMethod, (jint)bufferIndex);
        }
    });

    LOGI("Native engine created: %dx%d", width, height);
    return reinterpret_cast<jlong>(engine);
}

JNIEXPORT jint JNICALL
Java_com_omixlab_lckcontrol_streaming_NativeStreamingEngine_nativeAddDestination(
        JNIEnv* env, jobject thiz, jlong ptr, jstring rtmpUrl) {
    auto* engine = reinterpret_cast<StreamingEngine*>(ptr);
    if (!engine) return -1;

    const char* url = env->GetStringUTFChars(rtmpUrl, nullptr);
    int index = engine->AddDestination(url);
    env->ReleaseStringUTFChars(rtmpUrl, url);
    return index;
}

JNIEXPORT jboolean JNICALL
Java_com_omixlab_lckcontrol_streaming_NativeStreamingEngine_nativeStart(
        JNIEnv* env, jobject thiz, jlong ptr) {
    auto* engine = reinterpret_cast<StreamingEngine*>(ptr);
    if (!engine) return JNI_FALSE;
    return engine->Start() ? JNI_TRUE : JNI_FALSE;
}

JNIEXPORT void JNICALL
Java_com_omixlab_lckcontrol_streaming_NativeStreamingEngine_nativeSubmitVideoFrame(
        JNIEnv* env, jobject thiz, jlong ptr,
        jobject hardwareBuffer, jlong timestampNs, jint fenceFd, jint bufferIndex) {
    auto* engine = reinterpret_cast<StreamingEngine*>(ptr);
    if (!engine) return;

    AHardwareBuffer* buffer = AHardwareBuffer_fromHardwareBuffer(env, hardwareBuffer);
    if (!buffer) {
        LOGE("Failed to get AHardwareBuffer from Java HardwareBuffer");
        return;
    }

    engine->SubmitVideoFrame(buffer, timestampNs, fenceFd, bufferIndex);
}

JNIEXPORT void JNICALL
Java_com_omixlab_lckcontrol_streaming_NativeStreamingEngine_nativeSubmitAudioFrame(
        JNIEnv* env, jobject thiz, jlong ptr,
        jbyteArray pcmData, jlong timestampNs) {
    auto* engine = reinterpret_cast<StreamingEngine*>(ptr);
    if (!engine) return;

    jsize len = env->GetArrayLength(pcmData);
    jbyte* data = env->GetByteArrayElements(pcmData, nullptr);

    engine->SubmitAudioFrame(reinterpret_cast<const uint8_t*>(data), len, timestampNs);

    env->ReleaseByteArrayElements(pcmData, data, JNI_ABORT);
}

JNIEXPORT void JNICALL
Java_com_omixlab_lckcontrol_streaming_NativeStreamingEngine_nativeStop(
        JNIEnv* env, jobject thiz, jlong ptr) {
    auto* engine = reinterpret_cast<StreamingEngine*>(ptr);
    if (!engine) return;
    engine->Stop();
}

JNIEXPORT void JNICALL
Java_com_omixlab_lckcontrol_streaming_NativeStreamingEngine_nativeDestroy(
        JNIEnv* env, jobject thiz, jlong ptr) {
    auto* engine = reinterpret_cast<StreamingEngine*>(ptr);
    if (engine) {
        engine->Stop();
        delete engine;
        LOGI("Native engine destroyed");
    }
}

JNIEXPORT jboolean JNICALL
Java_com_omixlab_lckcontrol_streaming_NativeStreamingEngine_nativeIsRunning(
        JNIEnv* env, jobject thiz, jlong ptr) {
    auto* engine = reinterpret_cast<StreamingEngine*>(ptr);
    if (!engine) return JNI_FALSE;
    return engine->IsRunning() ? JNI_TRUE : JNI_FALSE;
}

// --- Preview surface ---

JNIEXPORT void JNICALL
Java_com_omixlab_lckcontrol_streaming_NativeStreamingEngine_nativeSetPreviewSurface(
        JNIEnv* env, jobject thiz, jlong ptr, jobject surface) {
    auto* engine = reinterpret_cast<StreamingEngine*>(ptr);
    if (!engine || !surface) return;

    ANativeWindow* window = ANativeWindow_fromSurface(env, surface);
    if (window) {
        engine->SetPreviewSurface(window);
        ANativeWindow_release(window); // SetPreviewSurface acquires its own ref
    }
}

JNIEXPORT void JNICALL
Java_com_omixlab_lckcontrol_streaming_NativeStreamingEngine_nativeRemovePreviewSurface(
        JNIEnv* env, jobject thiz, jlong ptr) {
    auto* engine = reinterpret_cast<StreamingEngine*>(ptr);
    if (!engine) return;
    engine->RemovePreviewSurface();
}

// --- Composition layers ---

JNIEXPORT jint JNICALL
Java_com_omixlab_lckcontrol_streaming_NativeStreamingEngine_nativeAddCompositionLayer(
        JNIEnv* env, jobject thiz, jlong ptr,
        jbyteArray rgbaData, jint w, jint h,
        jfloat posX, jfloat posY, jfloat scaleX, jfloat scaleY,
        jfloat rotation, jfloat opacity, jint zOrder, jstring tag) {
    auto* engine = reinterpret_cast<StreamingEngine*>(ptr);
    if (!engine) return -1;

    jsize len = env->GetArrayLength(rgbaData);
    jbyte* data = env->GetByteArrayElements(rgbaData, nullptr);
    const char* tagStr = env->GetStringUTFChars(tag, nullptr);

    int layerId = engine->AddCompositionLayer(
        reinterpret_cast<const uint8_t*>(data), w, h,
        posX, posY, scaleX, scaleY, rotation, opacity, zOrder,
        std::string(tagStr));

    env->ReleaseStringUTFChars(tag, tagStr);
    env->ReleaseByteArrayElements(rgbaData, data, JNI_ABORT);
    return layerId;
}

JNIEXPORT void JNICALL
Java_com_omixlab_lckcontrol_streaming_NativeStreamingEngine_nativeRemoveCompositionLayer(
        JNIEnv* env, jobject thiz, jlong ptr, jint layerId) {
    auto* engine = reinterpret_cast<StreamingEngine*>(ptr);
    if (!engine) return;
    engine->RemoveCompositionLayer(layerId);
}

JNIEXPORT void JNICALL
Java_com_omixlab_lckcontrol_streaming_NativeStreamingEngine_nativeUpdateCompositionLayerTransform(
        JNIEnv* env, jobject thiz, jlong ptr, jint layerId,
        jfloat posX, jfloat posY, jfloat scaleX, jfloat scaleY, jfloat rotation) {
    auto* engine = reinterpret_cast<StreamingEngine*>(ptr);
    if (!engine) return;
    engine->UpdateCompositionLayerTransform(layerId, posX, posY, scaleX, scaleY, rotation);
}

JNIEXPORT void JNICALL
Java_com_omixlab_lckcontrol_streaming_NativeStreamingEngine_nativeUpdateCompositionLayerOpacity(
        JNIEnv* env, jobject thiz, jlong ptr, jint layerId, jfloat opacity) {
    auto* engine = reinterpret_cast<StreamingEngine*>(ptr);
    if (!engine) return;
    engine->UpdateCompositionLayerOpacity(layerId, opacity);
}

JNIEXPORT void JNICALL
Java_com_omixlab_lckcontrol_streaming_NativeStreamingEngine_nativeSetCompositionLayerEnabled(
        JNIEnv* env, jobject thiz, jlong ptr, jint layerId, jboolean enabled) {
    auto* engine = reinterpret_cast<StreamingEngine*>(ptr);
    if (!engine) return;
    engine->SetCompositionLayerEnabled(layerId, enabled == JNI_TRUE);
}

} // extern "C"
