/**
 * Demucs audio processor - Core separation logic
 */

import { CONSTANTS } from './constants.js';
import { stft, istft, reflectPad } from './fft.js';

const { SAMPLE_RATE, FFT_SIZE, HOP_SIZE, TRAINING_SAMPLES, MODEL_SPEC_BINS, MODEL_SPEC_FRAMES, SEGMENT_OVERLAP, TRACKS } = CONSTANTS;

/**
 * Convert model frequency output to complex spectrogram per track
 */
export function standaloneMask(freqOutput) {
  const numTracks = 4;
  const numChannels = 4;
  const numBins = MODEL_SPEC_BINS;
  const numFrames = MODEL_SPEC_FRAMES;
  const result = [];

  for (let t = 0; t < numTracks; t++) {
    const trackSpec = {
      leftReal: new Float32Array(numBins * numFrames),
      leftImag: new Float32Array(numBins * numFrames),
      rightReal: new Float32Array(numBins * numFrames),
      rightImag: new Float32Array(numBins * numFrames)
    };

    for (let f = 0; f < numFrames; f++) {
      for (let b = 0; b < numBins; b++) {
        const baseIdx = t * numChannels * numBins * numFrames;
        const outIdx = b * numFrames + f;
        trackSpec.leftReal[outIdx] = freqOutput[baseIdx + 0 * numBins * numFrames + b * numFrames + f];
        trackSpec.leftImag[outIdx] = freqOutput[baseIdx + 1 * numBins * numFrames + b * numFrames + f];
        trackSpec.rightReal[outIdx] = freqOutput[baseIdx + 2 * numBins * numFrames + b * numFrames + f];
        trackSpec.rightImag[outIdx] = freqOutput[baseIdx + 3 * numBins * numFrames + b * numFrames + f];
      }
    }
    result.push(trackSpec);
  }

  return result;
}

/**
 * Convert complex spectrogram back to time domain (iSTFT with proper offsets)
 */
export function standaloneIspec(trackSpec, targetLength) {
  const numBins = MODEL_SPEC_BINS;
  const numFrames = MODEL_SPEC_FRAMES;
  const hopLength = HOP_SIZE;
  const paddedBins = numBins + 1;
  const paddedFrames = numFrames + 4;

  const padChannel = (real, imag) => {
    const paddedReal = new Float32Array(paddedFrames * paddedBins);
    const paddedImag = new Float32Array(paddedFrames * paddedBins);

    for (let f = 0; f < numFrames; f++) {
      for (let b = 0; b < numBins; b++) {
        const srcIdx = b * numFrames + f;
        const dstFrame = f + 2;
        const dstIdx = dstFrame * paddedBins + b;
        paddedReal[dstIdx] = real[srcIdx];
        paddedImag[dstIdx] = imag[srcIdx];
      }
    }
    return { real: paddedReal, imag: paddedImag };
  };

  const leftPadded = padChannel(trackSpec.leftReal, trackSpec.leftImag);
  const rightPadded = padChannel(trackSpec.rightReal, trackSpec.rightImag);

  const centerPad = FFT_SIZE / 2;
  const pad = Math.floor(hopLength / 2) * 3;
  const istftLength = (paddedFrames - 1) * hopLength + FFT_SIZE;

  const leftOut = istft(leftPadded.real, leftPadded.imag, paddedFrames, paddedBins, FFT_SIZE, hopLength, istftLength);
  const rightOut = istft(rightPadded.real, rightPadded.imag, paddedFrames, paddedBins, FFT_SIZE, hopLength, istftLength);

  const totalOffset = centerPad + pad;
  const left = leftOut.subarray(totalOffset, totalOffset + targetLength);
  const right = rightOut.subarray(totalOffset, totalOffset + targetLength);

  return { left: new Float32Array(left), right: new Float32Array(right) };
}

/**
 * Prepare model input from stereo audio
 */
