/*
 * Copyright 2025, Datadog, Inc.
 * SPDX-License-Identifier: Apache-2.0
 */

#include "callTraceHashTable.h"
#include "callTraceStorage.h"
#include "counters.h"
#include "os.h"
#include "arch.h"
#include "common.h"
#include "primeProbing.h"
#include <string.h>
#include <signal.h>
#include <pthread.h>

static const u32 INITIAL_CAPACITY = 65536;  // 64K initial table size (matches upstream)
static const u32 CALL_TRACE_CHUNK = 8 * 1024 * 1024;
static const u64 OVERFLOW_TRACE_ID = 0x7fffffffffffffffULL;  // Max 64-bit signed value

// Define the sentinel value for CallTraceSample
CallTrace* const CallTraceSample::PREPARING = reinterpret_cast<CallTrace*>(1);

class LongHashTable {
private:
  LongHashTable *_prev;
  void *_padding0;
  u32 _capacity;
  u32 _padding1[15];
  volatile u32 _size;
  u32 _padding2[15];

  static size_t getSize(u32 capacity) {
    size_t size = sizeof(LongHashTable) +
                  (sizeof(u64) + sizeof(CallTraceSample)) * capacity;
    return (size + OS::page_mask) & ~OS::page_mask;
  }

public:
  LongHashTable(LongHashTable *prev = nullptr, u32 capacity = 0, bool should_clean = true) 
    : _prev(prev), _padding0(nullptr), _capacity(capacity), _size(0) {
    memset(_padding1, 0, sizeof(_padding1));
    memset(_padding2, 0, sizeof(_padding2));
    if (should_clean) {
      clear();
    }
  }

  static LongHashTable *allocate(LongHashTable *prev, u32 capacity, LinearAllocator* allocator) {
    void *memory = allocator->alloc(getSize(capacity));
    if (memory != nullptr) {
      // Use placement new to invoke constructor in-place with parameters
      // LinearAllocator doesn't zero memory like OS::safeAlloc with anon mmap
      // so we need to explicitly clear the keys and values (should_clean = true)
      LongHashTable *table = new (memory) LongHashTable(prev, capacity, true);
      return table;
    }
    return nullptr;
  }

  LongHashTable *prev() { return _prev; }
  void setPrev(LongHashTable* prev) { _prev = prev; }

  u32 capacity() { return _capacity; }

  u32 size() { return _size; }

  u32 incSize() { return __sync_add_and_fetch(&_size, 1); }

  u64 *keys() { return (u64 *)(this + 1); }

  CallTraceSample *values() { return (CallTraceSample *)(keys() + _capacity); }

  void clear() {
    memset(keys(), 0, (sizeof(u64) + sizeof(CallTraceSample)) * _capacity);
    _size = 0;
  }
};

CallTrace CallTraceHashTable::_overflow_trace(false, 1, OVERFLOW_TRACE_ID);

// Static initializer for overflow trace frame
__attribute__((constructor))
static void init_overflow_trace() {
  CallTraceHashTable::_overflow_trace.frames[0] = {BCI_ERROR, LP64_ONLY(0 COMMA) (jmethodID)"storage_overflow"};
}

CallTraceHashTable::CallTraceHashTable() : _instance_id(0), _parent_storage(nullptr), _allocator(CALL_TRACE_CHUNK) {
  // Instance ID will be set externally via setInstanceId()
  
  // Start with initial capacity, allowing expansion as needed
  _table = LongHashTable::allocate(nullptr, INITIAL_CAPACITY, &_allocator);
  _overflow = 0;
}

CallTraceHashTable::~CallTraceHashTable() {
  // LinearAllocator handles all memory cleanup automatically
  // No need to explicitly destroy tables since they're allocated from LinearAllocator
  // Note: No synchronization needed here because CallTraceStorage ensures
  // no new operations can start by nullifying storage pointers first
  _table = nullptr;
}


