#include <jni.h>
#include <string>
#include <android/log.h>

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

// Genie C API declarations (from libGenie.so)
extern "C" {
    // Opaque types
    typedef void* GenieDialogConfig;
    typedef void* GenieDialog;
    typedef void* GenieTokenizer;
    typedef void* GenieSampler;

    // Version
    int Genie_getApiMajorVersion();
    int Genie_getApiMinorVersion();
    int Genie_getApiPatchVersion();

    // DialogConfig
    GenieDialogConfig GenieDialogConfig_createFromJson(const char* jsonPath);
    void GenieDialogConfig_free(GenieDialogConfig config);

    // Dialog
    GenieDialog GenieDialog_create(GenieDialogConfig config);
    void GenieDialog_free(GenieDialog dialog);
    const char* GenieDialog_query(GenieDialog dialog, const char* prompt);
    void GenieDialog_setStopSequence(GenieDialog dialog, const char* stopSeq);
    void GenieDialog_reset(GenieDialog dialog);
    void GenieDialog_signal(GenieDialog dialog);
    GenieTokenizer GenieDialog_getTokenizer(GenieDialog dialog);
    GenieSampler GenieDialog_getSampler(GenieDialog dialog);

    // Sampler callback
    typedef bool (*GenieSamplerCallback)(const char* token, void* userData);
    void GenieSampler_registerUserDataCallback(
        GenieSampler sampler,
        GenieSamplerCallback callback,
        void* userData
    );

    // Tokenizer
    const char* GenieTokenizer_decode(GenieTokenizer tokenizer, const int* tokens, int numTokens);
    int* GenieTokenizer_encode(GenieTokenizer tokenizer, const char* text, int* numTokens);
}

// Check if a pointer looks like a valid heap pointer (not an error code)
// Genie SDK returns small negative int values as error codes (e.g. -5 = 0xfffffffb)
// On ARM64 Android, valid heap pointers are always > 0x100000000 (above 4GB)
static bool isValidPointer(void* ptr) {
    uintptr_t addr = reinterpret_cast<uintptr_t>(ptr);
    if (ptr == nullptr) return false;
    // Any value that fits in 32 bits is likely an error code, not a heap pointer
    if (addr <= 0xFFFFFFFFULL) return false;
    return true;
}

// Callback context for token streaming
struct CallbackContext {
    JNIEnv* env;
    jobject callback;
    jmethodID onTokenMethod;
    std::string fullResponse;
    bool shouldStop;
};

static bool samplerCallback(const char* token, void* userData) {
    auto* ctx = reinterpret_cast<CallbackContext*>(userData);
    if (ctx->shouldStop || token == nullptr) return false;

    ctx->fullResponse += token;

    if (ctx->callback != nullptr) {
        JNIEnv* env = ctx->env;
        jstring jToken = env->NewStringUTF(token);
        jboolean continueGen = env->CallBooleanMethod(
            ctx->callback, ctx->onTokenMethod, jToken
        );
        env->DeleteLocalRef(jToken);

        if (!continueGen) {
            ctx->shouldStop = true;
            return false;
        }
    }
    return true;
}

extern "C" {

JNIEXPORT jlong JNICALL
Java_com_kazeia_llm_GenieJni_createDialog(JNIEnv* env, jobject, jstring configPath) {
    const char* path = env->GetStringUTFChars(configPath, nullptr);
    LOGI("Creating dialog from config: %s", path);

    GenieDialogConfig config = GenieDialogConfig_createFromJson(path);
    env->ReleaseStringUTFChars(configPath, path);

    if (!isValidPointer(config)) {
        LOGE("Failed to create dialog config (returned %p, likely error code %ld)",
             config, reinterpret_cast<long>(config));
        return 0;
    }

    LOGI("Dialog config created: %p", config);
    GenieDialog dialog = GenieDialog_create(config);
    GenieDialogConfig_free(config);

    if (!isValidPointer(dialog)) {
        LOGE("Failed to create dialog (returned %p, likely error code %ld)",
             dialog, reinterpret_cast<long>(dialog));
        return 0;
    }

    LOGI("Dialog created successfully: %p", dialog);
    return reinterpret_cast<jlong>(dialog);
}

JNIEXPORT jstring JNICALL
Java_com_kazeia_llm_GenieJni_query(
    JNIEnv* env, jobject, jlong dialogHandle, jstring prompt, jobject callback
) {
    auto* dialog = reinterpret_cast<GenieDialog>(dialogHandle);

    if (!isValidPointer(dialog)) {
        LOGE("Invalid dialog handle: %p", dialog);
        return env->NewStringUTF("[Erreur: modèle LLM non chargé]");
    }

    const char* promptStr = env->GetStringUTFChars(prompt, nullptr);

    CallbackContext ctx;
    ctx.env = env;
    ctx.callback = callback;
    ctx.shouldStop = false;

    if (callback != nullptr) {
        jclass cbClass = env->GetObjectClass(callback);
        ctx.onTokenMethod = env->GetMethodID(
            cbClass, "onToken", "(Ljava/lang/String;)Z"
        );

        // Register the sampler callback
        GenieSampler sampler = GenieDialog_getSampler(dialog);
        if (isValidPointer(sampler)) {
            GenieSampler_registerUserDataCallback(sampler, samplerCallback, &ctx);
        } else {
            LOGE("Invalid sampler pointer: %p", sampler);
        }
    }

    LOGI("Query: %.80s...", promptStr);
    const char* response = GenieDialog_query(dialog, promptStr);
    env->ReleaseStringUTFChars(prompt, promptStr);

    // Use callback accumulated response if available, otherwise use direct response
    std::string result;
    if (!ctx.fullResponse.empty()) {
        result = ctx.fullResponse;
    } else if (response != nullptr) {
        result = response;
    }

    LOGI("Response length: %zu chars", result.size());
    return env->NewStringUTF(result.c_str());
}

JNIEXPORT void JNICALL
Java_com_kazeia_llm_GenieJni_setStopSequence(
    JNIEnv* env, jobject, jlong dialogHandle, jstring stopSequence
) {
    auto* dialog = reinterpret_cast<GenieDialog>(dialogHandle);
    if (!isValidPointer(dialog)) return;
    const char* seq = env->GetStringUTFChars(stopSequence, nullptr);
    GenieDialog_setStopSequence(dialog, seq);
    env->ReleaseStringUTFChars(stopSequence, seq);
}

JNIEXPORT void JNICALL
Java_com_kazeia_llm_GenieJni_freeDialog(JNIEnv*, jobject, jlong dialogHandle) {
    auto* dialog = reinterpret_cast<GenieDialog>(dialogHandle);
    if (isValidPointer(dialog)) {
        GenieDialog_free(dialog);
        LOGI("Dialog freed");
    }
}

JNIEXPORT jstring JNICALL
Java_com_kazeia_llm_GenieJni_getVersion(JNIEnv* env, jobject) {
    int major = Genie_getApiMajorVersion();
    int minor = Genie_getApiMinorVersion();
    int patch = Genie_getApiPatchVersion();
    std::string version = std::to_string(major) + "." +
                          std::to_string(minor) + "." +
                          std::to_string(patch);
    return env->NewStringUTF(version.c_str());
}

} // extern "C"