/**
 * MiniMind (Mind2) JNI Bridge
 * Provides Java/Kotlin interface to llama.cpp inference engine
 */

#include <jni.h>
#include <android/log.h>
#include <android/asset_manager.h>
#include <android/asset_manager_jni.h>

#include <string>
#include <vector>
#include <memory>
#include <thread>
#include <atomic>
#include <mutex>

// If using llama.cpp, include these headers
// #include "llama.h"
// #include "ggml.h"

#define LOG_TAG "Mind2"
#define LOGI(...) __android_log_print(ANDROID_LOG_INFO, LOG_TAG, __VA_ARGS__)
#define LOGE(...) __android_log_print(ANDROID_LOG_ERROR, LOG_TAG, __VA_ARGS__)

namespace {

// Model context (placeholder - would use llama_context in real implementation)
struct Mind2Context {
    std::string model_path;
    int n_ctx = 2048;
    int n_threads = 4;
    bool loaded = false;
    std::atomic<bool> generating{false};
    std::mutex mutex;

    // llama_model* model = nullptr;
    // llama_context* ctx = nullptr;
};

std::unique_ptr<Mind2Context> g_context;

// Token callback for streaming
JavaVM* g_jvm = nullptr;
jobject g_callback = nullptr;
jmethodID g_callback_method = nullptr;

void stream_token(const std::string& token) {
    if (!g_jvm || !g_callback) return;

    JNIEnv* env = nullptr;
    bool attached = false;

    if (g_jvm->GetEnv((void**)&env, JNI_VERSION_1_6) != JNI_OK) {
        g_jvm->AttachCurrentThread(&env, nullptr);
        attached = true;
    }

    if (env && g_callback && g_callback_method) {
        jstring jtoken = env->NewStringUTF(token.c_str());
        env->CallVoidMethod(g_callback, g_callback_method, jtoken);
        env->DeleteLocalRef(jtoken);
    }

    if (attached) {
        g_jvm->DetachCurrentThread();
    }
}

} // anonymous namespace

extern "C" {

JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM* vm, void* reserved) {
    g_jvm = vm;
    LOGI("Mind2 JNI loaded");
    return JNI_VERSION_1_6;
}

JNIEXPORT void JNICALL JNI_OnUnload(JavaVM* vm, void* reserved) {
    g_context.reset();
    g_jvm = nullptr;
    LOGI("Mind2 JNI unloaded");
}

/**
 * Initialize the model
 */
JNIEXPORT jboolean JNICALL
Java_com_minimind_mind2_Mind2Model_nativeInit(
    JNIEnv* env,
    jobject thiz,
    jstring model_path,
    jint n_ctx,
    jint n_threads
) {
    const char* path = env->GetStringUTFChars(model_path, nullptr);
    LOGI("Initializing Mind2 with model: %s", path);

    g_context = std::make_unique<Mind2Context>();
    g_context->model_path = path;
    g_context->n_ctx = n_ctx;
    g_context->n_threads = n_threads > 0 ? n_threads : std::thread::hardware_concurrency();

    env->ReleaseStringUTFChars(model_path, path);

    // TODO: Actual llama.cpp initialization
    // llama_model_params model_params = llama_model_default_params();
    // g_context->model = llama_load_model_from_file(g_context->model_path.c_str(), model_params);
    // if (!g_context->model) {
    //     LOGE("Failed to load model");
    //     return JNI_FALSE;
    // }
    //
    // llama_context_params ctx_params = llama_context_default_params();
    // ctx_params.n_ctx = g_context->n_ctx;
    // ctx_params.n_threads = g_context->n_threads;
    // g_context->ctx = llama_new_context_with_model(g_context->model, ctx_params);

    g_context->loaded = true;
    LOGI("Mind2 initialized successfully (threads: %d, ctx: %d)",
         g_context->n_threads, g_context->n_ctx);

    return JNI_TRUE;
}

/**
 * Generate text from prompt
 */
