// Port of pivo/runtime.py + pivo/engine.py validation. Backend-agnostic: the
// encoder/decoder/conditioning/vocoder calls come from a backend object (see backends.js).

import { phonemize, SPECIAL, VOICE_IDS, FrontendError } from './frontend.ts';
import type { FrontendResult, Lexicon, Voice } from './frontend.ts';

type MaybePromise<T> = T | Promise<T>;
export interface Backend<Cache = unknown> {
  name: string;
  encoder(ids: number[], voice: number): MaybePromise<{ encoded: Float32Array; logDurations: Float32Array }>;
  emptyCache(): Cache;
  decoder(chunk: Float32Array, n: number, position: number, cache: Cache): MaybePromise<{ features: Float32Array; cache: Cache }>;
  trimCache?(cache: Cache, keep: number): Cache;
  resetVocoder(): void;
  conditioning(context: Float32Array, periods: Int32Array): MaybePromise<Float32Array>;
  vocoderFrame(cond: Float32Array, period: number, audio: Float32Array, base: number): MaybePromise<void>;
  release?(): void;
}
export interface WordTiming {
  text: string;
  index: number;
  startFrame: number;
  endFrame: number;
}
export interface Plan {
  front: FrontendResult;
  encoded: Float32Array;
  durations: Int32Array;
  tokenStart: Int32Array;
  frames: number;
  words: WordTiming[];
  voice: Voice;
  speed: number;
}
export interface RenderStats {
  decoder: number;
  conditioning: number;
  vocoder: number;
}
export interface RenderOptions {
  onFrames?: (pcm: Float32Array, firstFrame: number) => void;
  shouldStop?: () => boolean;
  onStats?: (stats: RenderStats) => void;
  yieldEvery?: () => Promise<void>;
}

export const SAMPLE_RATE = 16000;
export const HOP = 160;
const WIDTH = 96;
const PHASE_DIMS = 4;
const FEATURES = 20;
const MAX_TOKEN_FRAMES = 200;
const MAX_UTTERANCE_FRAMES = 3000;
const MAX_CHARS = 2000;
const MAX_TOKENS = 512;
const CHUNK_FRAMES = 16;
const KV_KEEP = 64; // decoder attention window; older cache rows are fully masked
const PREEMPHASIS = 0.85;

// numpy.rint: round half to even.
export function rint(x: number): number {
  const r = Math.round(x);
  return Math.abs(x % 1) === 0.5 ? 2 * Math.round(x / 2) : r;
}

const f32 = Math.fround;

// Runtime.periods, evaluated in float32 like NumPy does for a float32 input.
function period(feature18: number): number {
  const c = feature18 < -10 ? -10 : feature18 > 10 ? 10 : feature18;
  const p = f32(256 / f32(2 ** f32(c + 1.5)));
  const r = rint(p);
  return r < 32 ? 32 : r > 255 ? 255 : r;
}

export function validate(text: unknown, voice: string, speed: number): void {
  if (typeof text !== 'string' || !text.trim()) throw new FrontendError('Type something for Pi or Vo to say.');
  if (text.length > MAX_CHARS) throw new FrontendError(`Input exceeds ${MAX_CHARS} characters; split it into clauses.`);
  if (!(voice in VOICE_IDS)) throw new FrontendError('Voice must be Pi or Vo.');
  if (!Number.isFinite(speed) || speed < 0.5 || speed > 2) throw new FrontendError('Speed must be between 0.5 and 2.0.');
}

export class PiVoEngine<Cache = unknown> {
  backend: Backend<Cache>;
  dictionary: Lexicon;
  constructor(backend: Backend<Cache>, dictionary: Lexicon) {
    this.backend = backend;
    this.dictionary = dictionary;
  }

  // Frontend + encoder + duration rounding. Returns everything the UI needs for alignment.
  async plan(text: string, { voice = 'Pi' as Voice, speed = 1.0 } = {}): Promise<Plan> {
    validate(text, voice, speed);
    const front = phonemize(text, this.dictionary);
    if (!front.symbols.some((s) => !SPECIAL.includes(s))) throw new FrontendError('Type something for Pi or Vo to say.');
    if (front.ids.length > MAX_TOKENS) throw new FrontendError(`Input exceeds ${MAX_TOKENS} phoneme tokens; split it into clauses.`);

    const T = front.ids.length;
    const { encoded, logDurations } = await this.backend.encoder(front.ids, VOICE_IDS[voice]);
    const maxLog = Math.log1p(MAX_TOKEN_FRAMES);
    const durations = new Int32Array(T);
    let frames = 0;
    for (let i = 0; i < T; i++) {
      const ld = Math.min(Math.max(logDurations[i], 0), maxLog);
      let d = rint(Math.expm1(ld) / speed);
      d = Math.min(Math.max(d, 0), MAX_TOKEN_FRAMES);
      if (front.ids[i] >= SPECIAL.length) d = Math.max(d, 1);
      if (front.ids[i] <= 2) d = 0;
      durations[i] = d;
      frames += d;
    }
    if (frames > MAX_UTTERANCE_FRAMES) throw new FrontendError('That text is longer than 30 seconds of speech; split it into clauses.');

    const tokenStart = new Int32Array(T);
    for (let i = 1; i < T; i++) tokenStart[i] = tokenStart[i - 1] + durations[i - 1];
    const words = front.words.map((w) => ({
      text: w.text,
      index: w.index,
      startFrame: tokenStart[w.tokenStart],
      endFrame: w.tokenEnd < T ? tokenStart[w.tokenEnd - 1] + durations[w.tokenEnd - 1] : frames,
    }));
    return { front, encoded, durations, tokenStart, frames, words, voice, speed };
  }

