#include "mvgal_ai_features.h"
#include "scheduler_internal.h"

#include <string.h>
#include <stdint.h>

/* ── Constants ─────────────────────────────────────────────────────────── */

#define AI_FEATURE_GPU_HISTORY  8   /* matches mvgal_ai_features_t::gpu_utilization[8] etc. */
#define AI_WORKLOAD_HISTORY     16  /* matches mvgal_ai_features_t::workload_history[16] */
#define AI_MAX_GPUS             8   /* matches feature array sizes */
#define AI_MAX_QUEUE_DEPTH      64.0f

/* ── Internal state ───────────────────────────────────────────────────── */

static mvgal_ai_model_t g_ai_model = NULL;

/* ── Internal helpers ─────────────────────────────────────────────────── */

/* Normalize a value into [0.0, 1.0] range using min/max clamping */
static inline float normalize(float val, float min, float max) {
    if (max <= min) return 0.5f;
    float n = (val - min) / (max - min);
    if (n < 0.0f) n = 0.0f;
    if (n > 1.0f) n = 1.0f;
    return n;
}

/* Extract gpu_utilization from GPU state array (0.0-1.0 normalized) */
static void extract_utilization(const mvgal_gpu_state_t *gpu_states,
                                 uint32_t num_gpus,
                                 float *util_out,
                                 uint32_t max_entries)
{
    uint32_t count = (num_gpus < max_entries) ? num_gpus : max_entries;
    for (uint32_t i = 0; i < count; i++) {
        util_out[i] = normalize(gpu_states[i].utilization, 0.0f, 100.0f);
    }
    for (uint32_t i = count; i < max_entries; i++) {
        util_out[i] = 0.5f;
    }
}

/* Extract vram_pressure = used / total (0.0-1.0) */
static void extract_vram_pressure(const mvgal_gpu_state_t *gpu_states,
                                   uint32_t num_gpus,
                                   float *pressure_out,
                                   uint32_t max_entries)
{
    uint32_t count = (num_gpus < max_entries) ? num_gpus : max_entries;
    for (uint32_t i = 0; i < count; i++) {
        const mvgal_gpu_descriptor_t *info = &gpu_states[i].gpu_info;
        if (info->vram_total > 0) {
            float ratio = (float)info->vram_used / (float)info->vram_total;
            pressure_out[i] = (ratio > 1.0f) ? 1.0f : ratio;
        } else {
            pressure_out[i] = 0.0f;
        }
    }
    for (uint32_t i = count; i < max_entries; i++) {
        pressure_out[i] = 0.0f;
    }
}

/* Extract queue depth by walking ready + running linked lists.
 * Counts workloads assigned to each GPU. */
static void extract_queue_depth(const mvgal_scheduler_state_t *state,
                                 float *depth_out,
                                 uint32_t max_entries)
{
    uint32_t count = (state->gpu_count < max_entries) ? state->gpu_count : max_entries;
    uint32_t pending[AI_MAX_GPUS];
    memset(pending, 0, sizeof(pending));

    /* Walk ready queue */
    const struct mvgal_workload *w = state->ready_queue.head;
    while (w) {
        for (uint32_t g = 0; g < w->assigned_gpu_count && g < count; g++) {
            uint32_t gpu_idx = w->assigned_gpus[g];
            if (gpu_idx < count) pending[gpu_idx]++;
        }
        w = w->next;
    }

    /* Walk running queue */
    w = state->running_queue.head;
    while (w) {
        for (uint32_t g = 0; g < w->assigned_gpu_count && g < count; g++) {
            uint32_t gpu_idx = w->assigned_gpus[g];
            if (gpu_idx < count) pending[gpu_idx]++;
        }
        w = w->next;
    }

    for (uint32_t i = 0; i < count; i++) {
        depth_out[i] = normalize((float)pending[i], 0.0f, AI_MAX_QUEUE_DEPTH);
    }
    for (uint32_t i = count; i < max_entries; i++) {
        depth_out[i] = 0.0f;
    }
}

/* Encode workload type + duration + memory into workload_history[0..2] */
static void encode_workload_type(const mvgal_workload_descriptor_t *wkl,
                                  float *hist_out,
                                  uint32_t max_entries)
{
    if (max_entries > 0) {
        hist_out[0] = normalize((float)wkl->type, 0.0f, 16.0f);
    }
    if (max_entries > 1) {
        hist_out[1] = normalize((float)wkl->estimated_duration_ns,
                                0.0f, 1e9f);
    }
    if (max_entries > 2) {
        hist_out[2] = normalize((float)wkl->memory_required,
                                0.0f, (float)(16ULL * 1024 * 1024 * 1024));
    }
    for (uint32_t i = 3; i < max_entries; i++) {
        hist_out[i] = 0.0f;
    }
}

/* Extract temperature from GPU state (0.0-1.0 normalized) */
static void extract_temperature(const mvgal_gpu_state_t *gpu_states,
                                 uint32_t num_gpus,
                                 float *temp_out,
                                 uint32_t max_entries)
{
    uint32_t count = (num_gpus < max_entries) ? num_gpus : max_entries;
    for (uint32_t i = 0; i < count; i++) {
        temp_out[i] = normalize(gpu_states[i].temperature, 0.0f, 100.0f);
    }
    for (uint32_t i = count; i < max_entries; i++) {
        temp_out[i] = 0.0f;
    }
}

