/*
 * Copyright (c) 2016, Alliance for Open Media. All rights reserved.
 *
 * This source code is subject to the terms of the BSD 2 Clause License and
 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
 * was not distributed with this source code in the LICENSE file, you can
 * obtain it at www.aomedia.org/license/software. If the Alliance for Open
 * Media Patent License 1.0 was not distributed with this source code in the
 * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
 */

#ifndef AOM_AOM_DSP_BITREADER_H_
#define AOM_AOM_DSP_BITREADER_H_

#include <assert.h>
#include <limits.h>
#include <stdio.h>  // For trace file
#include <string.h> // For strncpy

#include "config/aom_config.h"

#include "aom/aomdx.h"
#include "aom/aom_integer.h"
#include "aom_dsp/entdec.h"
#include "aom_dsp/odintrin.h"
#include "aom_dsp/prob.h"

// --- HARDCORE REFACTOR START ---
// Force trace parameter everywhere
#define ACCT_STR_NAME acct_str
#define ACCT_STR_PARAM , const char *ACCT_STR_NAME
#define ACCT_STR_ARG(s) , s
// --- HARDCORE REFACTOR END ---

#define aom_read(r, prob, ACCT_STR_NAME) \
  aom_read_(r, prob ACCT_STR_ARG(ACCT_STR_NAME))
#define aom_read_bit(r, ACCT_STR_NAME) \
  aom_read_bit_(r ACCT_STR_ARG(ACCT_STR_NAME))
#define aom_read_tree(r, tree, probs, ACCT_STR_NAME) \
  aom_read_tree_(r, tree, probs ACCT_STR_ARG(ACCT_STR_NAME))
#define aom_read_literal(r, bits, ACCT_STR_NAME) \
  aom_read_literal_(r, bits ACCT_STR_ARG(ACCT_STR_NAME))
#define aom_read_cdf(r, cdf, nsymbs, ACCT_STR_NAME) \
  aom_read_cdf_(r, cdf, nsymbs ACCT_STR_ARG(ACCT_STR_NAME))
#define aom_read_symbol(r, cdf, nsymbs, ACCT_STR_NAME) \
  aom_read_symbol_(r, cdf, nsymbs ACCT_STR_ARG(ACCT_STR_NAME))