  // expanded = repeat(encoded, durations) ++ phoneme phase features (float32, like pivo/phase.py)
  expand(plan: Plan): Float32Array {
    const { encoded, durations, frames } = plan;
    const D = WIDTH + PHASE_DIMS;
    const out = new Float32Array(frames * D);
    const denom = f32(Math.log1p(MAX_TOKEN_FRAMES));
    let row = 0;
    for (let t = 0; t < durations.length; t++) {
      const n = durations[t];
      const counts = f32(n);
      const logc = f32(f32(Math.log1p(counts)) / denom);
      for (let k = 0; k < n; k++, row++) {
        const o = row * D;
        out.set(encoded.subarray(t * WIDTH, (t + 1) * WIDTH), o);
        const fraction = f32(f32(f32(k) + 0.5) / counts);
        const angle = f32(Math.PI * fraction);
        out[o + WIDTH] = f32(f32(2 * fraction) - 1);
        out[o + WIDTH + 1] = f32(Math.sin(angle));
        out[o + WIDTH + 2] = f32(Math.cos(angle));
        out[o + WIDTH + 3] = logc;
      }
    }
    return out;
  }

  // Streams audio: onFrames(pcm Float32Array, firstFrame) is called after every decoder chunk.
  async render(plan: Plan, { onFrames, shouldStop = () => false, onStats, yieldEvery }: RenderOptions = {}): Promise<Float32Array> {
    const b = this.backend;
    const expanded = this.expand(plan);
    const D = WIDTH + PHASE_DIMS;
    const total = plan.frames;
    const audio = new Float32Array(total * HOP);
    const context = new Float32Array(3 * FEATURES);
    const periods = new Int32Array(3);
    let deemph = 0;
    let produced = 0;
    let cache = b.emptyCache();
    b.resetVocoder();

    const queue: Float32Array[] = []; // Float32Array(20) feature frames
    let previous: Float32Array | null = null;
    const stats = { decoder: 0, conditioning: 0, vocoder: 0 };
    const now = () => performance.now();

    for (let position = 0; position < total; position += CHUNK_FRAMES) {
      const n = Math.min(CHUNK_FRAMES, total - position);
      let t0 = now();
      const dec = await b.decoder(expanded.subarray(position * D, (position + n) * D), n, position, cache);
      cache = b.trimCache ? b.trimCache(dec.cache, KV_KEEP) : dec.cache;
      stats.decoder += now() - t0;
      for (let f = 0; f < n; f++) queue.push(dec.features.slice(f * FEATURES, (f + 1) * FEATURES));
      const final = position + CHUNK_FRAMES >= total;
      const chunkStart = produced;
      while (queue.length >= 2 || (final && queue.length)) {
        const current = queue.shift() as Float32Array;
        const following = queue.length ? queue[0] : current;
        previous = previous ?? current;
        context.set(previous, 0);
        context.set(current, FEATURES);
        context.set(following, 2 * FEATURES);
        for (let k = 0; k < 3; k++) periods[k] = period(context[k * FEATURES + 18]);
        t0 = now();
        const cond = await b.conditioning(context, periods);
        const t1 = now();
        stats.conditioning += t1 - t0;
        const base = produced * HOP;
        await b.vocoderFrame(cond, periods[1], audio, base);
        stats.vocoder += now() - t1;
        // De-emphasis in float64 with persistent state (pivo/audio.py).
        for (let i = base; i < base + HOP; i++) {
          deemph = audio[i] + PREEMPHASIS * deemph;
          audio[i] = deemph;
        }
        previous = current;
        produced++;
      }
      if (produced > chunkStart && onFrames) onFrames(audio.subarray(chunkStart * HOP, produced * HOP), chunkStart);
      if (yieldEvery) await yieldEvery();
      if (shouldStop()) break;
    }
    if (onStats) onStats(stats);
    for (let i = 0; i < produced * HOP; i++) {
      if (!Number.isFinite(audio[i])) throw new Error('model produced non-finite audio');
    }
    return audio.subarray(0, produced * HOP);
  }
}

// PCM16 WAV, matching pivo.engine.wav_bytes (clip, *32767, round half even).
export function wavBytes(audio: Float32Array, sampleRate = SAMPLE_RATE): Uint8Array {
  const buffer = new ArrayBuffer(44 + audio.length * 2);
  const v = new DataView(buffer);
  const str = (o: number, s: string) => [...s].forEach((c, i) => v.setUint8(o + i, c.charCodeAt(0)));
  str(0, 'RIFF');
  v.setUint32(4, 36 + audio.length * 2, true);
  str(8, 'WAVE');
  str(12, 'fmt ');
  v.setUint32(16, 16, true);
  v.setUint16(20, 1, true);
  v.setUint16(22, 1, true);
  v.setUint32(24, sampleRate, true);
  v.setUint32(28, sampleRate * 2, true);
  v.setUint16(32, 2, true);
  v.setUint16(34, 16, true);
  str(36, 'data');
  v.setUint32(40, audio.length * 2, true);
  for (let i = 0; i < audio.length; i++) {
    const x = Math.min(1, Math.max(-1, audio[i]));
    v.setInt16(44 + i * 2, rint(Math.fround(x * 32767)), true);
  }
  return new Uint8Array(buffer);
}