// Inference backends for PiVoEngine.
//  - createLiteBackend: zero-dependency. onnx-lite interpreter for encoder/decoder and the
//    hand-written vocoder (WASM SIMD or plain JS kernel), or the interpreter for everything.
//  - createOrtBackend: onnxruntime-web (wasm or webgpu execution provider) for all four graphs.

import { Session, parseModel, tensor, setAccelerator } from './onnx-lite.ts';
import { WasmGemm, wasmSimdSupported } from './accel.ts';
import { HandVocoder } from './vocoder.ts';
import type { VocoderKind } from './vocoder.ts';
import type { Backend } from './engine.ts';
import type { AssetBuffers } from './assets.ts';
import type { Tensor } from './onnx-lite.ts';

export type LiteCache = Record<string, Tensor>;
export type ModelBuffers = Pick<AssetBuffers, 'encoder' | 'decoder' | 'conditioning' | 'vocoder_step'>;

// The parts of onnxruntime-web the benchmark uses (it is loaded at runtime, never bundled).
export interface OrtTensor {
  data: unknown;
  dims: readonly number[];
  location?: string;
  dispose(): void;
}
export type OrtFeeds = Record<string, OrtTensor>;
export interface OrtSession {
  run(feeds: OrtFeeds): Promise<OrtFeeds>;
  release?(): void;
}
export interface OrtModule {
  Tensor: new (type: string, data: Float32Array | BigInt64Array, dims: number[]) => OrtTensor;
  InferenceSession: { create(model: Uint8Array, options: object): Promise<OrtSession> };
  env: { wasm: { wasmPaths?: string; numThreads?: number; simd?: boolean } };
}

const CACHE_NAMES = ['key0', 'value0', 'key1', 'value1'];
const STATE: [string, number][] = [['memory', 256], ['gru1', 160], ['gru2', 128], ['gru3', 128], ['frame_state', 164]];

function trimRows(data: Float32Array, heads: number, rows: number, dim: number, keep: number): Float32Array {
  const out = new Float32Array(heads * keep * dim);
  for (let h = 0; h < heads; h++) {
    out.set(data.subarray((h * rows + rows - keep) * dim, (h * rows + rows) * dim), h * keep * dim);
  }
  return out;
}

// vocoder: 'wasm' | 'js' | 'interpreted'
export function createLiteBackend(
  buffers: ModelBuffers,
  { vocoder = (HandVocoder.wasmAvailable() ? 'wasm' : 'js') as VocoderKind | 'interpreted' } = {},
): Backend<LiteCache> {
  // The plain-JS and interpreted variants stay pure JS so they can be benchmarked as such.
  setAccelerator(vocoder === 'wasm' && wasmSimdSupported() ? new WasmGemm() : null);
  const encoder = Session.fromBuffer(buffers.encoder);
  const decoder = Session.fromBuffer(buffers.decoder);
  const vocGraph = parseModel(buffers.vocoder_step);
  const condGraph = parseModel(buffers.conditioning);
  const i64 = (dims: number[], values: ArrayLike<number>) => tensor('i64', dims, Float64Array.from(values));

  let frameFn: Backend<LiteCache>['vocoderFrame'];
  let condFn: Backend<LiteCache>['conditioning'];
  let resetFn: () => void;
  if (vocoder === 'interpreted') {
    const condSession = new Session(condGraph);
    const stepSession = new Session(vocGraph);
    let state: LiteCache;
    resetFn = () => {
      state = Object.fromEntries(STATE.map(([n, w]) => [n, tensor('f32', [1, w])]));
    };
    condFn = (context, periods) =>
      condSession.run({ features: tensor('f32', [1, 3, 20], Float32Array.from(context)), periods: i64([1, 3], periods) }).conditioning.data as Float32Array;
    frameFn = (cond, period, audio, base) => {
      for (let sub = 0; sub < 4; sub++) {
        const r = stepSession.run({ conditioning: tensor('f32', [1, 80], cond.slice(sub * 80, sub * 80 + 80)), period: i64([1], [period]), ...state });
        for (const [n] of STATE) state[n] = r[`${n}_out`];
        audio.set(r.excitation.data as Float32Array, base + sub * 40);
      }
    };
  } else {
    const hv = new HandVocoder(vocGraph, condGraph, vocoder);
    resetFn = () => hv.reset();
    condFn = (context, periods) => hv.conditioning(context, periods);
    frameFn = (cond, period, audio, base) => {
      for (let sub = 0; sub < 4; sub++) hv.step(cond, sub * 80, period, audio, base + sub * 40);
    };
  }

  return {
    name: vocoder === 'interpreted' ? 'JS interpreter (all graphs)' : `onnx-lite + hand-written vocoder (${vocoder === 'wasm' ? 'WASM SIMD' : 'JS'})`,
    encoder(ids, voice) {
      const T = ids.length;
      const r = encoder.run({ tokens: i64([1, T], ids), voice: i64([1], [voice]), lengths: i64([1], [T]) });
      return { encoded: r.encoded.data as Float32Array, logDurations: r.log_durations.data as Float32Array };
    },
    emptyCache() {
      return Object.fromEntries(CACHE_NAMES.map((n) => [n, tensor('f32', [1, 2, 0, 24])]));
    },
    decoder(chunk, n, position, cache) {
      const r = decoder.run({ expanded: tensor('f32', [1, n, 100], chunk), position: i64([], [position]), ...cache });
      return { features: r.features.data as Float32Array, cache: Object.fromEntries(CACHE_NAMES.map((name) => [name, r[`${name}_out`]])) };
    },
    trimCache(cache, keep) {
      const rows = cache.key0.dims[2];
      if (rows <= keep) return cache;
      return Object.fromEntries(CACHE_NAMES.map((n) => [n, tensor('f32', [1, 2, keep, 24], trimRows(cache[n].data as Float32Array, 2, rows, 24, keep))]));
    },
    resetVocoder: resetFn,
    conditioning: condFn,
    vocoderFrame: frameFn,
  };
}