export function prepareModelInput(leftChannel, rightChannel) {
  const inputLength = TRAINING_SAMPLES;

  const paddedLeft = new Float32Array(inputLength);
  const paddedRight = new Float32Array(inputLength);
  const copyLen = Math.min(leftChannel.length, inputLength);
  paddedLeft.set(leftChannel.subarray(0, copyLen));
  paddedRight.set(rightChannel.subarray(0, copyLen));

  const le = Math.ceil(inputLength / HOP_SIZE);
  const pad = Math.floor(HOP_SIZE / 2) * 3;
  const padRight = pad + le * HOP_SIZE - inputLength;

  const stftInputLeft = reflectPad(paddedLeft, pad, padRight);
  const stftInputRight = reflectPad(paddedRight, pad, padRight);

  const centerPad = FFT_SIZE / 2;
  const centeredLeft = reflectPad(stftInputLeft, centerPad, centerPad);
  const centeredRight = reflectPad(stftInputRight, centerPad, centerPad);

  const stftLeft = stft(centeredLeft, FFT_SIZE, HOP_SIZE);
  const stftRight = stft(centeredRight, FFT_SIZE, HOP_SIZE);

  const numBins = MODEL_SPEC_BINS;
  const numFrames = MODEL_SPEC_FRAMES;
  const frameOffset = 2;

  const magSpec = new Float32Array(4 * numBins * numFrames);

  for (let f = 0; f < numFrames; f++) {
    const srcFrame = f + frameOffset;
    for (let b = 0; b < numBins; b++) {
      const srcIdx = srcFrame * stftLeft.numBins + b;
      magSpec[0 * numBins * numFrames + b * numFrames + f] = stftLeft.real[srcIdx];
      magSpec[1 * numBins * numFrames + b * numFrames + f] = stftLeft.imag[srcIdx];
      magSpec[2 * numBins * numFrames + b * numFrames + f] = stftRight.real[srcIdx];
      magSpec[3 * numBins * numFrames + b * numFrames + f] = stftRight.imag[srcIdx];
    }
  }

  const waveform = new Float32Array(2 * inputLength);
  waveform.set(paddedLeft, 0);
  waveform.set(paddedRight, inputLength);

  return { waveform, magSpec, numBins, numFrames, originalLength: leftChannel.length };
}

/**
 * Main Demucs processor class
 */
export class DemucsProcessor {
  constructor(options = {}) {
    this.ort = options.ort || null;
    this.session = null;
    this.modelPath = options.modelPath || './htdemucs_embedded.onnx';
    this.sessionOptions = options.sessionOptions || {};
    this.onProgress = options.onProgress || (() => {});
    this.onLog = options.onLog || (() => {});
    this.onDownloadProgress = options.onDownloadProgress || (() => {});
  }

  async loadModel(modelPathOrBuffer) {
    if (!this.ort) {
      throw new Error('ONNX Runtime not provided. Pass ort in constructor options.');
    }

    this.onLog('model', 'Loading model...');

    let modelBuffer;
    if (modelPathOrBuffer instanceof ArrayBuffer) {
      modelBuffer = modelPathOrBuffer;
    } else {
      const response = await fetch(modelPathOrBuffer || this.modelPath);

      // Check if we can track progress
      const contentLength = response.headers.get('Content-Length');
      if (contentLength && response.body) {
        const totalSize = parseInt(contentLength, 10);
        const reader = response.body.getReader();
        const chunks = [];
        let loadedSize = 0;

        while (true) {
          const { done, value } = await reader.read();
          if (done) break;
          chunks.push(value);
          loadedSize += value.length;
          this.onDownloadProgress(loadedSize, totalSize);
        }

        // Combine chunks into single ArrayBuffer
        const combined = new Uint8Array(loadedSize);
        let offset = 0;
        for (const chunk of chunks) {
          combined.set(chunk, offset);
          offset += chunk.length;
        }
        modelBuffer = combined.buffer;
      } else {
        // Fallback: no progress tracking
        modelBuffer = await response.arrayBuffer();
      }
    }

    const defaultSessionOptions = {
      executionProviders: ['webgpu', 'wasm'],
      graphOptimizationLevel: 'basic'
    };

    this.session = await this.ort.InferenceSession.create(modelBuffer, {
      ...defaultSessionOptions,
      ...this.sessionOptions
    });

    this.onLog('model', 'Model loaded successfully');
    return this.session;
  }

