import { OBJECT, TOTAL_OVERHEAD } from "rt/common";
import { __heap_base, heap } from "memory";

// Buffer management constants
const SHRINK_EVERY_N_MASK: usize = 255; // check every 256 outputs
const MIN_BUFFER_SIZE: usize = 1024;

// Exponential moving average smoothing factor (0.0 to 1.0)
// Higher values = more responsive to recent sizes, lower = more stable
// Using 0.125 (1/8) for efficient bit-shift calculation
const EMA_ALPHA_SHIFT: usize = 3; // 1/8 = 0.125

/**
 * Central buffer namespace for managing memory operations.
 */
export namespace bs {
  /** Current unmanaged backing store pointer. */
  export let buffer: usize = heap.alloc(MIN_BUFFER_SIZE);

  /** Current offset within the buffer. */
  export let offset: usize = buffer;

  /** Byte length of the buffer. */
  export let bufferSize: usize = MIN_BUFFER_SIZE;

  /** Proposed size of output */
  export let stackSize: usize = 0;

  let pauseOffsets = new Array<usize>();
  let pauseStackSizes = new Array<usize>();

  // Exponential moving average of output sizes for adaptive buffer sizing
  // This provides smoother adaptation than simple averaging
  let typicalSize: usize = MIN_BUFFER_SIZE;
  let counter: usize = 0;

  /**
   * Updates the typical size using exponential moving average.
   * EMA formula: new_avg = alpha * new_value + (1 - alpha) * old_avg
   * Using bit shifts for efficiency: alpha = 1/8, so (1 - alpha) = 7/8
   * @param newSize - The new size to incorporate into the average
   */
  function updateTypicalSize(newSize: usize): void {
    // EMA: typicalSize = (newSize >> 3) + typicalSize - (typicalSize >> 3)
    // Simplified: typicalSize += (newSize - typicalSize) >> 3
    typicalSize += (newSize - typicalSize) >> EMA_ALPHA_SHIFT;
  }

  function renewBuffer(newSize: usize): void {
    const oldPtr = buffer;
    const relOffset = offset - oldPtr;
    const newPtr = heap.realloc(oldPtr, newSize);
    offset = newPtr + relOffset;
    buffer = newPtr;
    bufferSize = newSize;
  }

  function reserve(requiredSize: usize, extra: usize): void {
    if (requiredSize <= bufferSize) return;
    // Grow aggressively (2x) to minimize realloc frequency in hot serialization paths.
    let next = bufferSize << 1;
    const minNext = requiredSize + extra;
    if (next < minNext) next = minNext;
    renewBuffer(next);
  }

  function finalizeDynamicOutput(len: usize): void {
    counter += 1;
    updateTypicalSize(len);
    if ((counter & SHRINK_EVERY_N_MASK) == 0 && bufferSize > typicalSize << 2) {
      resize(u32(typicalSize << 1));
    }
    offset = buffer;
    stackSize = 0;
  }

  export let cacheOutput: usize = 0;
  export let cacheOutputLen: usize = 0;
  /**
   * Stores the state of the buffer, allowing further changes to be reset
   */
  export function saveState(): void {
    pauseOffsets.push(offset - buffer);
    pauseStackSizes.push(stackSize);
  }

  /**
   * Resets the buffer to the state it was in when `pause()` was called.
   * This allows for changes made after the pause to be discarded.
   */
  export function loadState(): void {
    const length = pauseOffsets.length;
    if (length == 0) return;
    const index = length - 1;
    offset = buffer + unchecked(pauseOffsets[index]);
    stackSize = unchecked(pauseStackSizes[index]);
    pauseOffsets.length = index;
    pauseStackSizes.length = index;
  }

  /**
   * Resets the buffer to a clean, empty state. Call this after a throw aborts a
   * serialize/deserialize op mid-flight: a partial run can leave `offset`
   * advanced and the pause stacks non-empty, which would corrupt the next op.
   */
  export function reset(): void {
    offset = buffer;
    stackSize = 0;
    pauseOffsets.length = 0;
    pauseStackSizes.length = 0;
  }

  /**
   * Proposes that the buffer size is should be greater than or equal to the proposed size.
   * If necessary, reallocates the buffer to the exact new size.
   * @param size - The size to propose.
   */
  export function ensureSize(size: u32): void {
    reserve(offset - buffer + usize(size), MIN_BUFFER_SIZE);
  }

