import type { KVCache } from './base.js';
import type { SlidingWindowKVCache } from './sliding-window.js';
import type {
  KVCacheConfig,
  KVGetResult,
  TieredGPUBuffersResult,
  MemoryStats,
  GPUContext,
} from './types.js';
import type { EmulatedVramStore } from '../../storage/emulated-vram.js';

/**
 * Tiered KV cache: hot ring buffer + cold paged cache.
 *
 * @module inference/kv-cache/tiered
 */
export class TieredKVCache {
  numLayers: number;
  numHeads: number;
  headDim: number;
  maxSeqLen: number;
  useGPU: boolean;
  layout: 'tiered';
  kvDtype: 'f16' | 'f32';
  bytesPerElem: number;
  kvSize: number;
  hotWindow: number;
  coldPageSize: number;
  coldDtype: 'f16' | 'f32';
  tieringMode: 'off' | 'fp16' | 'int8' | 'int4' | 'turboquant' | 'turboquant_prod';
  coldQuantMode: 'none' | 'int8' | 'int4' | 'turboquant' | 'turboquant_prod';
  coldPackedStride: number;
  compression: {
    mode: 'none' | 'int8' | 'int4' | 'turboquant' | 'turboquant_prod';
    blockSize: number;
    bitWidth?: number;
    prodMode?: boolean;
  };
  gating: { mode: 'auto' | 'force_on' | 'force_off'; minAluBwRatio: number };
  currentSeqLen: number;
  totalTokensSeen: number;
  memoryUsage: number;
  gpuContext: GPUContext | null;
  coldStore: EmulatedVramStore | null;
  coldStorePartition: string;
  coldStoreRegistered: boolean;
  coldStoreChunks: string[];
  hotCache: SlidingWindowKVCache;
  coldCache: KVCache | null;
  rotationMatrixBuffer: GPUBuffer | null;
  codebookCentroidsBuffer: GPUBuffer | null;
  codebookBoundariesBuffer: GPUBuffer | null;
  qjlMatrixBuffer: GPUBuffer | null;
  coldLayers: Array<{
    keysPackedGPU: GPUBuffer;
    valuesPackedGPU: GPUBuffer;
    scalesKGPU: GPUBuffer;
    scalesVGPU: GPUBuffer;
    residualKGPU?: GPUBuffer;
    residualVGPU?: GPUBuffer;
    residualNormsKGPU?: GPUBuffer;
    residualNormsVGPU?: GPUBuffer;
    seqLen: number;
  }> | null;

  constructor(
    config: KVCacheConfig,
    caches?: {
      hotCache: SlidingWindowKVCache;
      coldCache?: KVCache | null;
      coldLayers?: Array<{
        keysPackedGPU: GPUBuffer;
        valuesPackedGPU: GPUBuffer;
        scalesKGPU: GPUBuffer;
        scalesVGPU: GPUBuffer;
        residualKGPU?: GPUBuffer;
        residualVGPU?: GPUBuffer;
        residualNormsKGPU?: GPUBuffer;
        residualNormsVGPU?: GPUBuffer;
        seqLen: number;
      }> | null;
      coldStore?: EmulatedVramStore | null;
      coldStorePartition?: string;
    } | null
  );
  clear(): void;
  update(layerIdx: number, keys: Float32Array, values: Float32Array, startPos?: number): void;
  updateFromGPU(layerIdx: number, keysBuffer: GPUBuffer, valuesBuffer: GPUBuffer, startPos: number, numTokens: number, tokenIds?: number[] | null): Promise<void>;
  recordUpdateFromGPU(
    recorder: import('../../gpu/kernel-selector.js').CommandRecorder,
    layerIdx: number,
    keysBuffer: GPUBuffer,
    valuesBuffer: GPUBuffer,
    startPos: number,
    numTokens: number,
    tokenIds?: number[] | null
  ): Promise<void>;
  get(layerIdx: number): KVGetResult;
  getGPUBuffers(layerIdx: number): TieredGPUBuffersResult | null;
  hasGPUCache(): boolean;
  truncate(length: number): void;
  getMemoryStats(): MemoryStats;
  setGPUContext(gpuContext: GPUContext): void;
  destroy(): void;
  clone(): TieredKVCache;
}