/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/**
 * @file die_manager.c
 * @brief Implementation of the die manager (see die_manager.h). One
 *        process-global table (dies + bindings) behind a mutex, mirroring the
 *        singleton style of pn_memwatch. Placement policy detail lives in
 *        pn_dm_place below; everything else is straight accounting.
 *
 * @status Logical path tested on the Mac; CUDA and TT probes gated on GPU hardware.
 * @issues None
 * @todo None
 */

#include "die_manager.h"

#include <pthread.h>
#include <stdio.h>
#include <string.h>

#ifdef USE_CUDA
#include "cuda/fc_cuda.h"
#include "../core/mem/pn_memwatch.h"
#endif

#ifdef USE_TT
#include "tt/ft_tt_coord.h"
#include "tt/ft_tt_device.h"
#include "tt/ft_tt_runtime.h"
#include "../core/mem/pn_memwatch.h"
#endif

/* One resident model's reservation: which dies it holds and how many bytes it
 * put on each (total split evenly across the set). */
typedef struct {
    bool     used;
    char     name[PN_DM_NAME_MAX];
    int      dies[PN_DM_MAX_DIES];
    int      n_dies;
    uint64_t bytes_per_die;
} dm_binding;

typedef struct {
    bool     present;
    uint64_t vram_budget;
    uint64_t vram_used;
} dm_die;

static struct {
    pthread_mutex_t lock;
    bool            running;
    int             n_dies;
    dm_die          die[PN_DM_MAX_DIES];
    dm_binding      binding[PN_DM_MAX_BINDINGS];
} g_dm = { .lock = PTHREAD_MUTEX_INITIALIZER };

/* --- init / shutdown --------------------------------------------------- */

int pn_dm_init_logical(int n_dies, uint64_t budget_bytes)
{
    if (n_dies <= 0 || n_dies > PN_DM_MAX_DIES || budget_bytes == 0)
        return -1;
    pthread_mutex_lock(&g_dm.lock);
    if (g_dm.running) {                     /* idempotent while running */
        pthread_mutex_unlock(&g_dm.lock);
        return 0;
    }
    memset(g_dm.die, 0, sizeof(g_dm.die));
    memset(g_dm.binding, 0, sizeof(g_dm.binding));
    g_dm.n_dies = n_dies;
    for (int i = 0; i < n_dies; i++) {
        g_dm.die[i].present = true;
        g_dm.die[i].vram_budget = budget_bytes;
    }
    g_dm.running = true;
    pthread_mutex_unlock(&g_dm.lock);
    return 0;
}

#ifdef USE_CUDA
int pn_dm_init_cuda(int n_dies)
{
    if (n_dies <= 0 || n_dies > PN_DM_MAX_DIES)
        return -1;
    /* Driver-only gate on purpose: this runs at serve startup, before any
     * model has a placement, so it must not be the thing that forces a
     * process-default context onto whatever die happens to be ordinal 0. */
    if (!fc_driver_available())
        return -1;
    pthread_mutex_lock(&g_dm.lock);
    if (g_dm.running) {
        int n = g_dm.n_dies;
        pthread_mutex_unlock(&g_dm.lock);
        return n;
    }
    memset(g_dm.die, 0, sizeof(g_dm.die));
    memset(g_dm.binding, 0, sizeof(g_dm.binding));
    int found = 0;
    uint64_t reserve = pn_die_reserve_bytes();
    for (int i = 0; i < n_dies; i++) {
        /* Transient bare context per die — created, read, destroyed. Probing
         * with fc_cuda_init_device() stranded a full context (cuBLAS + 12 JIT'd
         * PTX modules, ~124 MiB) on EVERY die the serve enumerated, including
         * dies owned by another serve. A probe must leave no trace. */
        size_t die_free = 0, die_total = 0;
        if (fc_cuda_probe_device(i, &die_free, &die_total) != 0 || die_total == 0)
            break;                      /* ordinals are contiguous — stop */
        uint64_t budget = die_total > reserve ? die_total - reserve : 0;
        g_dm.die[i].present = true;
        g_dm.die[i].vram_budget = budget;
        /* Anything already gone from `free` beyond the reserve is occupancy we
         * did not create (another process / stranded VRAM) — count it so
         * placement doesn't over-commit the die. */
        uint64_t external = (die_total - die_free) > reserve
                          ? (die_total - die_free) - reserve : 0;
        g_dm.die[i].vram_used = external < budget ? external : budget;
        /* 0 bytes ours: the probe context is already destroyed, and no model
         * has been placed on this die yet. */
        if (pn_memwatch_running())
            pn_memwatch_set_die(i, die_free, budget, 0);
        found++;
    }
    if (found == 0) {
        pthread_mutex_unlock(&g_dm.lock);
        return -1;
    }
    g_dm.n_dies = found;
    g_dm.running = true;
    pthread_mutex_unlock(&g_dm.lock);
    fprintf(stderr, "die_manager: tracking %d CUDA die%s\n",
            found, found == 1 ? "" : "s");
    return found;
}
#endif /* USE_CUDA */