  /**
   * Proposes that the buffer size is should be greater than or equal to the proposed size.
   * If necessary, reallocates the buffer to the exact new size.
   * @param size - The size to propose.w
   */
  export function proposeSize(size: u32): void {
    stackSize += size;
    reserve(stackSize, 0);
  }

  /**
   * Increases the proposed size by n + MIN_BUFFER_SIZE if necessary.
   * If necessary, reallocates the buffer to the exact new size.
   * @param size - The size to grow by.
   */
  export function growSize(size: u32): void {
    stackSize += size;
    reserve(stackSize, MIN_BUFFER_SIZE);
  }

  /**
   * Resizes the buffer to the specified size.
   * @param newSize - The new buffer size.
   */
  export function resize(newSize: u32): void {
    const oldPtr = buffer;
    const relOffset = offset - oldPtr;
    const newPtr = heap.realloc(buffer, newSize);
    buffer = newPtr;
    bufferSize = newSize;
    offset = buffer + relOffset;
  }

  /**
   * Shrinks the buffer using the same adaptive sizing policy as normal output
   * finalization. Keeps enough capacity for the recent typical output size while
   * releasing clearly excess memory.
   */
  export function shrink(): void {
    let next = typicalSize << 1;
    if (next < MIN_BUFFER_SIZE) next = MIN_BUFFER_SIZE;
    if (bufferSize > next) {
      resize(u32(next));
    } else {
      offset = buffer;
      stackSize = 0;
    }
  }

  /**
   * Copies the buffer's content to a new object of a specified type. Does not shrink the buffer.
   * @returns The new object containing the buffer's content.
   */
  export function cpyOut<T>(): T {
    if (pauseOffsets.length == 0) {
      const len = offset - buffer;
      // @ts-expect-error: __new is a runtime builtin
      const _out = __new(len, idof<T>());
      memory.copy(_out, buffer, len);
      return changetype<T>(_out);
    } else {
      const pauseOffset =
        buffer + unchecked(pauseOffsets[pauseOffsets.length - 1]);
      const len = offset - pauseOffset;
      // @ts-expect-error: __new is a runtime builtin
      const _out = __new(len, idof<T>());
      memory.copy(_out, pauseOffset, len);
      bs.loadState();
      return changetype<T>(_out);
    }
  }

  /**
   * Copies the slice starting at a caller-provided relative buffer offset and restores
   * `offset` back to that slice start.
   *
   * This is intended for deserializers that borrow `bs` as temporary scratch space and
   * already know their local slice start as `bs.offset - bs.buffer`.
   *
   * Note: this restores only `offset`. Deserialization paths do not currently depend on
   * `stackSize`, which is tracked for serialization growth heuristics.
   */
  export function sliceOut<T>(start: usize): T {
    const sliceStart = buffer + start;
    const len = offset - sliceStart;
    // @ts-expect-error: __new is a runtime builtin
    const _out = __new(len, idof<T>());
    memory.copy(_out, sliceStart, len);
    offset = sliceStart;
    return changetype<T>(_out);
  }

  /**
   * Copies the slice starting at a caller-provided relative buffer offset into a string field
   * and restores `offset` back to that slice start.
   */
  export function toField(start: usize, dstFieldPtr: usize): void {
    const sliceStart = buffer + start;
    const byteLength = <u32>(offset - sliceStart);
    if (byteLength == 0) {
      store<usize>(dstFieldPtr, changetype<usize>(""));
      offset = sliceStart;
      return;
    }

    const current = load<usize>(dstFieldPtr);
    let stringPtr: usize;
    if (current >= __heap_base) {
      if (changetype<OBJECT>(current - TOTAL_OVERHEAD).rtSize == byteLength) {
        stringPtr = current;
      } else {
        // @ts-expect-error: __renew is a runtime builtin
        stringPtr = __renew(current, byteLength);
        store<usize>(dstFieldPtr, stringPtr);
      }
    } else {
      // @ts-expect-error: __new is a runtime builtin
      stringPtr = __new(byteLength, idof<string>());
      store<usize>(dstFieldPtr, stringPtr);
    }

    memory.copy(stringPtr, sliceStart, byteLength);
    offset = sliceStart;
  }

