import * as tf from "@tensorflow/tfjs";
import type { Prediction, WorkerInMessage, WorkerOutMessage } from "./types";

const MODEL_PATH = "/models/birdnet/model.json";
const WINDOW_SAMPLES = 144000;
const ALPHA = 5.0;

let model: tf.LayersModel | null = null;

// ── Custom MelSpecLayerSimple ──────────────────────────────────────────────
// Ported from https://github.com/birdnet-team/real-time-pwa (MIT)
class MelSpecLayerSimple extends tf.layers.Layer {
  sampleRate: number;
  specShape: number[];
  frameStep: number;
  frameLength: number;
  melFilterbank: tf.Tensor2D;
  magScale!: tf.LayerVariable;

  constructor(config: Record<string, unknown>) {
    super(config as tf.serialization.ConfigDict);
    this.sampleRate = config.sampleRate as number;
    this.specShape = config.specShape as number[];
    this.frameStep = config.frameStep as number;
    this.frameLength = config.frameLength as number;
    this.melFilterbank = tf.tensor2d(config.melFilterbank as number[][]);
  }

  build(_inputShape: tf.Shape | tf.Shape[]) {
    this.magScale = this.addWeight(
      "magnitude_scaling",
      [],
      "float32",
      tf.initializers.constant({ value: 1.23 }),
    );
    super.build(_inputShape);
  }

  computeOutputShape(inputShape: tf.Shape): tf.Shape {
    return [inputShape[0], this.specShape[0], this.specShape[1], 1];
  }

  call(inputs: tf.Tensor | tf.Tensor[]): tf.Tensor {
    return tf.tidy(() => {
      const x = Array.isArray(inputs) ? inputs[0] : inputs;
      const frameLength = this.frameLength;
      const frameStep = this.frameStep;
      return tf.stack(
        x.split(x.shape[0]).map((input) => {
          let spec = input.squeeze();
          spec = tf.sub(spec, tf.min(spec, -1, true));
          spec = tf.div(spec, tf.max(spec, -1, true).add(1e-6));
          spec = tf.sub(spec, 0.5).mul(2.0);
          // eslint-disable-next-line @typescript-eslint/no-explicit-any
          spec = (tf.engine() as any).runKernel("STFT", {
            signal: spec,
            frameLength,
            frameStep,
          });
          spec = tf.matMul(spec as tf.Tensor2D, this.melFilterbank).pow(2.0);
          spec = spec.pow(
            tf.div(1.0, tf.add(1.0, tf.exp(this.magScale.read()))),
          );
          spec = tf.reverse(spec, -1);
          spec = tf.transpose(spec).expandDims(-1);
          return spec;
        }),
      );
    });
  }

  static get className() {
    return "MelSpecLayerSimple";
  }
}