JNIEXPORT jstring JNICALL
Java_com_minimind_mind2_Mind2Model_nativeGenerate(
    JNIEnv* env,
    jobject thiz,
    jstring prompt,
    jint max_tokens,
    jfloat temperature,
    jfloat top_p,
    jint top_k
) {
    if (!g_context || !g_context->loaded) {
        LOGE("Model not initialized");
        return env->NewStringUTF("");
    }

    std::lock_guard<std::mutex> lock(g_context->mutex);

    const char* prompt_str = env->GetStringUTFChars(prompt, nullptr);
    std::string result;

    LOGI("Generating with prompt: %.50s...", prompt_str);

    // TODO: Actual generation with llama.cpp
    // This is a placeholder that returns the prompt
    result = std::string(prompt_str) + "\n\n[Generated response would appear here]";

    // Actual implementation would be:
    // std::vector<llama_token> tokens = llama_tokenize(g_context->ctx, prompt_str, true);
    // for (int i = 0; i < max_tokens; i++) {
    //     llama_token new_token = llama_sample_token(g_context->ctx, ...);
    //     if (new_token == llama_token_eos(g_context->ctx)) break;
    //     result += llama_token_to_piece(g_context->ctx, new_token);
    //     stream_token(llama_token_to_piece(g_context->ctx, new_token));
    // }

    env->ReleaseStringUTFChars(prompt, prompt_str);

    return env->NewStringUTF(result.c_str());
}

/**
 * Generate with streaming callback
 */
JNIEXPORT void JNICALL
Java_com_minimind_mind2_Mind2Model_nativeGenerateStream(
    JNIEnv* env,
    jobject thiz,
    jstring prompt,
    jint max_tokens,
    jfloat temperature,
    jfloat top_p,
    jint top_k,
    jobject callback
) {
    if (!g_context || !g_context->loaded) {
        LOGE("Model not initialized");
        return;
    }

    // Store callback reference
    g_callback = env->NewGlobalRef(callback);
    jclass callback_class = env->GetObjectClass(callback);
    g_callback_method = env->GetMethodID(callback_class, "onToken", "(Ljava/lang/String;)V");

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

    g_context->generating = true;

    // TODO: Actual streaming generation
    // Simulated streaming for now
    std::vector<std::string> demo_tokens = {
        "Hello", "!", " ", "I", "'m", " ", "Mind2", ",",
        " ", "a", " ", "lightweight", " ", "AI", " ", "assistant", "."
    };

    for (const auto& token : demo_tokens) {
        if (!g_context->generating) break;
        stream_token(token);
        std::this_thread::sleep_for(std::chrono::milliseconds(50));
    }

    // Signal completion
    jmethodID complete_method = env->GetMethodID(callback_class, "onComplete", "()V");
    if (complete_method) {
        env->CallVoidMethod(callback, complete_method);
    }

    env->ReleaseStringUTFChars(prompt, prompt_str);
    env->DeleteGlobalRef(g_callback);
    g_callback = nullptr;
}

/**
 * Stop ongoing generation
 */
JNIEXPORT void JNICALL
Java_com_minimind_mind2_Mind2Model_nativeStop(
    JNIEnv* env,
    jobject thiz
) {
    if (g_context) {
        g_context->generating = false;
        LOGI("Generation stopped");
    }
}

/**
 * Release model resources
 */
JNIEXPORT void JNICALL
Java_com_minimind_mind2_Mind2Model_nativeRelease(
    JNIEnv* env,
    jobject thiz
) {
    if (g_context) {
        std::lock_guard<std::mutex> lock(g_context->mutex);

        // TODO: Release llama.cpp resources
        // if (g_context->ctx) llama_free(g_context->ctx);
        // if (g_context->model) llama_free_model(g_context->model);

        g_context->loaded = false;
        LOGI("Mind2 resources released");
    }
}

/**
 * Get model info
 */
JNIEXPORT jstring JNICALL
Java_com_minimind_mind2_Mind2Model_nativeGetInfo(
    JNIEnv* env,
    jobject thiz
) {
    if (!g_context) {
        return env->NewStringUTF("{}");
    }

    char info[512];
    snprintf(info, sizeof(info),
        "{\"loaded\": %s, \"model\": \"%s\", \"n_ctx\": %d, \"n_threads\": %d}",
        g_context->loaded ? "true" : "false",
        g_context->model_path.c_str(),
        g_context->n_ctx,
        g_context->n_threads
    );

    return env->NewStringUTF(info);
}

/**
 * Benchmark inference speed
 */
JNIEXPORT jfloat JNICALL
Java_com_minimind_mind2_Mind2Model_nativeBenchmark(
    JNIEnv* env,
    jobject thiz,
    jint n_tokens
) {
    if (!g_context || !g_context->loaded) {
        return 0.0f;
    }

    // TODO: Actual benchmark
    // Simulated result
    float tokens_per_second = 25.0f + (rand() % 10);

    LOGI("Benchmark: %.1f tokens/sec", tokens_per_second);
    return tokens_per_second;
}

} // extern "C"