import { fftInPlace } from "./fft.js";
import { cooperativeYield } from "./cooperative-yield.js";

function throwIfAborted(signal) {
  if (signal?.aborted) {
    const error = new Error("STFT 계산이 취소되었습니다.");
    error.name = "AbortError";
    throw error;
  }
}

export function hannWindow(size) {
  if (!Number.isInteger(size) || size < 2) throw new Error("Hann window 크기는 2 이상이어야 합니다.");
  const window = new Float64Array(size);
  for (let index = 0; index < size; index += 1) {
    window[index] = 0.5 - 0.5 * Math.cos((2 * Math.PI * index) / (size - 1));
  }
  return window;
}

const STFT_WORKSPACE_KIND = "micarrsim.stft-workspace/1";

export function createStftWorkspace(fftSize) {
  return {
    kind: STFT_WORKSPACE_KIND,
    fftSize,
    window: hannWindow(fftSize),
    real: new Float64Array(fftSize),
    imaginary: new Float64Array(fftSize),
  };
}

function resolveStftWorkspace(fftSize, workspace) {
  if (workspace === undefined) return createStftWorkspace(fftSize);
  if (
    workspace?.kind !== STFT_WORKSPACE_KIND ||
    workspace.fftSize !== fftSize ||
    workspace.window?.length !== fftSize ||
    workspace.real?.length !== fftSize ||
    workspace.imaginary?.length !== fftSize
  ) {
    throw new RangeError("STFT workspace must match fftSize.");
  }
  return workspace;
}

export function getFrameCount(sampleCount, fftSize, hopSize) {
  if (!Number.isInteger(hopSize) || hopSize < 1 || hopSize > fftSize) {
    throw new Error("hopSize는 1 이상 fftSize 이하여야 합니다.");
  }
  return Math.max(1, 1 + Math.ceil(Math.max(0, sampleCount - fftSize) / hopSize));
}

export async function extractStftBins(channels, options) {
  const {
    fftSize,
    hopSize,
    bins,
    signal,
    onProgress,
    transformBatchSize = 32,
    workspace: providedWorkspace,
  } = options;
  if (!Array.isArray(channels) || channels.length === 0) throw new Error("마이크 채널이 필요합니다.");
  const sampleCount = channels[0].length;
  if (channels.some((channel) => channel.length !== sampleCount)) throw new Error("마이크 채널 길이가 다릅니다.");
  const uniqueBins = [...new Set(bins)].sort((a, b) => a - b);
  for (const bin of uniqueBins) {
    if (!Number.isInteger(bin) || bin < 0 || bin > fftSize / 2) throw new Error(`유효하지 않은 FFT bin: ${bin}`);
  }

  const frameCount = getFrameCount(sampleCount, fftSize, hopSize);
  const microphoneCount = channels.length;
  if (!Number.isInteger(transformBatchSize) || transformBatchSize < 1) {
    throw new Error("transformBatchSize는 1 이상의 정수여야 합니다.");
  }
  const spectra = new Map();
  for (const bin of uniqueBins) {
    spectra.set(bin, {
      real: new Float32Array(frameCount * microphoneCount),
      imaginary: new Float32Array(frameCount * microphoneCount),
    });
  }

  const workspace = resolveStftWorkspace(fftSize, providedWorkspace);
  const { window, real, imaginary } = workspace;
  const totalTransforms = frameCount * microphoneCount;
  let completedTransforms = 0;
  for (let frameIndex = 0; frameIndex < frameCount; frameIndex += 1) {
    const frameStart = frameIndex * hopSize;
    for (let microphoneIndex = 0; microphoneIndex < microphoneCount; microphoneIndex += 1) {
      const channel = channels[microphoneIndex];
      real.fill(0);
      imaginary.fill(0);
      for (let offset = 0; offset < fftSize; offset += 1) {
        real[offset] = (channel[frameStart + offset] ?? 0) * window[offset];
      }
      fftInPlace(real, imaginary);
      const spectrumIndex = frameIndex * microphoneCount + microphoneIndex;
      for (const bin of uniqueBins) {
        const spectrum = spectra.get(bin);
        spectrum.real[spectrumIndex] = real[bin];
        spectrum.imaginary[spectrumIndex] = imaginary[bin];
      }
      completedTransforms += 1;
      if (
        completedTransforms % transformBatchSize === 0 &&
        completedTransforms < totalTransforms
      ) {
        onProgress?.({ completedTransforms, totalTransforms });
        await cooperativeYield();
        throwIfAborted(signal);
      }
    }
  }
  onProgress?.({ completedTransforms, totalTransforms });
  throwIfAborted(signal);
  return { frameCount, microphoneCount, spectra };
}

export function binsForBand(band, sampleRate, fftSize) {
  const binWidth = sampleRate / fftSize;
  const firstBin = Math.max(0, Math.ceil(band.minHz / binWidth));
  const lastBin = Math.min(fftSize / 2, Math.floor(band.maxHz / binWidth));
  if (firstBin > lastBin) {
    throw new Error(`${band.minHz}–${band.maxHz} Hz에 해당하는 FFT bin이 없습니다.`);
  }
  return Array.from({ length: lastBin - firstBin + 1 }, (_, index) => firstBin + index);
}