import { spawn } from "child_process";
import { randomBytes } from "crypto";
import { tmpdir } from "os";
import { join } from "path";

export class AudioPreprocessor {
  private readonly targetSampleRate = 16000;
  private readonly channels = 1;

  async preprocess(audioBuffer: Buffer, inputSampleRate?: number): Promise<PreprocessedAudio> {
    const tempInput = this.writeTempFile(audioBuffer, ".wav");
    const tempOutput = this.writeTempFile(Buffer.alloc(0), ".wav");

    const outputSampleRate = inputSampleRate || this.detectSampleRate(audioBuffer);

    try {
      await this.runPythonPreprocess(tempInput, tempOutput, outputSampleRate);

      const processed = await this.readFile(tempOutput);
      return {
        audio: processed,
        sampleRate: this.targetSampleRate,
        channels: this.channels,
        durationSec: processed.length / (this.targetSampleRate * this.channels * 2),
      };
    } catch (err) {
      const fallback = await this.jsFallback(audioBuffer, outputSampleRate);
      return fallback;
    }
  }

  async preprocessBatch(
    inputs: Array<{ buffer: Buffer; sampleRate?: number }>
  ): Promise<PreprocessedAudio[]> {
    const promises = inputs.map(async (input) => {
      return this.preprocess(input.buffer, input.sampleRate);
    });
    return Promise.all(promises);
  }

  applyVAD(audioBuffer: Buffer, sampleRate: number): Buffer {
    const int16Array = this.bufferToInt16(audioBuffer);
    const silenceThreshold = 500;
    const windowSize = Math.floor(sampleRate * 0.02);

    const segments: number[][] = [];
    let currentSegment: number[] = [];

    for (let i = 0; i < int16Array.length; i += windowSize) {
      const window = int16Array.slice(i, i + windowSize);
      const rms = Math.sqrt(
        window.reduce((sum, val) => sum + val * val, 0) / window.length
      );

      if (rms > silenceThreshold) {
        currentSegment.push(...window);
      } else if (currentSegment.length > 0) {
        segments.push(currentSegment);
        currentSegment = [];
      }
    }

    if (currentSegment.length > 0) {
      segments.push(currentSegment);
    }

    const merged = segments.flat();
    return this.int16ToBuffer(merged);
  }

  normalizeAudio(audioBuffer: Buffer): Buffer {
    const int16Array = this.bufferToInt16(audioBuffer);
    const maxAmplitude = Math.max(...int16Array.map((v) => Math.abs(v)));
    const targetMax = 32767 * 0.95;

    if (maxAmplitude === 0) return audioBuffer;

    const scale = targetMax / maxAmplitude;
    const normalized = int16Array.map((v) => Math.round(v * scale));
    return this.int16ToBuffer(normalized);
  }

  async extractFeatures(audioBuffer: Buffer): Promise<AudioFeatures> {
    const int16Array = this.bufferToInt16(audioBuffer);
    const sampleRate = this.targetSampleRate;
    const windowSize = Math.floor(sampleRate * 0.025);
    const hopLength = Math.floor(sampleRate * 0.01);

    const mfccs: number[][] = [];
    const numCoeffs = 13;

    for (let i = 0; i < int16Array.length - windowSize; i += hopLength) {
      const frame = int16Array.slice(i, i + windowSize);
      const coeffs = this.computeMFCC(frame, numCoeffs);
      mfccs.push(coeffs);
    }

    const zeroCrossingRate = this.computeZCR(int16Array);
    const spectralCentroid = this.computeSpectralCentroid(int16Array);
    const spectralRollOff = this.computeSpectralRollOff(int16Array);

    return {
      mfccs,
      zeroCrossingRate,
      spectralCentroid,
      spectralRollOff,
      durationSec: int16Array.length / sampleRate,
    };
  }

  private async runPythonPreprocess(
    inputPath: string,
    outputPath: string,
    inputSampleRate: number
  ): Promise<void> {
    return new Promise((resolve, reject) => {
      const mlServiceUrl = process.env.VOICEPRINT_ML_URL || "http://localhost:8001";
      const proc = spawn("python3", [
        "-c",
        `
import urllib.request, json, sys
req = urllib.request.Request(
    "${mlServiceUrl}/preprocess",
    data=json.dumps({"input_path": "${inputPath}", "output_path": "${outputPath}", "input_sr": ${inputSampleRate}}).encode(),
    headers={"Content-Type": "application/json"}
)
try:
    with urllib.request.urlopen(req, timeout=30) as resp:
        sys.exit(0) if resp.status == 200 else sys.exit(1)
except:
    sys.exit(1)
`,
      ]);

      proc.on("close", (code) => {
        code === 0 ? resolve() : reject(new Error(`Python preprocess exited with code ${code}`));
      });
    });
  }

  private async jsFallback(
    audioBuffer: Buffer,
    inputSampleRate: number
  ): Promise<PreprocessedAudio> {
    let processed = audioBuffer;

    if (inputSampleRate !== this.targetSampleRate) {
      processed = this.resample(audioBuffer, inputSampleRate, this.targetSampleRate);
    }

    processed = this.applyVAD(processed, this.targetSampleRate);
    processed = this.normalizeAudio(processed);

    return {
      audio: processed,
      sampleRate: this.targetSampleRate,
      channels: this.channels,
      durationSec: processed.length / (this.targetSampleRate * this.channels * 2),
    };
  }