#ifdef USE_TT
bool pn_dm_tt_driver_explicit(void)
{
    ft_tt_runtime_driver driver = FT_TT_RUNTIME_DRIVER_KMD;
    bool explicit_choice = false;

    /* A REFUSED value (neither kmd nor udrv) counts as explicit: somebody asked
     * for a transport and did not get one, which must not read as "no card". */
    if (ft_tt_runtime_driver_select(&driver, &explicit_choice) != FT_TT_OK)
        return true;
    return explicit_choice;
}

/* Capacity of one Blackhole device, in bytes, or 0 if it cannot be measured.
 *
 * The probe leaves no trace: it opens the device, reads the grid and closes
 * again. That matters for the same reason it did on CUDA — this runs at serve
 * startup, before any model has a placement, and a probe that held a context
 * or a TLB window would be holding it against whichever job takes the card.
 *
 * It goes through ft_tt_runtime_open() rather than ft_tt_init_device() so it
 * probes over WHICHEVER TRANSPORT the process is configured for. Opening the
 * chardev here would have refused a machine with no vendor module before the
 * serve ever reached the seam — which is exactly the defect this round exists
 * to remove, one layer up from where it was first seen.
 *
 * The capacity is NOT a constant and is deliberately not defaulted. Bank count
 * comes from ARC telemetry and bank size from the architecture table selected
 * at open; ft_tt_runtime_dram_bytes() answers 0 rather than reporting the
 * nominal eight banks as if it had measured them, so a telemetry failure ends
 * the scan instead of registering a die with an invented budget. */
static uint64_t tt_probe_device_bytes(int device_index)
{
    ft_tt_runtime runtime;
    uint64_t total;

    if (ft_tt_runtime_open(device_index, &runtime) != FT_TT_OK)
        return 0;
    /* ft_tt_runtime_open() already refuses a topology ft_tt_coord cannot map,
     * on both transports, so a card we would not dare address never gets here
     * (spec §7.3). */
    total = ft_tt_runtime_dram_bytes(&runtime);
    ft_tt_runtime_close(&runtime);
    return total;
}

int pn_dm_init_tt(int n_devices)
{
    if (n_devices <= 0 || n_devices > PN_DM_MAX_DIES)
        return -1;
    pthread_mutex_lock(&g_dm.lock);
    if (g_dm.running) {                     /* idempotent while running */
        int n = g_dm.n_dies;
        pthread_mutex_unlock(&g_dm.lock);
        return n;
    }
    memset(g_dm.die, 0, sizeof(g_dm.die));
    memset(g_dm.binding, 0, sizeof(g_dm.binding));
    int found = 0;
    uint64_t reserve = pn_die_reserve_bytes();
    for (int i = 0; i < n_devices; i++) {
        uint64_t total = tt_probe_device_bytes(i);
        if (total == 0)
            break;                          /* ordinals are contiguous — stop */
        g_dm.die[found].present = true;
        g_dm.die[found].vram_budget = total > reserve ? total - reserve : 0;
        /* No external-occupancy term, unlike the CUDA probe: the card exposes
         * no cross-process allocation view, so `free` from a fresh context is
         * always `total`. A second process holding DRAM is invisible here and
         * is a reason to run one serve per card, not a number to fake. */
        g_dm.die[found].vram_used = 0;
        if (pn_memwatch_running())
            pn_memwatch_set_die(found, total, g_dm.die[found].vram_budget, 0);
        found++;
    }
    if (found == 0) {
        pthread_mutex_unlock(&g_dm.lock);
        return -1;
    }
    g_dm.n_dies = found;
    g_dm.running = true;
    pthread_mutex_unlock(&g_dm.lock);
    fprintf(stderr, "die_manager: tracking %d Tenstorrent device%s (one die each)\n",
            found, found == 1 ? "" : "s");
    return found;
}
#endif /* USE_TT */

