/*
 * CDQ Dump - Export coefficient data in CDQ format
 * For VPU hardware verification
 *
 * Per-tile output with 36-bit format, flags at [3:0]
 */

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

// Global CDQ context for simple integration
static CdqDumpCtx g_cdq_ctx = {0};

CdqDumpCtx *cdq_get_global_ctx(void) {
    return &g_cdq_ctx;
}

// Flag names for JSONL output
static const char *flag_to_str(int flags) {
    switch (flags) {
        case 0x0: return "COEF";
        case 0x1: return "SOU";
        case 0x2: return "SOT";
        case 0x4: return "COEF+EOT";
        case 0x6: return "SOT+EOT";
        case 0x9: return "SOU+EOU";
        case 0xC: return "COEF+EOT+EOU";
        case 0xE: return "SOT+EOT+EOU";
        default:  return "???";
    }
}

// Emit SOU word (internal helper)
static void emit_sou(CdqDumpCtx *ctx, int sb_row, int sb_col, int add_eou) {
    if (!ctx->fp) return;

    int flags = CDQ_FLAG_SOU;
    if (add_eou) flags |= CDQ_FLAG_EOU;

    uint64_t word = cdq_build_sou(flags, ctx->bit_depth, sb_row, sb_col, ctx->mode, ctx->sb_size);

    fprintf(ctx->fp, "{\"data\":\"0x%09llX\",\"ctrl\":\"%s\",\"bit_depth\":%d,\"sb_row\":%d,\"sb_col\":%d,\"mode\":%d,\"sb_size\":%d}\n",
            (unsigned long long)word, flag_to_str(flags), ctx->bit_depth, sb_row, sb_col, ctx->mode, ctx->sb_size);
}

// Flush pending data element - SOT or COEF (internal)
static void flush_pending_data(CdqDumpCtx *ctx, int add_eou) {
    if (!ctx->fp || ctx->pending_type == CDQ_PENDING_NONE) return;

    int flags = ctx->pending_flags;
    if (add_eou) flags |= CDQ_FLAG_EOU;

    if (ctx->pending_type == CDQ_PENDING_SOT) {
        uint64_t word = cdq_build_sot(flags, ctx->pending_plane_id, ctx->pending_dc_qidx,
                                       ctx->pending_ac_qidx, ctx->pending_tx_type, ctx->pending_tx_size);
        fprintf(ctx->fp, "{\"data\":\"0x%09llX\",\"ctrl\":\"%s\",\"plane_id\":%d,\"dc_qidx\":%d,\"ac_qidx\":%d,\"tx_type\":%d,\"tx_size\":%d}\n",
                (unsigned long long)word, flag_to_str(flags), ctx->pending_plane_id,
                ctx->pending_dc_qidx, ctx->pending_ac_qidx, ctx->pending_tx_type, ctx->pending_tx_size);
    } else if (ctx->pending_type == CDQ_PENDING_COEF) {
        uint64_t word = cdq_build_coef(flags, ctx->pending_scan_pos, ctx->pending_coef);
        fprintf(ctx->fp, "{\"data\":\"0x%09llX\",\"ctrl\":\"%s\",\"scan_pos\":%d,\"coef\":%d}\n",
                (unsigned long long)word, flag_to_str(flags), ctx->pending_scan_pos, (int)ctx->pending_coef);
    }

    ctx->pending_type = CDQ_PENDING_NONE;
}

// Flush pending SOU and data - called when SB ends
static void flush_pending_sou(CdqDumpCtx *ctx) {
    if (!ctx->fp) return;

    // If SOU is still pending (no data was emitted for this SB)
    if (ctx->pending_sou) {
        if (ctx->pending_type != CDQ_PENDING_NONE) {
            // This shouldn't happen - if there's pending data, SOU should have been emitted
            // But handle it anyway
            emit_sou(ctx, ctx->pending_sou_row, ctx->pending_sou_col, 0);
            flush_pending_data(ctx, 1);  // add EOU to data
        } else {
            // No data in this SB - emit SOU+EOU (skip SB)
            emit_sou(ctx, ctx->pending_sou_row, ctx->pending_sou_col, 1);
        }
        ctx->pending_sou = 0;
    } else {
        // SOU already emitted - just flush pending data with EOU
        flush_pending_data(ctx, 1);
    }
}