  async separate(leftChannel, rightChannel) {
    if (!this.session) {
      throw new Error('Model not loaded. Call loadModel() first.');
    }

    const totalSamples = leftChannel.length;
    const stride = Math.floor(TRAINING_SAMPLES * (1 - SEGMENT_OVERLAP));
    const numSegments = Math.ceil((totalSamples - TRAINING_SAMPLES) / stride) + 1;

    const outputs = TRACKS.map(() => ({
      left: new Float32Array(totalSamples),
      right: new Float32Array(totalSamples)
    }));
    const weights = new Float32Array(totalSamples);

    let segmentIdx = 0;

    for (let start = 0; start < totalSamples; start += stride) {
      const end = Math.min(start + TRAINING_SAMPLES, totalSamples);
      const segmentLength = end - start;

      const segLeft = new Float32Array(TRAINING_SAMPLES);
      const segRight = new Float32Array(TRAINING_SAMPLES);

      for (let i = 0; i < segmentLength; i++) {
        segLeft[i] = leftChannel[start + i];
        segRight[i] = rightChannel[start + i];
      }

      const input = prepareModelInput(segLeft, segRight);

      const waveformTensor = new this.ort.Tensor('float32', input.waveform, [1, 2, TRAINING_SAMPLES]);
      const magSpecTensor = new this.ort.Tensor('float32', input.magSpec, [1, 4, MODEL_SPEC_BINS, MODEL_SPEC_FRAMES]);

      const feeds = {};
      feeds[this.session.inputNames[0]] = waveformTensor;
      if (this.session.inputNames.length > 1) {
        feeds[this.session.inputNames[1]] = magSpecTensor;
      }

      const inferResults = await this.session.run(feeds);

      let timeData = null, timeShape = null;
      let freqData = null;

      for (const name of this.session.outputNames) {
        const tensor = inferResults[name];
        if (tensor.dims.length === 4 && tensor.dims[2] === 2) {
          timeData = tensor.data;
          timeShape = tensor.dims;
        } else if (tensor.dims.length === 5 && tensor.dims[2] === 4) {
          freqData = tensor.data;
        }
      }

      if (!timeData) {
        throw new Error('Could not find time-domain output tensor');
      }

      let combinedOutputs = null;
      if (freqData) {
        const trackSpecs = standaloneMask(freqData);
        combinedOutputs = [];

        for (let t = 0; t < 4; t++) {
          const freqOutput = standaloneIspec(trackSpecs[t], TRAINING_SAMPLES);
          const numChannels = timeShape[2];
          const samples = timeShape[3];
          const timeLeft = new Float32Array(samples);
          const timeRight = new Float32Array(samples);

          for (let i = 0; i < samples; i++) {
            timeLeft[i] = timeData[t * numChannels * samples + 0 * samples + i];
            timeRight[i] = timeData[t * numChannels * samples + 1 * samples + i];
          }

          const combined = {
            left: new Float32Array(samples),
            right: new Float32Array(samples)
          };
          for (let i = 0; i < samples; i++) {
            combined.left[i] = timeLeft[i] + (freqOutput.left[i] || 0);
            combined.right[i] = timeRight[i] + (freqOutput.right[i] || 0);
          }
          combinedOutputs.push(combined);
        }
      }

      const numTracks = timeShape[1];
      const numChannels = timeShape[2];
      const samples = timeShape[3];

      const overlapWindow = new Float32Array(segmentLength);
      for (let i = 0; i < segmentLength; i++) {
        const fadeIn = Math.min(i / (stride * 0.5), 1);
        const fadeOut = Math.min((segmentLength - i) / (stride * 0.5), 1);
        overlapWindow[i] = Math.min(fadeIn, fadeOut);
      }

      for (let t = 0; t < numTracks; t++) {
        for (let i = 0; i < segmentLength && start + i < totalSamples; i++) {
          let leftVal, rightVal;
          if (combinedOutputs) {
            leftVal = combinedOutputs[t].left[i];
            rightVal = combinedOutputs[t].right[i];
          } else {
            const leftIdx = t * numChannels * samples + 0 * samples + i;
            const rightIdx = t * numChannels * samples + 1 * samples + i;
            leftVal = timeData[leftIdx];
            rightVal = timeData[rightIdx];
          }
          outputs[t].left[start + i] += leftVal * overlapWindow[i];
          outputs[t].right[start + i] += rightVal * overlapWindow[i];
        }
      }

      for (let i = 0; i < segmentLength && start + i < totalSamples; i++) {
        weights[start + i] += overlapWindow[i];
      }

      segmentIdx++;
      this.onProgress({
        progress: segmentIdx / numSegments,
        currentSegment: segmentIdx,
        totalSegments: numSegments
      });
    }

    for (let t = 0; t < TRACKS.length; t++) {
      for (let i = 0; i < totalSamples; i++) {
        if (weights[i] > 0) {
          outputs[t].left[i] /= weights[i];
          outputs[t].right[i] /= weights[i];
        }
      }
    }

    return {
      drums: outputs[0],
      bass: outputs[1],
      other: outputs[2],
      vocals: outputs[3]
    };
  }
}