  /**
   * Copies the buffer's content to a new object of a specified type.
   * Uses exponential moving average to track typical output sizes for
   * adaptive buffer management - shrinks buffer when consistently oversized.
   * @returns The new object containing the buffer's content.
   */
  export function out<T>(): T {
    let out: usize;
    if (cacheOutput === 0) {
      const len = offset - buffer;
      // @ts-expect-error: __new is a runtime builtin
      out = __new(len, idof<T>());
      memory.copy(out, buffer, len);
      finalizeDynamicOutput(len);
    } else {
      // zero-copy path
      // @ts-expect-error: __new is a runtime builtin
      out = __new(cacheOutputLen, idof<T>());
      memory.copy(out, cacheOutput, cacheOutputLen);
      // reset arena flag
      cacheOutput = 0;
      cacheOutputLen = 0;
      offset = buffer;
      stackSize = 0;
    }
    return changetype<T>(out);
  }

  /**
   * Like `out<T>()`, but writes the finished bytes into the existing object
   * `target`, reusing its allocation: an in-place overwrite when the size
   * already matches, `__renew` when it differs. Falls back to a fresh
   * `out<T>()` when `target` is null or a static (`< __heap_base`), so callers
   * can pass an empty/uninitialized target on the first call. Mirrors the reuse
   * policy of `toField`, but for the top-level return value.
   */
  export function outTo<T>(target: usize): T {
    if (target < __heap_base) return out<T>();
    let len: usize;
    let src: usize;
    if (cacheOutput === 0) {
      len = offset - buffer;
      src = buffer;
    } else {
      len = cacheOutputLen;
      src = cacheOutput;
    }
    let dst = target;
    if (changetype<OBJECT>(target - TOTAL_OVERHEAD).rtSize != <u32>len) {
      // @ts-expect-error: __renew is a runtime builtin
      dst = __renew(target, len);
    }
    memory.copy(dst, src, len);
    if (cacheOutput === 0) {
      finalizeDynamicOutput(len);
    } else {
      cacheOutput = 0;
      cacheOutputLen = 0;
      offset = buffer;
      stackSize = 0;
    }
    return changetype<T>(dst);
  }
}

/**
 * String Caching (sc) namespace for optimizing repeated string serialization.
 *
 * This caching system can significantly boost performance when serializing
 * objects with repeated string values. When enabled, serialized strings are
 * cached and reused on subsequent serializations of the same string reference.
 *
 * ## Configuration
 *
 * Enable caching by setting the `JSON_CACHE` environment variable:
 * ```bash
 * JSON_CACHE=true npx asc your-file.ts --transform json-as/transform
 * ```
 *
 * You can also configure cache size directly:
 * - `JSON_CACHE=<int>` => bytes
 * - `JSON_CACHE=<int>kb|mb|gb` => kilobits/megabits/gigabits
 * - `JSON_CACHE=<int>KB|MB|GB` => kilobytes/megabytes/gigabytes
 *
 * ## Memory Configuration
 *
 * The cache uses a size budget from `JSON_CACHE` (or defaults to 1MB when
 * enabled with `true/on/yes`):
 *
 * - **CACHE_SIZE**: Number of direct-mapped cache slots. Derived from budget.
 * - **ARENA_SIZE**: Circular buffer for storing cached serialized strings.
 *   When full, wraps around and overwrites oldest entries. Larger values retain
 *   more cached data but consume more memory.
 *
 * - **MIN_CACHE_LEN** (128 bytes): Minimum serialized output size to cache.
 *   Smaller outputs aren't cached as the overhead isn't worth it.
 *   Only strings producing outputs >= 128 bytes are cached.
 *
 * ## Trade-offs
 *
 * - **Memory**: ~1.05MB fixed overhead when enabled (arena + entry table)
 * - **Best for**: Applications serializing the same string references repeatedly
 * - **Not ideal for**: One-time serializations or highly unique string content
 *
 * ## Performance
 *
 * When effective, caching can achieve >22 GB/s serialization throughput by
 * avoiding re-serialization of previously seen strings.
 */
