/*
 * MemConstrainedReader experiment.
 *
 * Given a memory limit and the table schema, automatically derives the
 * batch_size to use for each column count:
 *
 *   batch_size = floor( (M_limit - M_fixed - N_cols * C_page) / s_row )
 *
 * where
 *   N_cols  = total queried columns (TAG + FIELD)
 *   C_page  = decompression buffer per column ≈ page_writer_max_memory_bytes_
 *   s_row   = 8 (timestamp) + sum sizeof(col_type) for all queried columns
 *   M_fixed = estimated fixed overhead (LRU cache, decoder objects, etc.)
 *
 * For each (N_field_cols, derived_batch_size) pair the experiment:
 *   1. Queries the tsfile with the derived batch_size.
 *   2. Records peak memory via ModStat (if ENABLE_MEM_STAT is defined).
 *   3. Writes results to CSV for comparison with the formula prediction.
 *
 * Prerequisite: a tsfile generated by the write_memory or writer_bench
 *               experiment (schema: table "mem_bench", 8 FIELD columns).
 *
 * Build:
 *   cmake -DENABLE_MEM_STAT=ON -DBUILD_TEST=OFF ..
 *   cmake --build . --target reader_bench
 *
 * Usage:
 *   ./reader_bench [tsfile_path] [mem_limit_mb] [m_fixed_mb] [csv_path]
 *
 * Default:
 *   tsfile_path  = /Users/colin/dev/tsfile_b1/cpp/experiment/experiment.tsfile
 *   mem_limit    = 16 MB
 *   m_fixed      = 4 MB
 *   csv_path     = reader_bench.csv
 */

#include <sys/resource.h>
#include <sys/time.h>

#include <chrono>
#include <cstdint>
#include <cstdio>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <string>
#include <vector>

#include "common/allocator/alloc_base.h"
#include "common/config/config.h"
#include "common/tsblock/tsblock.h"
#include "reader/table_result_set.h"
#include "reader/tsfile_reader.h"

using namespace storage;
using namespace common;

// ---------------------------------------------------------------------------
// Schema for "mem_bench" table
// ---------------------------------------------------------------------------
//  TAG  columns: id1 (STRING), id2 (STRING)
//  FIELD columns (in order): s1(INT32), s2(INT32), s3(INT64), s4(INT64),
//                             s5(FLOAT), s6(FLOAT), s7(DOUBLE), s8(DOUBLE)
// ---------------------------------------------------------------------------
static const char* kTable = "mem_bench";
static const char* kTagCols[] = {"id1", "id2"};
static const int kNumTags = 2;
static const char* kFieldCols[] = {"s1", "s2", "s3", "s4",
                                   "s5", "s6", "s7", "s8"};
// kNumFields = 8 (unused directly; used as upper bound for field_counts[])
static const int kMaxFields = 8;

// sizeof of each field column type (bytes)
static const int kFieldSizes[] = {4, 4, 8, 8, 4, 4, 8, 8};

// sizeof of TAG column type.  STRING stored as pointer/index in TsBlock: use 8.
static const int kTagSize = 8;

// Default C_page: one page's decompression buffer per column.
// Matches page_writer_max_memory_bytes_ (128 KB default).
static const int64_t kDefaultCPage = 128LL * 1024;

// ---------------------------------------------------------------------------
// Derived batch size
// ---------------------------------------------------------------------------
struct ReadParams {
    int n_field_cols;  // FIELD columns queried
    int n_total_cols;  // TAG + FIELD + 1 time col (for s_row)
    int64_t s_row;     // bytes per row in TsBlock
    int64_t c_page;    // decompression buffer per column
    int64_t m_fixed;   // fixed overhead (user supplied)
    int64_t mem_limit;
    int64_t m_data_budget;  // mem_limit - m_fixed - n_total_cols * c_page
    int32_t batch_size;     // derived
};