void CallTraceHashTable::decrementCounters() {
#ifdef COUNTERS
  // Compute and decrement the global counters for everything in this table.
  // Safe to call when (a) this is a standby/scratch table (never _active_storage,
  // so no signal-handler put() can target it), or (b) the active-table path is
  // guarded by lockAll() — both conditions are enforced by the only caller,
  // clearTableOnly().  The _prev traversal is safe because waitForRefCountToClear(this)
  // in clearTableOnly() has already drained any in-flight put() operations.
  // Use a set to deduplicate: put() may store the same CallTrace* pointer in
  // both a newer and an older table (when findCallTrace finds it in prev()),
  // but the counter was only incremented once, so we must only count it once.
  const size_t header_size = sizeof(CallTrace) - sizeof(ASGCT_CallFrame);
  long long freed_bytes = 0;
  long long freed_traces = 0;
  size_t estimated_entries = 0;
  for (LongHashTable *t = _table; t != nullptr; t = t->prev()) {
    estimated_entries += t->size();
  }
  std::unordered_set<CallTrace*> seen;
  seen.reserve(estimated_entries);
  for (LongHashTable *t = _table; t != nullptr; t = t->prev()) {
    u64 *keys = t->keys();
    CallTraceSample *values = t->values();
    u32 capacity = t->capacity();
    for (u32 slot = 0; slot < capacity; slot++) {
      if (keys[slot] != 0) {
        CallTrace *trace = values[slot].acquireTrace();
        if (trace != nullptr && trace != CallTraceSample::PREPARING) {
          if (seen.insert(trace).second) {
            freed_bytes += header_size + trace->num_frames * sizeof(ASGCT_CallFrame);
            freed_traces++;
          }
        }
      }
    }
  }
  Counters::increment(CALLTRACE_STORAGE_BYTES, -freed_bytes);
  Counters::increment(CALLTRACE_STORAGE_TRACES, -freed_traces);
#endif // COUNTERS
}

ChunkList CallTraceHashTable::clearTableOnly() {
  // Wait only for in-flight put() operations that hold a RefCountGuard on THIS
  // table.  Waiting globally (waitForAllRefCountsToClear) would block on
  // unrelated puts to the currently-active table, causing 500 ms timeouts under
  // sustained wall-clock profiling and leaving collect() racing with a still-
  // running put().  Since standby and scratch tables never appear as the
  // _active_storage, this wait returns instantly for them; for the active table
  // (called from clear() -> clearTableOnly()) the protection comes from the caller
  // holding lockAll() (which blocks signal-handler puts) and from this in-function
  // targeted wait — there is no prior caller-side drain.
  RefCountGuard::waitForRefCountToClear(this);
  decrementCounters();

  // Disconnect the full _prev chain before freeing chunks.  The advance step
  // must use a pre-saved pointer because setPrev(nullptr) clears the link that
  // the original loop used for advancement, causing early termination after only
  // the first node on an expanded (multi-node) table.
  for (LongHashTable *table = __atomic_load_n(&_table, __ATOMIC_ACQUIRE);
       table != nullptr; ) {
    LongHashTable *next = table->prev();
    table->setPrev(nullptr);
    table = next;
  }

  // Detach chunks for deferred deallocation - keeps trace memory alive
  ChunkList detached_chunks = _allocator.detachChunks();

  // Reinitialize with fresh table (using the new chunk from detachChunks)
  // Note: If detachChunks() failed to allocate a fresh chunk, the allocator's
  // _tail will be nullptr. LongHashTable::allocate will try to allocate,
  // which will call LinearAllocator::alloc(), which needs to handle nullptr _tail.
  // This is already handled in alloc() by checking _tail before use.
  // RELEASE: pairs with ACQUIRE loads in collect() and put() to ensure the
  // freshly-initialised table is visible on weakly-ordered architectures (aarch64).
  __atomic_store_n(&_table,
      LongHashTable::allocate(nullptr, INITIAL_CAPACITY, &_allocator),
      __ATOMIC_RELEASE);
  _overflow = 0;

  return detached_chunks;
}

void CallTraceHashTable::clear() {
  // Clear table and immediately free chunks (original behavior)
  ChunkList chunks = clearTableOnly();
  LinearAllocator::freeChunks(chunks);
}