// ── Custom STFT WebGL kernel ───────────────────────────────────────────────
// Ported from https://github.com/birdnet-team/real-time-pwa (MIT)
function registerStftKernel() {
  const kernelFunc: tf.KernelFunc = ({ backend, inputs }) => {
    const b = backend as Record<string, (...args: unknown[]) => unknown>;
    const { signal, frameLength, frameStep } = inputs;
    const fl = frameLength as unknown as number;
    const fs = frameStep as unknown as number;
    const innerDim = fl / 2;
    const log2Inner = Math.log2(innerDim);
    const batch =
      ((signal as unknown as { size: number }).size - fl + fs) / fs | 0;

    // Stage 1: windowing + bit-reversal
    let cur = b.runWebGLProgram(
      {
        variableNames: ["x"],
        outputShape: [batch, fl],
        userCode: `void main(){
            ivec2 c=getOutputCoords();
            int p=c[1]%${innerDim};
            int k=0;
            for(int i=0;i<${log2Inner};++i){
              if((p & (1<<i))!=0){ k|=(1<<(${log2Inner - 1}-i)); }
            }
            int i=2*k;
            if(c[1]>=${innerDim}){ i=2*(k%${innerDim})+1; }
            int q=c[0]*${fl}+i;
            float val=getX((q/${fl})*${fs}+ q % ${fl});
            float cosArg=${(2.0 * Math.PI) / fl}*float(q);
            float mul=0.5-0.5*cos(cosArg);
            setOutput(val*mul);
          }`,
      } as unknown,
      [signal],
      "float32",
    ) as unknown;

    // Stage 2: FFT butterflies
    for (let len = 1; len < innerDim; len *= 2) {
      const prev = cur;
      cur = b.runWebGLProgram(
        {
          variableNames: ["x"],
          outputShape: [batch, innerDim * 2],
          userCode: `void main(){
              ivec2 c=getOutputCoords();
              int b=c[0];
              int i=c[1];
              int k=i%${innerDim};
              int isHigh=(k%${len * 2})/${len};
              int highSign=(1 - isHigh*2);
              int baseIndex=k - isHigh*${len};
              float t=${Math.PI / len}*float(k%${len});
              float a=cos(t);
              float bsin=sin(-t);
              float oddK_re=getX(b, baseIndex+${len});
              float oddK_im=getX(b, baseIndex+${len + innerDim});
              if(i<${innerDim}){
                float evenK_re=getX(b, baseIndex);
                setOutput(evenK_re + (oddK_re*a - oddK_im*bsin)*float(highSign));
              } else {
                float evenK_im=getX(b, baseIndex+${innerDim});
                setOutput(evenK_im + (oddK_re*bsin + oddK_im*a)*float(highSign));
              }
            }`,
        } as unknown,
        [prev],
        "float32",
      ) as unknown;
      b.disposeIntermediateTensorInfo(prev);
    }

    // Stage 3: real RFFT output
    const real = b.runWebGLProgram(
      {
        variableNames: ["x"],
        outputShape: [batch, innerDim + 1],
        userCode: `void main(){
            ivec2 c=getOutputCoords();
            int b=c[0];
            int i=c[1];
            int zI=i%${innerDim};
            int conjI=(${innerDim}-i)%${innerDim};
            float Zk0=getX(b,zI);
            float Zk1=getX(b,zI+${innerDim});
            float Zk_conj0=getX(b,conjI);
            float Zk_conj1=-getX(b,conjI+${innerDim});
            float t=${-2.0 * Math.PI}*float(i)/float(${innerDim * 2});
            float diff0=Zk0 - Zk_conj0;
            float diff1=Zk1 - Zk_conj1;
            float result=(Zk0+Zk_conj0 + cos(t)*diff1 + sin(t)*diff0)*0.5;
            setOutput(result);
          }`,
      } as unknown,
      [cur],
      "float32",
    ) as unknown;
    b.disposeIntermediateTensorInfo(cur);
    return real as tf.TensorInfo;
  };
  tf.registerKernel({ kernelName: "STFT", backendName: "webgl", kernelFunc });
}

// ── Init ───────────────────────────────────────────────────────────────────

async function init() {
  try {
    await tf.setBackend("webgl");
    registerStftKernel();
    tf.serialization.registerClass(MelSpecLayerSimple);

    model = await tf.loadLayersModel(MODEL_PATH);
    console.log(
      "BirdNET model inputs:",
      model.inputs.map((t) => t.shape),
      "outputs:",
      model.outputs.map((t) => t.shape),
    );

    tf.tidy(() => {
      (model as tf.LayersModel).predict(tf.zeros([1, WINDOW_SAMPLES]));
    });

    const out: WorkerOutMessage = { type: "ready" };
    self.postMessage(out);
  } catch (err) {
    const out: WorkerOutMessage = {
      type: "error",
      message: String(err),
    };
    self.postMessage(out);
  }
}

// ── Analyze ────────────────────────────────────────────────────────────────

async function analyze(samples: Float32Array) {
  if (!model) return;
  try {
    const audioTensor = tf.tensor2d(samples, [1, WINDOW_SAMPLES]);
    const resTensor = model.predict(audioTensor) as tf.Tensor;
    const rawPreds = (await resTensor.array()) as number[][];
    resTensor.dispose();
    audioTensor.dispose();

    const frame = rawPreds[0];
    const sumsExp = frame.map((p) => Math.exp(ALPHA * p));
    const pooled = sumsExp.map((s) => Math.log(s) / ALPHA);

    const predictions: Prediction[] = pooled
      .map((confidence, labelIndex) => ({ confidence, labelIndex }))
      .filter((p) => p.confidence > 0.1)
      .sort((a, b) => b.confidence - a.confidence)
      .slice(0, 10);

    const out: WorkerOutMessage = { type: "results", predictions };
    self.postMessage(out);
  } catch (err) {
    const out: WorkerOutMessage = { type: "error", message: String(err) };
    self.postMessage(out);
  }
}

self.addEventListener("message", (e: MessageEvent<WorkerInMessage>) => {
  const msg = e.data;
  if (msg.type === "init") {
    init();
  } else if (msg.type === "analyze") {
    analyze(msg.samples);
  }
});