/* ── Public API ───────────────────────────────────────────────────────── */

int mvgal_ai_features_extract(const mvgal_scheduler_state_t *sched_state,
                               const mvgal_workload_descriptor_t *workload,
                               mvgal_ai_features_t *features)
{
    if (!sched_state || !features) return -1;
    if (sched_state->gpu_count == 0) return -1;

    memset(features, 0, sizeof(*features));

    uint32_t num_gpus = sched_state->gpu_count;
    if (num_gpus > AI_MAX_GPUS)
        num_gpus = AI_MAX_GPUS;

    features->num_gpus = num_gpus;

    extract_utilization(sched_state->gpus, num_gpus,
                        features->gpu_utilization, AI_FEATURE_GPU_HISTORY);
    extract_vram_pressure(sched_state->gpus, num_gpus,
                          features->vram_pressure, AI_FEATURE_GPU_HISTORY);
    extract_queue_depth(sched_state, features->queue_depth,
                        AI_FEATURE_GPU_HISTORY);

    if (workload) {
        encode_workload_type(workload, features->workload_history,
                             AI_WORKLOAD_HISTORY);
        features->workload_type = (uint32_t)workload->type;
    } else {
        for (uint32_t i = 0; i < AI_WORKLOAD_HISTORY; i++) {
            features->workload_history[i] = 0.5f;
        }
        features->workload_type = 0;
    }

    extract_temperature(sched_state->gpus, num_gpus,
                        features->temperature, AI_FEATURE_GPU_HISTORY);

    return 0;
}

int mvgal_ai_features_predict(const mvgal_ai_features_t *features,
                               mvgal_ai_prediction_t *prediction)
{
    if (!features || !prediction) return -1;

    memset(prediction, 0, sizeof(*prediction));

    if (g_ai_model) {
        int rc = mvgal_ai_model_predict(g_ai_model, features, prediction);
        if (rc == 0) return 0;
    }

    /* Model unavailable or predict failed: fall back to heuristic.
     * Pick GPU with lowest queue depth + highest free VRAM. */
    uint32_t best = 0;
    float best_score = -1.0f;
    for (uint32_t i = 0; i < features->num_gpus && i < AI_MAX_GPUS; i++) {
        float score = (1.0f - features->queue_depth[i]) * 0.6f +
                      (1.0f - features->vram_pressure[i]) * 0.4f;
        if (score > best_score) {
            best_score = score;
            best = i;
        }
    }

    prediction->recommended_gpu = (int32_t)best;
    prediction->confidence = best_score;
    for (uint32_t i = 0; i < features->num_gpus && i < AI_MAX_GPUS; i++) {
        prediction->gpu_scores[i] = (1.0f - features->queue_depth[i]);
    }
    prediction->inference_time_us = 0;

    return 1; /* fallback used */
}

int mvgal_ai_features_dispatch(mvgal_scheduler_state_t *sched_state,
                                const mvgal_workload_descriptor_t *workload)
{
    if (!sched_state) return -1;

    mvgal_ai_features_t features;
    mvgal_ai_prediction_t prediction;

    int rc = mvgal_ai_features_extract(sched_state, workload, &features);
    if (rc != 0) return -1;

    rc = mvgal_ai_features_predict(&features, &prediction);
    if (rc != 0 && rc != 1) return -1;

    /* If AI prediction has sufficient confidence, use recommended GPU */
    if (prediction.confidence >= 0.3f &&
        prediction.recommended_gpu >= 0 &&
        (uint32_t)prediction.recommended_gpu < sched_state->gpu_count)
    {
        return (int)prediction.recommended_gpu;
    }

    /* Fallback: use scheduler's find_best_gpu */
    struct mvgal_workload *wkl = mvgal_workload_create_internal(NULL);
    if (!wkl) {
        /* Last resort: pick first available GPU */
        for (uint32_t i = 0; i < sched_state->gpu_count; i++) {
            if (sched_state->gpus[i].available) return (int)i;
        }
        return 0;
    }

    uint32_t gpu_idx = 0;
    rc = mvgal_scheduler_find_best_gpu(wkl, &gpu_idx);
    mvgal_workload_release(wkl);

    if (rc != MVGAL_SUCCESS) {
        /* Last resort */
        for (uint32_t i = 0; i < sched_state->gpu_count; i++) {
            if (sched_state->gpus[i].available) return (int)i;
        }
        return 0;
    }

    return (int)gpu_idx;
}

bool mvgal_ai_features_ready(void) {
    return (g_ai_model != NULL);
}

mvgal_ai_model_t mvgal_ai_features_reload_model(const char *model_path) {
    /* Unload existing model if any */
    if (g_ai_model) {
        mvgal_ai_model_unload(g_ai_model);
        g_ai_model = NULL;
    }

    g_ai_model = mvgal_ai_model_load(model_path);
    return g_ai_model;
}