// Adaptation of MurmurHash64A by Austin Appleby
u64 CallTraceHashTable::calcHash(int num_frames, ASGCT_CallFrame *frames,
                               bool truncated) {
  const u64 M = 0xc6a4a7935bd1e995ULL;
  const int R = 47;

  int len = num_frames * sizeof(ASGCT_CallFrame);
  u64 h = len * M * (truncated ? 1 : 2);

  const u64 *data = (const u64 *)frames;
  const u64 *end = data + len / sizeof(u64);

  while (data != end) {
    u64 k = *data++;
    k *= M;
    k ^= k >> R;
    k *= M;
    h ^= k;
    h *= M;
  }

  if (len & 4) {
    h ^= *(u32 *)data;
    h *= M;
  }

  h ^= h >> R;
  h *= M;
  h ^= h >> R;

  return h;
}

CallTrace *CallTraceHashTable::storeCallTrace(int num_frames,
                                            ASGCT_CallFrame *frames,
                                            bool truncated, u64 trace_id) {
  const size_t header_size = sizeof(CallTrace) - sizeof(ASGCT_CallFrame);
  const size_t total_size = header_size + num_frames * sizeof(ASGCT_CallFrame);
  void *memory = _allocator.alloc(total_size);
  CallTrace *buf = nullptr;
  if (memory != nullptr) {
    // Use placement new to invoke constructor in-place
    buf = new (memory) CallTrace(truncated, num_frames, trace_id);
    // Do not use memcpy inside signal handler
    for (int i = 0; i < num_frames; i++) {
      buf->frames[i] = frames[i];
    }
    Counters::increment(CALLTRACE_STORAGE_BYTES, total_size);
    Counters::increment(CALLTRACE_STORAGE_TRACES);
  }
  return buf;
}

CallTrace *CallTraceHashTable::findCallTrace(LongHashTable *table, u64 hash) {
  u64 *keys = table->keys();
  HashProbe probe(hash, table->capacity());

  u32 slot = probe.slot();
  while (true) {
    // Use atomic load: keys[] can be written concurrently via CAS in put()
    // when a table is promoted to prev but still has in-flight insertions.
    u64 key = __atomic_load_n(&keys[slot], __ATOMIC_ACQUIRE);
    if (key == hash) {
      // Use acquireTrace() to pair with the RELEASE store in setTrace().
      // If still PREPARING, treat as not found: callers will create a new entry.
      CallTrace *trace = table->values()[slot].acquireTrace();
      if (trace == CallTraceSample::PREPARING) {
        return nullptr;
      }
      return trace;
    }
    if (key == 0) {
      return nullptr;
    }
    if (!probe.hasNext()) {
      break;
    }
    slot = probe.next();
  };

  return nullptr;
}

void CallTraceHashTable::expandTableIfNeeded(LongHashTable* table, u32 size) {
  u32 capacity = table->capacity();

  // EXPANSION LOGIC: Check if load ratio reached after incrementing size
  if (size >= (u32) (capacity * LOAD_RATIO)  &&
      table == __atomic_load_n(&_table, __ATOMIC_RELAXED)) { // quick check, if other thread already expanded the table
    // Allocate new table with double capacity using LinearAllocator
    LongHashTable* new_table = LongHashTable::allocate(table, capacity * 2, &_allocator);
    if (new_table != nullptr) {
      // Atomic table swap - only one thread succeeds
      __atomic_compare_exchange_n(&_table, &table, new_table, false, __ATOMIC_ACQ_REL, __ATOMIC_RELAXED);
    }
  }
}