void pn_dm_shutdown(void)
{
    pthread_mutex_lock(&g_dm.lock);
    g_dm.running = false;
    g_dm.n_dies = 0;
    memset(g_dm.die, 0, sizeof(g_dm.die));
    memset(g_dm.binding, 0, sizeof(g_dm.binding));
    pthread_mutex_unlock(&g_dm.lock);
}

/* --- queries ------------------------------------------------------------ */

int pn_dm_count(void)
{
    pthread_mutex_lock(&g_dm.lock);
    int n = g_dm.running ? g_dm.n_dies : 0;
    pthread_mutex_unlock(&g_dm.lock);
    return n;
}

/* Locked helper: does `b` cover die `id`? */
static bool binding_has_die(const dm_binding *b, int id)
{
    for (int i = 0; i < b->n_dies; i++)
        if (b->dies[i] == id)
            return true;
    return false;
}

int pn_dm_get(int die_id, pn_dm_die *out)
{
    if (!out)
        return -1;
    pthread_mutex_lock(&g_dm.lock);
    if (!g_dm.running || die_id < 0 || die_id >= g_dm.n_dies ||
        !g_dm.die[die_id].present) {
        pthread_mutex_unlock(&g_dm.lock);
        return -1;
    }
    memset(out, 0, sizeof(*out));
    out->id = die_id;
    out->vram_budget = g_dm.die[die_id].vram_budget;
    out->vram_used = g_dm.die[die_id].vram_used;
    for (int i = 0; i < PN_DM_MAX_BINDINGS; i++) {
        if (!g_dm.binding[i].used || !binding_has_die(&g_dm.binding[i], die_id))
            continue;
        if (out->n_models == 0) {
            strncpy(out->model_name, g_dm.binding[i].name,
                    PN_DM_NAME_MAX - 1);
            out->model_name[PN_DM_NAME_MAX - 1] = '\0';
        }
        out->n_models++;
    }
    pthread_mutex_unlock(&g_dm.lock);
    return 0;
}

/* --- placement ----------------------------------------------------------- */

/* Locked helpers. */
static uint64_t die_free_bytes(int id)
{
    const dm_die *d = &g_dm.die[id];
    return d->vram_budget > d->vram_used ? d->vram_budget - d->vram_used : 0;
}

static bool die_occupied(int id)
{
    return g_dm.die[id].vram_used > 0;
}

/* The card mate of a die under the fixed K80 pairing (2i <-> 2i+1); -1 when
 * the mate is outside the tracked range. */
static int die_mate(int id)
{
    int mate = (id % 2 == 0) ? id + 1 : id - 1;
    return (mate >= 0 && mate < g_dm.n_dies) ? mate : -1;
}