static ReadParams derive_params(int n_field_cols, int64_t mem_limit,
                                int64_t m_fixed, int64_t c_page) {
    ReadParams p;
    p.n_field_cols = n_field_cols;
    // total columns seen in the query = 2 TAGs + n_field FIELDs
    // (timestamp is implicit in TsBlock but counts for decompression)
    p.n_total_cols = kNumTags + n_field_cols;
    p.c_page = c_page;
    p.m_fixed = m_fixed;
    p.mem_limit = mem_limit;

    // s_row = 8 (timestamp) + 2*kTagSize + sum(field_sizes[0..n_field-1])
    p.s_row = 8 + static_cast<int64_t>(kNumTags) * kTagSize;
    for (int i = 0; i < n_field_cols; i++) p.s_row += kFieldSizes[i];

    p.m_data_budget =
        mem_limit - m_fixed - static_cast<int64_t>(p.n_total_cols) * c_page;
    if (p.m_data_budget <= 0) {
        p.batch_size = 1;
    } else {
        int64_t bs = p.m_data_budget / p.s_row;
        p.batch_size = static_cast<int32_t>(bs > INT32_MAX ? INT32_MAX : bs);
        if (p.batch_size < 1) p.batch_size = 1;
    }
    return p;
}

// ---------------------------------------------------------------------------
// ModStat helpers (compile-time guarded)
// ---------------------------------------------------------------------------
#ifdef ENABLE_MEM_STAT
static const int kModCount = __LAST_MOD_ID;

static int64_t total_modstat() {
    auto& ms = ModStat::get_instance();
    int64_t t = 0;
    for (int i = 0; i < kModCount; i++) t += ms.get_stat(i);
    return t;
}

// Record baseline before a query and track peak increment above it.
static int64_t g_baseline = 0;
static int64_t g_peak_above_baseline = 0;

static void begin_query_mem() {
    g_baseline = total_modstat();
    g_peak_above_baseline = 0;
}

static void update_peak() {
    int64_t delta = total_modstat() - g_baseline;
    if (delta > g_peak_above_baseline) g_peak_above_baseline = delta;
}

static int64_t query_peak_mem() { return g_peak_above_baseline; }
#endif

// ---------------------------------------------------------------------------
// CSV header / row
// ---------------------------------------------------------------------------
static void write_csv_header(std::ofstream& csv) {
    csv << "n_field_cols,n_total_cols,s_row,c_page,m_fixed_mb,"
           "mem_limit_mb,m_data_budget,"
           "batch_size_derived,"
           "rows_read,"
           "peak_memory_actual,"
           "m_data_formula,"
           "modstat_available\n";
}

static void write_csv_row(std::ofstream& csv, const ReadParams& p,
                          int64_t rows_read, int64_t peak_actual) {
    bool has_modstat = false;
#ifdef ENABLE_MEM_STAT
    has_modstat = true;
#endif
    // Formula prediction for TsBlock buffer at batch_size rows
    int64_t m_data_formula = static_cast<int64_t>(p.batch_size) * p.s_row;

    csv << p.n_field_cols << "," << p.n_total_cols << "," << p.s_row << ","
        << p.c_page << "," << (p.m_fixed / (1024 * 1024)) << ","
        << (p.mem_limit / (1024 * 1024)) << "," << p.m_data_budget << ","
        << p.batch_size << "," << rows_read << "," << peak_actual << ","
        << m_data_formula << "," << (has_modstat ? "yes" : "no") << "\n";
}