// Ensure SOU is emitted before data
static void ensure_sou_emitted(CdqDumpCtx *ctx) {
    if (!ctx->fp || !ctx->pending_sou) return;

    // First flush any previous pending data (without EOU - belongs to prev SB)
    // Actually this shouldn't happen - pending data should be flushed in flush_pending_sou
    // But just in case...
    flush_pending_data(ctx, 0);

    // Emit SOU (without EOU - this SB has data)
    emit_sou(ctx, ctx->pending_sou_row, ctx->pending_sou_col, 0);
    ctx->pending_sou = 0;
}

void cdq_dump_init(CdqDumpCtx *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 CDQ output directory: %s\n", output_dir);
        }
        ctx->output_dir = strdup(output_dir);
        ctx->enabled = 1;
    }
}

void cdq_dump_close(CdqDumpCtx *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 cdq_dump_start_frame(CdqDumpCtx *ctx, int frame_num, int mode, int sb_size, int bit_depth) {
    if (!ctx->enabled) return;

    // Close previous 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->sb_count = 0;
    ctx->block_count = 0;
    ctx->coef_count = 0;
    ctx->sb_size = sb_size;
    ctx->mode = mode;
    ctx->bit_depth = bit_depth;
    ctx->pending_sou = 0;
    ctx->pending_type = CDQ_PENDING_NONE;
}

void cdq_dump_end_frame(CdqDumpCtx *ctx) {
    if (!ctx->enabled) return;

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

    fprintf(stderr, "CDQ dump: frame %d complete, %d superblocks, %d blocks, %d coefficients\n",
            ctx->frame_num, ctx->sb_count, ctx->block_count, ctx->coef_count);
}

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

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

    ctx->tile_row = tile_row;
    ctx->tile_col = tile_col;
    ctx->pending_sou = 0;
    ctx->pending_type = CDQ_PENDING_NONE;

    // Open new file for this tile (same naming as MSAC trace)
    char filename[512];
    snprintf(filename, sizeof(filename), "%s/frame_%04d_tile_%02d_%02d_cdq.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 CDQ dump file: %s\n", filename);
        return;
    }
}

void cdq_dump_end_tile(CdqDumpCtx *ctx) {
    if (!ctx->enabled || !ctx->fp) return;

    // Flush last SB of this tile
    flush_pending_sou(ctx);

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

void cdq_dump_start_sb(CdqDumpCtx *ctx, int sb_row, int sb_col) {
    if (!ctx->enabled || !ctx->fp) return;

    // Flush previous SB (will emit SOU+EOU if it was empty, or add EOU to last data)
    flush_pending_sou(ctx);

    // Store new SB as pending
    ctx->pending_sou = 1;
    ctx->pending_sou_row = sb_row;
    ctx->pending_sou_col = sb_col;
    ctx->sb_count++;
}

void cdq_dump_sot(CdqDumpCtx *ctx, int plane_id, int dc_qidx, int ac_qidx,
                  int tx_type, int tx_size, int is_skip) {
    if (!ctx->enabled || !ctx->fp) return;

    // Make sure SOU is emitted first
    ensure_sou_emitted(ctx);

    // Flush previous pending data element (without EOU - not end of SB yet)
    flush_pending_data(ctx, 0);

    int flags = CDQ_FLAG_SOT;
    if (is_skip) flags |= CDQ_FLAG_EOT;

    // Store as pending
    ctx->pending_type = CDQ_PENDING_SOT;
    ctx->pending_flags = flags;
    ctx->pending_plane_id = plane_id;
    ctx->pending_dc_qidx = dc_qidx;
    ctx->pending_ac_qidx = ac_qidx;
    ctx->pending_tx_type = tx_type;
    ctx->pending_tx_size = tx_size;

    ctx->block_count++;
}

void cdq_dump_coef(CdqDumpCtx *ctx, int scan_pos, int32_t coef, int is_eot) {
    if (!ctx->enabled || !ctx->fp) return;

    // Make sure SOU is emitted first
    ensure_sou_emitted(ctx);

    // Flush previous pending data element (without EOU)
    flush_pending_data(ctx, 0);

    int flags = 0;
    if (is_eot) flags |= CDQ_FLAG_EOT;

    // Store as pending
    ctx->pending_type = CDQ_PENDING_COEF;
    ctx->pending_flags = flags;
    ctx->pending_scan_pos = scan_pos;
    ctx->pending_coef = coef;

    ctx->coef_count++;
}