int pn_dm_place(uint64_t need_bytes, int n_dies_wanted, int *dies_out)
{
    if (!dies_out || n_dies_wanted < 1 || n_dies_wanted > 2)
        return -1;   /* >2-die pipeline placement arrives with Phase 6.1 */

    pthread_mutex_lock(&g_dm.lock);
    if (!g_dm.running) {
        pthread_mutex_unlock(&g_dm.lock);
        return -1;
    }
    uint64_t per_die = need_bytes / (uint64_t)n_dies_wanted;
    int rc = -1;

    if (n_dies_wanted == 2) {
        /* Same-card pair with room on both halves. Prefer a fully-free pair;
         * tie-break on the most free space (least loaded). */
        int best = -1;
        bool best_free = false;
        uint64_t best_room = 0;
        for (int a = 0; a + 1 < g_dm.n_dies; a += 2) {
            int b = a + 1;
            if (!g_dm.die[a].present || !g_dm.die[b].present)
                continue;
            if (die_free_bytes(a) < per_die || die_free_bytes(b) < per_die)
                continue;
            bool pair_free = !die_occupied(a) && !die_occupied(b);
            uint64_t room = die_free_bytes(a) + die_free_bytes(b);
            if (best < 0 || (pair_free && !best_free) ||
                (pair_free == best_free && room > best_room)) {
                best = a;
                best_free = pair_free;
                best_room = room;
            }
        }
        if (best >= 0) {
            dies_out[0] = best;
            dies_out[1] = best + 1;
            rc = 0;
        }
    } else {
        /* Single die: pack onto a card whose mate is already occupied so free
         * same-card pairs stay available for 2-die models; among candidates
         * take the least-loaded (most free). */
        int best = -1;
        bool best_packed = false;
        uint64_t best_room = 0;
        for (int i = 0; i < g_dm.n_dies; i++) {
            if (!g_dm.die[i].present || die_free_bytes(i) < per_die)
                continue;
            int mate = die_mate(i);
            bool packed = (mate >= 0 && die_occupied(mate)) || die_occupied(i);
            uint64_t room = die_free_bytes(i);
            if (best < 0 || (packed && !best_packed) ||
                (packed == best_packed && room > best_room)) {
                best = i;
                best_packed = packed;
                best_room = room;
            }
        }
        if (best >= 0) {
            dies_out[0] = best;
            rc = 0;
        }
    }
    pthread_mutex_unlock(&g_dm.lock);
    return rc;
}

/* --- bind / release ------------------------------------------------------ */

int pn_dm_bind(const int *dies, int n, const char *model_name,
               uint64_t total_bytes)
{
    if (!dies || n < 1 || n > PN_DM_MAX_DIES || !model_name || !*model_name)
        return -1;

    pthread_mutex_lock(&g_dm.lock);
    if (!g_dm.running)
        goto refuse;

    /* Same-name multi-card placement: one model id may be bound to several
     * DISJOINT die-sets (one binding per card) for load-balanced serving.
     * Overlapping dies are still rejected by the per-die room check below, so
     * here we only need a free binding slot — no duplicate-name refusal. Each
     * placement gets its own binding; the registry releases one per entry on
     * unload (pn_dm_release is by-name, first-match — N entries → N releases). */
    dm_binding *slot = NULL;
    for (int i = 0; i < PN_DM_MAX_BINDINGS; i++) {
        if (!g_dm.binding[i].used) { slot = &g_dm.binding[i]; break; }
    }
    if (!slot)
        goto refuse;                             /* binding table full */

    uint64_t per_die = total_bytes / (uint64_t)n;
    for (int i = 0; i < n; i++) {
        int id = dies[i];
        if (id < 0 || id >= g_dm.n_dies || !g_dm.die[id].present ||
            die_free_bytes(id) < per_die)
            goto refuse;                         /* unknown die / no room */
    }

    slot->used = true;
    strncpy(slot->name, model_name, PN_DM_NAME_MAX - 1);
    slot->name[PN_DM_NAME_MAX - 1] = '\0';
    slot->n_dies = n;
    slot->bytes_per_die = per_die;
    for (int i = 0; i < n; i++) {
        slot->dies[i] = dies[i];
        g_dm.die[dies[i]].vram_used += per_die;
    }
    pthread_mutex_unlock(&g_dm.lock);
    return 0;

refuse:
    pthread_mutex_unlock(&g_dm.lock);
    return -1;
}

int pn_dm_release(const char *model_name)
{
    if (!model_name)
        return -1;
    pthread_mutex_lock(&g_dm.lock);
    for (int i = 0; i < PN_DM_MAX_BINDINGS; i++) {
        dm_binding *b = &g_dm.binding[i];
        if (!b->used || strncmp(b->name, model_name, PN_DM_NAME_MAX) != 0)
            continue;
        for (int j = 0; j < b->n_dies; j++) {
            dm_die *d = &g_dm.die[b->dies[j]];
            d->vram_used = d->vram_used > b->bytes_per_die
                         ? d->vram_used - b->bytes_per_die : 0;
        }
        memset(b, 0, sizeof(*b));
        pthread_mutex_unlock(&g_dm.lock);
        return 0;
    }
    pthread_mutex_unlock(&g_dm.lock);
    return -1;
}