u64 CallTraceHashTable::put(int num_frames, ASGCT_CallFrame *frames,
                          bool truncated, u64 weight) {
  u64 hash = calcHash(num_frames, frames, truncated);

  // ACQUIRE pairs with the ACQ_REL CAS in the expansion path below, ensuring
  // that if another thread published a new (expanded) table we see its fully
  // initialised contents.
  LongHashTable *table = __atomic_load_n(&_table, __ATOMIC_ACQUIRE);
  if (table == nullptr) {
    // Table allocation failed or was cleared - drop sample
    Counters::increment(CALLTRACE_STORAGE_DROPPED);
    return CallTraceStorage::DROPPED_TRACE_ID;
  }
  
  u64 *keys = table->keys();
  HashProbe probe(hash, table->capacity());

  u32 slot = probe.slot();
  while (true) {
    u64 key_value = __atomic_load_n(&keys[slot], __ATOMIC_RELAXED);
    if (key_value == hash) { 
      // Hash matches - wait for the preparing thread to complete
      CallTrace* current_trace = table->values()[slot].acquireTrace();
      
      // If another thread is preparing this slot, wait for completion
      if (current_trace == CallTraceSample::PREPARING) {
        // Wait for the preparing thread to complete, with timeout
        int wait_cycles = 0;
        const int MAX_WAIT_CYCLES = 1000; // ~1000 cycles should be enough for allocation
        
        do {
          // Brief spin-wait to allow preparing thread to complete
          for (volatile int i = 0; i < 10; i++) {
            spinPause(); // Architecture-specific pause instruction
          }
          
          current_trace = table->values()[slot].acquireTrace();
          wait_cycles++;
          
          // Check if key was cleared (preparation failed)
          if (__atomic_load_n(&keys[slot], __ATOMIC_RELAXED) != hash) {
            break; // Key cleared, restart search
          }
          
        } while (current_trace == CallTraceSample::PREPARING && wait_cycles < MAX_WAIT_CYCLES);
        
        // If still preparing after timeout, something is wrong - continue search
        if (current_trace == CallTraceSample::PREPARING) {
          continue;
        }
      }
      
      // Check final state after waiting
      if (current_trace != nullptr && current_trace != CallTraceSample::PREPARING) {
        // Trace is ready, use it
        return current_trace->trace_id;
      } else {
        // Trace is nullptr but hash exists - preparation failed
        u64 recheck_key = __atomic_load_n(&keys[slot], __ATOMIC_ACQUIRE);
        if (recheck_key != hash) {
          continue; // Key was cleared, retry
        }
        Counters::increment(CALLTRACE_STORAGE_DROPPED);
        return CallTraceStorage::DROPPED_TRACE_ID;
      }
    }
    
    if (key_value == 0) {
      u64 expected = 0;
      if (!__atomic_compare_exchange_n(&keys[slot], &expected, hash, false, __ATOMIC_ACQ_REL, __ATOMIC_RELAXED)) {
        continue; // another thread claimed it, go to next slot
      }
      
      // Mark the slot as being prepared so other threads know to wait
      if (!table->values()[slot].markPreparing()) {
        // Failed to mark as preparing (shouldn't happen), clear key with full barrier and retry
        __atomic_thread_fence(__ATOMIC_SEQ_CST);
        __atomic_store_n(&keys[slot], 0, __ATOMIC_RELEASE);
        continue;
      }

      // Increment size counter for statistics and check for expansion
      u32 new_size = table->incSize();

      probe.updateCapacity(new_size);

      expandTableIfNeeded(table, new_size);

      // Check if trace exists in previous tables to avoid duplication
      CallTrace *trace = nullptr;
      if (table->prev() != nullptr) {
        trace = findCallTrace(table->prev(), hash);
      }
      
      if (trace == nullptr) {
        // Generate unique trace ID: upper 32 bits = instance_id, lower 32 bits = slot
        // ACQUIRE ordering synchronizes with RELEASE store in setInstanceId() to ensure
        // visibility of new instance_id on weakly-ordered architectures (aarch64, POWER)
        u64 instance_id = _instance_id.load(std::memory_order_acquire);
        u64 trace_id = (instance_id << 32) | slot;
        trace = storeCallTrace(num_frames, frames, truncated, trace_id);
        if (trace == nullptr) {
          // Allocation failure - reset trace first, then clear key
          table->values()[slot].setTrace(nullptr);
          __atomic_thread_fence(__ATOMIC_SEQ_CST);
          __atomic_store_n(&keys[slot], 0, __ATOMIC_RELEASE);
          Counters::increment(CALLTRACE_STORAGE_DROPPED);
          return CallTraceStorage::DROPPED_TRACE_ID;
        }
      }
      
      // Set the actual trace (this changes state from PREPARING to ready)
      table->values()[slot].setTrace(trace);
      return trace->trace_id;
    }

    if (!probe.hasNext()) {
      // Table overflow - very unlikely with expansion logic
      atomicIncRelaxed(_overflow);
      return OVERFLOW_TRACE_ID;
    }
    // Prime probing for better distribution
    slot = probe.next();
  }
}