// ---------------------------------------------------------------------------
// main
// ---------------------------------------------------------------------------
int main(int argc, char* argv[]) {
    std::string tsfile_path =
        "/Users/colin/dev/tsfile_b1/cpp/experiment/experiment.tsfile";
    int64_t mem_limit_mb = 16;
    int64_t m_fixed_mb = 4;
    std::string csv_path = "reader_bench.csv";

    if (argc > 1) tsfile_path = argv[1];
    if (argc > 2) mem_limit_mb = std::atoll(argv[2]);
    if (argc > 3) m_fixed_mb = std::atoll(argv[3]);
    if (argc > 4) csv_path = argv[4];

    int64_t mem_limit = mem_limit_mb * 1024LL * 1024;
    int64_t m_fixed = m_fixed_mb * 1024LL * 1024;
    int64_t c_page = kDefaultCPage;

    libtsfile_init();

#ifdef ENABLE_MEM_STAT
    std::cout << "ModStat: ENABLED (actual peak memory will be measured)\n";
#else
    std::cout << "ModStat: DISABLED (build with -DENABLE_MEM_STAT=ON for "
                 "actual measurement)\n";
#endif

    std::cout << "=== MemConstrainedReader Experiment ===\n"
              << "  tsfile:      " << tsfile_path << "\n"
              << "  mem_limit:   " << mem_limit_mb << " MB\n"
              << "  m_fixed:     " << m_fixed_mb << " MB\n"
              << "  c_page:      " << c_page / 1024 << " KB\n"
              << "  csv:         " << csv_path << "\n\n";

    // Print derived parameters for all field column counts
    std::cout << std::setw(12) << "N_field" << std::setw(12) << "N_total"
              << std::setw(10) << "s_row" << std::setw(14) << "data_budget"
              << std::setw(14) << "batch_size\n";
    std::cout << std::string(62, '-') << "\n";

    int field_counts[] = {1, 2, 4, 6, 8};
    for (int n : field_counts) {
        auto p = derive_params(n, mem_limit, m_fixed, c_page);
        std::cout << std::setw(12) << p.n_field_cols << std::setw(12)
                  << p.n_total_cols << std::setw(10) << p.s_row << std::setw(14)
                  << p.m_data_budget << std::setw(14) << p.batch_size << "\n";
    }
    std::cout << "\n";

    std::ofstream csv(csv_path);
    if (!csv.is_open()) {
        std::cerr << "cannot open csv: " << csv_path << "\n";
        return 1;
    }
    write_csv_header(csv);

    // The file covers timestamps 0 .. 200M-1 (write_memory default)
    // Use INT64 max to read all rows.
    int64_t t_start = 0;
    int64_t t_end = INT64_MAX;

    for (int n_field : field_counts) {
        auto p = derive_params(n_field, mem_limit, m_fixed, c_page);

        // Build column list: tag cols + first n_field FIELD cols
        std::vector<std::string> cols;
        for (int t = 0; t < kNumTags; t++) cols.push_back(kTagCols[t]);
        for (int f = 0; f < n_field && f < kMaxFields; f++)
            cols.push_back(kFieldCols[f]);

        std::cout << "  [cols=" << std::setw(2) << n_field
                  << " batch=" << std::setw(7) << p.batch_size << "] opening..."
                  << std::flush;

        TsFileReader reader;
        if (reader.open(tsfile_path) != 0) {
            std::cerr << "\nFailed to open: " << tsfile_path << "\n";
            return 1;
        }

#ifdef ENABLE_MEM_STAT
        begin_query_mem();
#endif

        ResultSet* rs = nullptr;
        int ret = reader.query(kTable, cols, t_start, t_end, rs, p.batch_size);
        if (ret != 0 || rs == nullptr) {
            std::cerr << "\nquery failed: " << ret << "\n";
            reader.close();
            continue;
        }

        std::cout << " reading..." << std::flush;

        TsBlock* block = nullptr;
        int64_t rows = 0;
        while (rs->get_next_tsblock(block) == E_OK && block) {
            rows += block->get_row_count();
#ifdef ENABLE_MEM_STAT
            update_peak();
#endif
        }

#ifdef ENABLE_MEM_STAT
        update_peak();
        int64_t peak_actual = query_peak_mem();
#else
        int64_t peak_actual = -1;  // not available
#endif

        rs->close();
        reader.close();

        write_csv_row(csv, p, rows, peak_actual);

        std::cout << " rows=" << rows << " peak="
                  << (peak_actual >= 0
                          ? std::to_string(peak_actual / (1024 * 1024)) + " MB"
                          : "N/A")
                  << "\n";
    }

    csv.close();
    std::cout << "\nCSV: " << csv_path << "\n";

    libtsfile_destroy();
    return 0;
}