export namespace sc {
  export const ENTRY_KEY = offsetof<sc.Entry>("key");
  export const ENTRY_PTR = offsetof<sc.Entry>("ptr");
  export const ENTRY_LEN = offsetof<sc.Entry>("len");

  // @ts-expect-error: JSON_CACHE may not be defined. If so, it will default to false.
  export const CACHE_ENABLED: bool = isDefined(JSON_CACHE) ? JSON_CACHE : false;
  // @ts-expect-error: JSON_CACHE_SIZE may not be defined. If so, it will default to 1MB.
  export const CACHE_BYTES: usize = isDefined(JSON_CACHE_SIZE)
    ? <usize>JSON_CACHE_SIZE
    : 1 << 20;
  /** Minimum serialized length to cache - smaller outputs aren't worth caching */
  export const MIN_CACHE_LEN: usize = 128;
  /** Size of the circular arena buffer for cached strings */
  export const ARENA_SIZE: usize =
    CACHE_BYTES >= MIN_CACHE_LEN ? CACHE_BYTES : MIN_CACHE_LEN;

  /** Number of cache slots (power of 2 for efficient masking). Set to 0 when caching disabled. */
  const CACHE_SIZE_BASE: i32 = CACHE_ENABLED
    ? i32(ARENA_SIZE >> 10 >= 1 ? ARENA_SIZE >> 10 : 1)
    : 0;
  export const CACHE_SIZE: usize = CACHE_ENABLED
    ? <usize>(1 << (32 - clz<i32>(CACHE_SIZE_BASE - 1)))
    : 0;
  /** Bitmask for fast modulo operation on cache index */
  export const CACHE_MASK: usize = CACHE_SIZE > 0 ? CACHE_SIZE - 1 : 0;

  /** Cache entry structure - stores pointer to string, cached output location, and length */
  @unmanaged
  export class Entry {
    /** Original string pointer (used as cache key) */
    key!: usize;
    /** Pointer to cached serialized output in arena */
    ptr!: usize;
    /** Length of cached serialized output */
    len!: usize;
  }

  /** Static array of cache entries */
  export const entries = new StaticArray<sc.Entry>(i32(CACHE_SIZE));
  /** Circular buffer arena for storing cached serialized strings */
  export const arena = new ArrayBuffer(i32(ARENA_SIZE));
  /** Current write position in the arena */
  export let arenaPtr: usize = changetype<usize>(arena);
  /** End boundary of the arena */
  export let arenaEnd: usize = arenaPtr + ARENA_SIZE;

  /**
   * Computes cache index for a given string pointer.
   * Uses pointer address shifted right by 4 bits (aligned to 16-byte boundaries)
   * masked to fit within cache size.
   */

  export function indexFor(ptr: usize): usize {
    return (ptr >> 4) & CACHE_MASK;
  }

  /**
   * Attempts to retrieve a cached serialization for the given string pointer.
   * If found, sets up the buffer system to use the cached output.
   * @param key - The string pointer to look up
   * @returns true if cache hit, false if cache miss
   */

  export function tryEmitCached(key: usize): bool {
    const e = unchecked(entries[indexFor(key)]);
    if (e.key == key) {
      bs.cacheOutput = e.ptr;
      bs.cacheOutputLen = e.len;
      return true;
    }
    return false;
  }

  /**
   * Stores a serialized string output in the cache for future reuse.
   * Only caches outputs >= MIN_CACHE_LEN bytes to avoid overhead for small strings.
   * Uses a circular arena buffer - when full, wraps around and overwrites oldest data.
   * @param str - Original string pointer (used as cache key)
   * @param start - Start of serialized output to cache
   * @param len - Length of serialized output
   */
  export function insertCached(str: usize, start: usize, len: usize): void {
    if (!CACHE_ENABLED || ARENA_SIZE == 0) return;
    if (len < MIN_CACHE_LEN) return;
    if (len > ARENA_SIZE) return;
    if (arenaPtr + len > arenaEnd) {
      // Wrap around to beginning of arena (circular buffer)
      arenaPtr = changetype<usize>(arena);
    }

    memory.copy(arenaPtr, start, len);

    const e = unchecked(entries[i32((str >> 4) & CACHE_MASK)]);
    e.key = str;
    e.ptr = arenaPtr;
    e.len = len;

    arenaPtr += len;
  }
}