void CallTraceHashTable::collect(std::unordered_set<CallTrace *> &traces, std::function<void(CallTrace*)> trace_hook) {
  // Lock-free collection for read-only tables.
  // Use ACQUIRE to pair with the ACQ_REL CAS in put()'s expansion path and the
  // RELEASE store in clearTableOnly(); ensures we see the fully-initialised table
  // on weakly-ordered architectures (aarch64).
  for (LongHashTable *table = __atomic_load_n(&_table, __ATOMIC_ACQUIRE);
       table != nullptr; table = table->prev()) {
    u64 *keys = table->keys();
    CallTraceSample *values = table->values();
    u32 capacity = table->capacity();
    for (u32 slot = 0; slot < capacity; slot++) {
      if (keys[slot] != 0) {
        CallTrace *trace = values[slot].acquireTrace();
        if (trace != nullptr && trace != CallTraceSample::PREPARING) {
          if (trace_hook) {
            trace_hook(trace);  // Call hook first if provided
          }
          traces.insert(trace);
        }
      }
    }
  }
  
  // Handle overflow trace
  if (_overflow > 0) {
    if (trace_hook) {
      trace_hook(&_overflow_trace);  // Call hook for overflow trace too
    }
    traces.insert(&_overflow_trace);
  }
}


void CallTraceHashTable::putWithExistingId(CallTrace* source_trace, u64 weight) {
  // Trace preservation for standby tables (no contention with new puts)
  // This is safe because new put() operations go to the new active table
  
  u64 hash = calcHash(source_trace->num_frames, source_trace->frames, source_trace->truncated);
  
  // First check if trace already exists in any table in the chain.
  // Use ACQUIRE to match the RELEASE store in clearTableOnly(); putWithExistingId()
  // is only called on scratch/standby tables with no concurrent writers, so the
  // load is safe, but consistent ordering prevents latent issues if callers change.
  for (LongHashTable *search_table = __atomic_load_n(&_table, __ATOMIC_ACQUIRE);
       search_table != nullptr; search_table = search_table->prev()) {
    CallTrace *existing_trace = findCallTrace(search_table, hash);
    if (existing_trace != nullptr) {
      return;
    }
  }

  LongHashTable *table = __atomic_load_n(&_table, __ATOMIC_ACQUIRE);
  if (table == nullptr) {
    return; // Table allocation failed
  }
  
  u64 *keys = table->keys();
  u32 capacity = table->capacity();

  HashProbe probe(hash, capacity);
  u32 slot = probe.slot();

  // Look for existing entry or empty slot - no locking needed
  while (true) {
    u64 key_value = __atomic_load_n(&keys[slot], __ATOMIC_RELAXED);
    if (key_value == 0) {
      // Found empty slot - claim it atomically
      u64 expected = 0;
      if (__atomic_compare_exchange_n(&keys[slot], &expected, hash, false, __ATOMIC_ACQ_REL, __ATOMIC_RELAXED)) {
        // Successfully claimed the slot
        // Create a copy of the source trace preserving its exact ID
        const size_t header_size = sizeof(CallTrace) - sizeof(ASGCT_CallFrame);
        const size_t total_size = header_size + source_trace->num_frames * sizeof(ASGCT_CallFrame);
        void *memory = _allocator.alloc(total_size);
        if (memory != nullptr) {
          // Use placement new to invoke constructor in-place
          CallTrace* copied_trace = new (memory) CallTrace(source_trace->truncated, source_trace->num_frames, source_trace->trace_id);
          // memcpy safe since not in signal handler
          memcpy(copied_trace->frames, source_trace->frames, source_trace->num_frames * sizeof(ASGCT_CallFrame));
          table->values()[slot].setTrace(copied_trace);
          Counters::increment(CALLTRACE_STORAGE_BYTES, total_size);
          Counters::increment(CALLTRACE_STORAGE_TRACES);

          // Increment table size
          table->incSize();
        } else {
          // Allocation failure - clear the key we claimed
          __atomic_store_n(&keys[slot], 0, __ATOMIC_RELEASE);
        }
        break;
      }
    }
    if (probe.hasNext()) {
      slot = probe.next();
    } else {
      // No more slots. The sample is dropped
      Counters::increment(CALLTRACE_STORAGE_DROPPED);
      break;
    }
  }
}