export async function createOrtBackend(ort: OrtModule, buffers: ModelBuffers, { ep = 'wasm' } = {}): Promise<Backend<OrtFeeds>> {
  const gpu = ep === 'webgpu';
  const options = { executionProviders: [ep], graphOptimizationLevel: 'all' };
  const create = (buf: Uint8Array, extra = {}) => ort.InferenceSession.create(new Uint8Array(buf), { ...options, ...extra });
  const encoder = await create(buffers.encoder);
  const decoder = await create(buffers.decoder);
  const conditioning = await create(buffers.conditioning);
  // On WebGPU keep the recurrent state on the GPU between steps; only excitation comes back.
  const stepOutputs = gpu ? Object.fromEntries(STATE.map(([n]) => [`${n}_out`, 'gpu-buffer'])) : undefined;
  const step = await create(buffers.vocoder_step, gpu ? { preferredOutputLocation: { ...stepOutputs, excitation: 'cpu' } } : {});

  const big = (values: number[]) => BigInt64Array.from(values, (v) => BigInt(v));
  // Reused input tensors: their backing arrays are rewritten in place before each run.
  const condBuf = new Float32Array(80);
  const condTensor = new ort.Tensor('float32', condBuf, [1, 80]);
  const periodBuf = new BigInt64Array(1);
  const periodTensor = new ort.Tensor('int64', periodBuf, [1]);
  const ctxBuf = new Float32Array(60);
  const ctxTensor = new ort.Tensor('float32', ctxBuf, [1, 3, 20]);
  const perBuf = new BigInt64Array(3);
  const perTensor = new ort.Tensor('int64', perBuf, [1, 3]);
  let state: OrtFeeds | undefined;

  const disposeState = () => {
    if (!state) return;
    for (const t of Object.values(state)) if (t.location === 'gpu-buffer') t.dispose();
  };

  return {
    name: `onnxruntime-web ${ep}`,
    async encoder(ids, voice) {
      const T = ids.length;
      const r = await encoder.run({
        tokens: new ort.Tensor('int64', big(ids), [1, T]),
        voice: new ort.Tensor('int64', big([voice]), [1]),
        lengths: new ort.Tensor('int64', big([T]), [1]),
      });
      return { encoded: r.encoded.data as Float32Array, logDurations: r.log_durations.data as Float32Array };
    },
    emptyCache() {
      return Object.fromEntries(CACHE_NAMES.map((n) => [n, new ort.Tensor('float32', new Float32Array(0), [1, 2, 0, 24])]));
    },
    async decoder(chunk, n, position, cache) {
      const r = await decoder.run({
        expanded: new ort.Tensor('float32', chunk.slice(), [1, n, 100]),
        position: new ort.Tensor('int64', big([position]), []),
        ...cache,
      });
      return { features: r.features.data as Float32Array, cache: Object.fromEntries(CACHE_NAMES.map((name) => [name, r[`${name}_out`]])) };
    },
    trimCache(cache, keep) {
      const rows = cache.key0.dims[2];
      if (rows <= keep) return cache;
      return Object.fromEntries(CACHE_NAMES.map((n) => [n, new ort.Tensor('float32', trimRows(cache[n].data as Float32Array, 2, rows, 24, keep), [1, 2, keep, 24])]));
    },
    resetVocoder() {
      disposeState();
      state = Object.fromEntries(STATE.map(([n, w]) => [n, new ort.Tensor('float32', new Float32Array(w), [1, w])]));
    },
    async conditioning(context, periods) {
      ctxBuf.set(context);
      for (let k = 0; k < 3; k++) perBuf[k] = BigInt(periods[k]);
      const r = await conditioning.run({ features: ctxTensor, periods: perTensor });
      return r.conditioning.data as Float32Array;
    },
    async vocoderFrame(cond, period, audio, base) {
      periodBuf[0] = BigInt(period);
      for (let sub = 0; sub < 4; sub++) {
        condBuf.set(cond.subarray(sub * 80, sub * 80 + 80));
        const r = await step.run({ conditioning: condTensor, period: periodTensor, ...state });
        const next: OrtFeeds = {};
        for (const [n] of STATE) next[n] = r[`${n}_out`];
        if (gpu) for (const t of Object.values(state as OrtFeeds)) if (t.location === 'gpu-buffer') t.dispose();
        state = next;
        audio.set(r.excitation.data as Float32Array, base + sub * 40);
      }
    },
    release() {
      disposeState();
      for (const s of [encoder, decoder, conditioning, step]) s.release?.();
    },
  };
}