  private resample(buffer: Buffer, fromRate: number, toRate: number): Buffer {
    const int16 = this.bufferToInt16(buffer);
    const ratio = fromRate / toRate;
    const newLength = Math.round(int16.length / ratio);
    const resampled: number[] = [];

    for (let i = 0; i < newLength; i++) {
      const srcIdx = Math.floor(i * ratio);
      const nextIdx = Math.min(srcIdx + 1, int16.length - 1);
      const frac = (i * ratio) - srcIdx;
      resampled.push(Math.round(int16[srcIdx] * (1 - frac) + int16[nextIdx] * frac));
    }

    return this.int16ToBuffer(resampled);
  }

  private detectSampleRate(buffer: Buffer): number {
    if (buffer.length < 44) return 16000;

    const fmtOffset = buffer.toString("ascii", 12, 16).trim() === "fmt " ? 16 : 40;
    if (fmtOffset + 4 <= buffer.length) {
      const sr = buffer.readUInt32LE(fmtOffset + 4);
      if (sr >= 8000 && sr <= 48000) return sr;
    }
    return 16000;
  }

  private bufferToInt16(buffer: Buffer): number[] {
    const arr: number[] = new Array(buffer.length / 2);
    for (let i = 0; i < arr.length; i++) {
      arr[i] = buffer.readInt16LE(i * 2);
    }
    return arr;
  }

  private int16ToBuffer(arr: number[]): Buffer {
    const buf = Buffer.alloc(arr.length * 2);
    for (let i = 0; i < arr.length; i++) {
      buf.writeInt16LE(arr[i], i * 2);
    }
    return buf;
  }

  private computeMFCC(frame: number[], numCoeffs: number): number[] {
    const n = frame.length;
    const fftSize = Math.pow(2, Math.ceil(Math.log2(n)) + 1);
    const spectrum: number[] = new Array(fftSize / 2 + 1);

    for (let k = 0; k < spectrum.length; k++) {
      let real = 0, imag = 0;
      for (let n = 0; n < frame.length; n++) {
        const angle = (2 * Math.PI * k * n) / fftSize;
        real += frame[n] * Math.cos(angle);
        imag -= frame[n] * Math.sin(angle);
      }
      spectrum[k] = Math.sqrt(real * real + imag * imag);
    }

    const numFilters = 20;
    const filterbank = this.createFilterbank(spectrum.length, numFilters);
    const logEnergies: number[] = [];

    for (let m = 0; m < numFilters; m++) {
      let energy = 0;
      for (let k = 0; k < spectrum.length; k++) {
        energy += filterbank[m][k] * spectrum[k] * spectrum[k];
      }
      logEnergies.push(Math.log(Math.max(energy, 1e-10)));
    }

    const mfccs: number[] = [];
    for (let d = 0; d < numCoeffs; d++) {
      let coeff = 0;
      for (let m = 0; m < numFilters; m++) {
        coeff += logEnergies[m] * Math.cos(((2 * m + 1) * (d + 1) * Math.PI) / (2 * numFilters));
      }
      mfccs.push(coeff);
    }

    return mfccs;
  }

  private createFilterbank(numBins: number, numFilters: number): number[][] {
    const filterbank: number[][] = Array.from({ length: numFilters }, () =>
      new Array(numBins).fill(0)
    );
    const low = Math.floor(20 * numBins / 8000);
    const high = Math.floor(5500 * numBins / 8000);
    const spacing = (high - low) / (numFilters + 1);

    for (let m = 1; m <= numFilters; m++) {
      const center = Math.floor(low + m * spacing);
      const prev = Math.floor(low + (m - 1) * spacing);
      const next = Math.floor(low + (m + 1) * spacing);

      for (let k = prev; k < center; k++) {
        if (k > 0) filterbank[m - 1][k] = (k - prev) / (center - prev);
      }
      for (let k = center; k < next; k++) {
        if (next - center > 0) filterbank[m - 1][k] = (next - k) / (next - center);
      }
    }

    return filterbank;
  }

  private computeZCR(int16: number[]): number {
    let crossings = 0;
    for (let i = 1; i < int16.length; i++) {
      if ((int16[i] > 0 && int16[i - 1] <= 0) || (int16[i] <= 0 && int16[i - 1] > 0)) {
        crossings++;
      }
    }
    return crossings / int16.length;
  }

  private computeSpectralCentroid(int16: number[]): number {
    const n = int16.length;
    let num = 0, den = 0;
    for (let i = 0; i < n; i++) {
      num += i * int16[i] * int16[i];
      den += int16[i] * int16[i];
    }
    return den > 0 ? num / den : 0;
  }

  private computeSpectralRollOff(int16: number[]): number {
    const n = int16.length;
    let totalEnergy = 0;
    for (let i = 0; i < n; i++) {
      totalEnergy += int16[i] * int16[i];
    }
    const threshold = totalEnergy * 0.85;
    let cumulative = 0;
    for (let i = 0; i < n; i++) {
      cumulative += int16[i] * int16[i];
      if (cumulative >= threshold) return i;
    }
    return n - 1;
  }

  private writeTempFile(content: Buffer, ext: string): string {
    const file = join(tmpdir(), `vp_${randomBytes(8).toString("hex")}${ext}`);
    require("fs").writeFileSync(file, content);
    return file;
  }

  private async readFile(path: string): Promise<Buffer> {
    return require("fs").readFileSync(path);
  }
}

export interface PreprocessedAudio {
  audio: Buffer;
  sampleRate: number;
  channels: number;
  durationSec: number;
}

export interface AudioFeatures {
  mfccs: number[][];
  zeroCrossingRate: number;
  spectralCentroid: number;
  spectralRollOff: number;
  durationSec: number;
}
