/*
 * MSAC Trace - Export MSAC symbol decoding data for hardware verification
 */

#define _POSIX_C_SOURCE 200809L
#include "common/msac_trace.h"
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <errno.h>

// Global MSAC trace context
static MsacTraceCtx g_msac_trace_ctx = {0};

MsacTraceCtx *msac_trace_get_global_ctx(void) {
    return &g_msac_trace_ctx;
}

void msac_trace_init(MsacTraceCtx *ctx, const char *output_dir) {
    memset(ctx, 0, sizeof(*ctx));
    if (output_dir) {
        // Create output directory if it doesn't exist
        if (mkdir(output_dir, 0755) != 0 && errno != EEXIST) {
            fprintf(stderr, "Warning: Could not create MSAC trace output directory: %s\n", output_dir);
        }
        ctx->output_dir = strdup(output_dir);
        ctx->enabled = 1;
    }
}

void msac_trace_close(MsacTraceCtx *ctx) {
    if (ctx->fp) {
        fclose(ctx->fp);
        ctx->fp = NULL;
    }
    if (ctx->output_dir) {
        free(ctx->output_dir);
        ctx->output_dir = NULL;
    }
    ctx->enabled = 0;
}

void msac_trace_start_frame(MsacTraceCtx *ctx, int frame_num) {
    if (!ctx->enabled) return;

    // Close previous tile file if open
    if (ctx->fp) {
        fclose(ctx->fp);
        ctx->fp = NULL;
    }

    ctx->frame_num = frame_num;
    ctx->tile_row = -1;
    ctx->tile_col = -1;
    ctx->symbol_count = 0;
}

void msac_trace_end_frame(MsacTraceCtx *ctx) {
    if (!ctx->enabled) return;

    if (ctx->fp) {
        fclose(ctx->fp);
        ctx->fp = NULL;
    }

    fprintf(stderr, "MSAC trace: frame %d complete, %d symbols traced\n",
            ctx->frame_num, ctx->symbol_count);
}

void msac_trace_start_tile(MsacTraceCtx *ctx, int tile_row, int tile_col) {
    if (!ctx->enabled) return;

    // Close previous tile file if open
    if (ctx->fp) {
        fclose(ctx->fp);
        ctx->fp = NULL;
    }

    ctx->tile_row = tile_row;
    ctx->tile_col = tile_col;

    // Open new file for this tile
    char filename[512];
    snprintf(filename, sizeof(filename), "%s/frame_%04d_tile_%02d_%02d_msac.jsonl",
             ctx->output_dir, ctx->frame_num, tile_row, tile_col);
    ctx->fp = fopen(filename, "w");
    if (!ctx->fp) {
        fprintf(stderr, "Warning: Could not open MSAC trace file: %s\n", filename);
        return;
    }
}

void msac_trace_end_tile(MsacTraceCtx *ctx) {
    if (!ctx->enabled || !ctx->fp) return;

    fclose(ctx->fp);
    ctx->fp = NULL;
}

void msac_trace_set_sb(MsacTraceCtx *ctx, int sb_row, int sb_col) {
    if (!ctx->enabled) return;
    ctx->sb_row = sb_row;
    ctx->sb_col = sb_col;
}

void msac_trace_set_mi(MsacTraceCtx *ctx, int mi_row, int mi_col) {
    if (!ctx->enabled) return;
    ctx->mi_row = mi_row;
    ctx->mi_col = mi_col;
}

// Helper to write CDF array as JSON
static void write_cdf_json(FILE *fp, const uint16_t *cdf, int nsymbs) {
    fprintf(fp, "[");
    for (int i = 0; i < nsymbs; i++) {
        if (i > 0) fprintf(fp, ",");
        fprintf(fp, "%u", cdf[i]);
    }
    fprintf(fp, "]");
}

void msac_trace_symbol(MsacTraceCtx *ctx,
                        const char *name,
                        uint32_t peek32_before,
                        uint16_t rng_before,
                        int16_t cnt,
                        const uint16_t *cdf,
                        int nsymbs,
                        int value,
                        uint16_t rng_after,
                        int shift) {
    if (!ctx->enabled || !ctx->fp) return;

    // Calculate thresholds from CDF (for hardware comparison)
    // threshold[i] = (rng >> 8) * (cdf[i] >> EC_PROB_SHIFT) >> (7 - EC_PROB_SHIFT) + EC_MIN_PROB * (N - i)
    // But we'll just output the raw CDF and let the verification tool calculate

    fprintf(ctx->fp, "{\"name\":\"%s\"", name);
    fprintf(ctx->fp, ",\"peek32\":\"0x%08X\"", peek32_before);
    fprintf(ctx->fp, ",\"rng_before\":%u", rng_before);
    fprintf(ctx->fp, ",\"cnt\":%d", cnt);
    fprintf(ctx->fp, ",\"nsymbs\":%d", nsymbs);

    if (cdf != NULL) {
        fprintf(ctx->fp, ",\"cdf\":");
        write_cdf_json(ctx->fp, cdf, nsymbs);
    }

    fprintf(ctx->fp, ",\"value\":%d", value);
    fprintf(ctx->fp, ",\"rng_after\":%u", rng_after);
    fprintf(ctx->fp, ",\"shift\":%d", shift);

    // Position context
    fprintf(ctx->fp, ",\"sb_row\":%d", ctx->sb_row);
    fprintf(ctx->fp, ",\"sb_col\":%d", ctx->sb_col);
    fprintf(ctx->fp, ",\"mi_row\":%d", ctx->mi_row);
    fprintf(ctx->fp, ",\"mi_col\":%d", ctx->mi_col);

    fprintf(ctx->fp, "}\n");

    ctx->symbol_count++;
}

// Separate before/after calls (for when we need to trace in two steps)
static struct {
    const char *name;
    uint32_t peek32;
    uint16_t rng;
    int16_t cnt;
    uint16_t cdf[MSAC_TRACE_MAX_CDF];
    int nsymbs;
    int has_pending;
} pending_trace = {0};

void msac_trace_before_symbol(MsacTraceCtx *ctx,
                               const char *name,
                               uint32_t peek32,
                               uint16_t rng,
                               int16_t cnt,
                               const uint16_t *cdf,
                               int nsymbs) {
    if (!ctx->enabled) return;

    pending_trace.name = name;
    pending_trace.peek32 = peek32;
    pending_trace.rng = rng;
    pending_trace.cnt = cnt;
    pending_trace.nsymbs = nsymbs;

    if (cdf != NULL && nsymbs <= MSAC_TRACE_MAX_CDF) {
        for (int i = 0; i < nsymbs; i++) {
            pending_trace.cdf[i] = cdf[i];
        }
    }

    pending_trace.has_pending = 1;
}

void msac_trace_after_symbol(MsacTraceCtx *ctx,
                              int value,
                              uint16_t new_rng,
                              int shift) {
    if (!ctx->enabled || !pending_trace.has_pending) return;

    msac_trace_symbol(ctx,
                       pending_trace.name,
                       pending_trace.peek32,
                       pending_trace.rng,
                       pending_trace.cnt,
                       pending_trace.cdf,
                       pending_trace.nsymbs,
                       value,
                       new_rng,
                       shift);

    pending_trace.has_pending = 0;
}