#ifdef __cplusplus
extern "C" {
#endif

struct aom_reader {
  const uint8_t *buffer;
  const uint8_t *buffer_end;
  od_ec_dec ec;
#if CONFIG_ACCOUNTING
  struct Accounting *accounting;
#endif
  uint8_t allow_update_cdf;
};

typedef struct aom_reader aom_reader;

int aom_reader_init(aom_reader *r, const uint8_t *buffer, size_t size);

const uint8_t *aom_reader_find_begin(aom_reader *r);

const uint8_t *aom_reader_find_end(aom_reader *r);

int aom_reader_has_overflowed(const aom_reader *r);

uint32_t aom_reader_tell(const aom_reader *r);

uint32_t aom_reader_tell_frac(const aom_reader *r);

// --- MSAC TRACE (JSONL per tile) ---
// Global state for MSAC trace (defined in bitreader.c)
extern FILE* g_msac_trace_fp;
extern int g_msac_frame_num;
extern int g_msac_tile_row;
extern int g_msac_tile_col;
extern int g_msac_sb_row;
extern int g_msac_sb_col;
extern char g_msac_output_dir[256];
extern int g_msac_trace_enabled;
extern int g_msac_count_only;
extern uint64_t g_msac_symbol_count;

static inline void aom_msac_trace_set_output_dir(const char *dir) {
    if (dir) {
        strncpy(g_msac_output_dir, dir, sizeof(g_msac_output_dir) - 1);
        g_msac_output_dir[sizeof(g_msac_output_dir) - 1] = '\0';
    }
}

static inline void aom_msac_trace_start_tile(int frame_num, int tile_row, int tile_col) {
    if (!g_msac_trace_enabled) return;

    // Close previous tile file
    if (g_msac_trace_fp) {
        fclose(g_msac_trace_fp);
        g_msac_trace_fp = NULL;
    }

    g_msac_frame_num = frame_num;
    g_msac_tile_row = tile_row;
    g_msac_tile_col = tile_col;

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

static inline void aom_msac_trace_end_tile(void) {
    if (g_msac_trace_fp) {
        fclose(g_msac_trace_fp);
        g_msac_trace_fp = NULL;
    }
}

// Dump raw tile payload bytes to .bin file and metadata to .json
static inline void aom_msac_trace_dump_tile_payload(int frame_num,
                                                     int tile_row, int tile_col,
                                                     const uint8_t *data,
                                                     size_t size) {
    if (!g_msac_trace_enabled || g_msac_count_only) return;

    // Write .bin payload
    char filename[512];
    snprintf(filename, sizeof(filename), "%s/frame_%04d_tile_%02d_%02d.bin",
             g_msac_output_dir, frame_num, tile_row, tile_col);
    FILE *fp = fopen(filename, "wb");
    if (fp) {
        fwrite(data, 1, size, fp);
        fclose(fp);
    }

    // Write .json metadata
    snprintf(filename, sizeof(filename), "%s/frame_%04d_tile_%02d_%02d.json",
             g_msac_output_dir, frame_num, tile_row, tile_col);
    fp = fopen(filename, "w");
    if (fp) {
        fprintf(fp, "{\"frame\":%d,\"tile_row\":%d,\"tile_col\":%d,\"payload_bytes\":%zu}\n",
                frame_num, tile_row, tile_col, size);
        fclose(fp);
    }
}

static inline void aom_msac_trace_set_sb(int sb_row, int sb_col) {
    g_msac_sb_row = sb_row;
    g_msac_sb_col = sb_col;
}

// Peek next 32 bits from entropy decoder buffer
static inline uint32_t aom_msac_peek32(aom_reader *r) {
    const unsigned char *bptr = r->ec.bptr;
    const unsigned char *end = r->ec.end;
    uint32_t value = 0;
    for (int i = 0; i < 4; i++) {
        if (bptr < end) {
            value = (value << 8) | *bptr++;
        } else {
            value <<= 8;
        }
    }
    return value;
}

// Trace a decoded symbol with full state (CDF, type, pre-decode state)
// type: "ms" = multi-symbol CDF, "bo" = bool, "lt" = literal raw bits
// cdf: iCDF array (nsymbs entries, last=0), NULL for lt
// rng_before, cnt_before, dif_before: decoder state BEFORE this decode
static inline void trace_symbol_full(const char *name, const char *type,
                                      int val, int nsymbs,
                                      const uint16_t *cdf,
                                      uint16_t rng_before, int16_t cnt_before,
                                      uint32_t dif_before) {
    if (g_msac_count_only) {
        g_msac_symbol_count++;
        return;
    }
    if (!g_msac_trace_fp) return;

    fprintf(g_msac_trace_fp, "{\"name\":\"%s\",\"type\":\"%s\"",
            name ? name : "_", type);
    fprintf(g_msac_trace_fp, ",\"nsymbs\":%d", nsymbs);

    if (cdf != NULL) {
        fprintf(g_msac_trace_fp, ",\"cdf\":[");
        for (int i = 0; i < nsymbs; i++) {
            if (i > 0) fprintf(g_msac_trace_fp, ",");
            fprintf(g_msac_trace_fp, "%u", cdf[i]);
        }
        fprintf(g_msac_trace_fp, "]");
    }

    fprintf(g_msac_trace_fp, ",\"rng\":%u,\"cnt\":%d,\"dif\":\"0x%08X\"",
            rng_before, cnt_before, dif_before);
    fprintf(g_msac_trace_fp, ",\"value\":%d", val);
    fprintf(g_msac_trace_fp, ",\"sb\":[%d,%d]}\n", g_msac_sb_row, g_msac_sb_col);
}

static inline int aom_read_(aom_reader *r, int prob ACCT_STR_PARAM) {
  // Capture state BEFORE decode
  uint16_t rng0 = r->ec.rng;
  int16_t cnt0 = r->ec.cnt;
  uint32_t dif0 = (uint32_t)r->ec.dif;

  int p = (0x7FFFFF - (prob << 15) + prob) >> 8;
  int bit = od_ec_decode_bool_q15(&r->ec, p);

  if (ACCT_STR_NAME) {
    uint16_t bo_cdf[2] = { (uint16_t)p, 0 };
    trace_symbol_full(ACCT_STR_NAME, "bo", bit, 2, bo_cdf, rng0, cnt0, dif0);
  }
  return bit;
}

static inline int aom_read_bit_(aom_reader *r ACCT_STR_PARAM) {
  int ret;
  ret = aom_read(r, 128, ACCT_STR_NAME);
  return ret;
}

static inline int aom_read_literal_(aom_reader *r, int bits ACCT_STR_PARAM) {
  // Capture state BEFORE the literal decode sequence
  uint16_t rng0 = r->ec.rng;
  int16_t cnt0 = r->ec.cnt;
  uint32_t dif0 = (uint32_t)r->ec.dif;

  int literal = 0, bit;
  for (bit = bits - 1; bit >= 0; bit--) literal |= aom_read_bit(r, NULL) << bit;

  if (ACCT_STR_NAME) {
    trace_symbol_full(ACCT_STR_NAME, "lt", literal, bits, NULL, rng0, cnt0, dif0);
  }
  return literal;
}

static inline int aom_read_cdf_(aom_reader *r, const aom_cdf_prob *cdf,
                                int nsymbs ACCT_STR_PARAM) {
  // Capture state BEFORE decode
  uint16_t rng0 = r->ec.rng;
  int16_t cnt0 = r->ec.cnt;
  uint32_t dif0 = (uint32_t)r->ec.dif;

  int symb;
  assert(cdf != NULL);
  symb = od_ec_decode_cdf_q15(&r->ec, cdf, nsymbs);

  trace_symbol_full(ACCT_STR_NAME, "ms", symb, nsymbs, cdf, rng0, cnt0, dif0);
  return symb;
}

static inline int aom_read_symbol_(aom_reader *r, aom_cdf_prob *cdf,
                                   int nsymbs ACCT_STR_PARAM) {
  int ret;
  ret = aom_read_cdf(r, cdf, nsymbs, ACCT_STR_NAME);
  if (r->allow_update_cdf) update_cdf(cdf, ret, nsymbs);
  return ret;
}

#ifdef __cplusplus
}
#endif

#endif  // AOM_AOM_DSP_BITREADER_H_