/// <reference lib="webworker" />
import * as ort from 'onnxruntime-web';

ort.env.wasm.wasmPaths = '/assets/wasm/';
ort.env.wasm.numThreads = 1;

let onnxSession: ort.InferenceSession | null = null;
let tfliteSession: any = null;

async function initOnnxSession(): Promise<ort.InferenceSession> {
  if (onnxSession) return onnxSession;
  onnxSession = await ort.InferenceSession.create('/assets/models/onnx/best.onnx', {
    executionProviders: ['wasm'],
  });
  return onnxSession;
}

async function initTfliteSession(): Promise<any> {
  if (tfliteSession) return tfliteSession;

  if (!(globalThis as any).Module) {
    (globalThis as any).Module = {
      locateFile: (path: string) => {
        // Direct the internal loader to fetch the correct compiled _cc binaries explicitly
        return `/assets/tflite-wasm/${path}`;
      },
    };
  }

  if (!(globalThis as any).tfweb) {
    // Mask the path in a runtime variable to hide it from esbuild static analysis
    const runtimeAssetPath = '/assets/tflite-wasm/tflite_web_api_client.js';
    await import(/* @vite-ignore */ runtimeAssetPath as any);
  }

  const tfwebGlobal = (globalThis as any).tfweb;
  if (!tfwebGlobal || !tfwebGlobal.TFLiteWebModelRunner) {
    throw new Error('[TFLite] TFLiteWebModelRunner constructor not found on global scope');
  }

  if (tfwebGlobal.tflite_web_api && typeof tfwebGlobal.tflite_web_api.setWasmPath === 'function') {
    tfwebGlobal.tflite_web_api.setWasmPath('/assets/tflite-wasm/');
  }

  tfliteSession = await tfwebGlobal.TFLiteWebModelRunner.create(
    '/assets/models/tflite/best_float16.tflite',
  );
  return tfliteSession;
}

const MPOB_CLASSES = ['Empty_Bunch', 'Underripe', 'Abnormal', 'Ripe', 'Unripe', 'Overripe'];
const CONF_THRESHOLD = 0.25;

addEventListener('message', async ({ data }) => {
  const { frameId, batchId, imageDataUrl, tensor, processingStart, mode } = data;

  try {
    const inferenceStart = performance.now();
    let detections: any[] = [];
    const industrialSummary: Record<string, number> = {
      Empty_Bunch: 0, Underripe: 0, Abnormal: 0, Ripe: 0, Unripe: 0, Overripe: 0,
    };

    // ── Engine Routing Branch Matrix ──────────────────────────────────────────
    if (mode === 'local-tflite') {
      const runner = await initTfliteSession();

      // Reconstruct typed view over the transferred ArrayBuffer — prevents undefined/NaN indexing
      const rawPlanarData = new Float32Array(tensor);
      const pixelCount = 640 * 640;

      // Build a clean intermediate HWC buffer before touching WASM memory
      const hwcData = new Float32Array(pixelCount * 3);
      for (let i = 0; i < pixelCount; i++) {
        hwcData[i * 3]     = rawPlanarData[i];                  // R
        hwcData[i * 3 + 1] = rawPlanarData[pixelCount + i];     // G
        hwcData[i * 3 + 2] = rawPlanarData[2 * pixelCount + i]; // B
      }

      const inputTensor = runner.getInputs()[0];
      if (!inputTensor) {
        throw new Error('[TFLite] Missing input tensor description definition');
      }

      // Resolve the true underlying memory buffer — data may be a function closure over WASM heap
      const inputBuffer = typeof inputTensor.data === 'function'
        ? (inputTensor.data as any)()
        : inputTensor.data;

      if (inputBuffer) {
        if (typeof inputBuffer.set === 'function') {
          inputBuffer.set(hwcData);
        } else {
          for (let i = 0; i < hwcData.length; i++) {
            inputBuffer[i] = hwcData[i];
          }
        }
      } else {
        throw new Error('[TFLite] Failed to initialize a valid WebAssembly memory view for input data');
      }

      runner.infer();

      const outputs = runner.getOutputs();
      console.log('[TFLite Debug] Input Matrix Shape:', runner.getInputs()[0].shape);

      const outputTensor = outputs[0];
      let outputShape: number[] = [];

      if (typeof outputTensor.shape === 'string') {
        outputShape = (outputTensor.shape as any).split(',').map((v: string) => parseInt(v, 10));
      } else if (Array.isArray(outputTensor.shape)) {
        outputShape = outputTensor.shape;
      } else if (outputTensor.shape && typeof (outputTensor.shape as any).toArray === 'function') {
        outputShape = (outputTensor.shape as any).toArray();
      } else {
        outputShape = [1, 300, 6];
      }
      console.log('[TFLite Debug] Final Parsed Dimensions Array:', outputShape);

      // Native engine accessor API — bypasses function-closure indirection over WASM heap
      let rawValues: any = null;
      if (typeof (runner as any).getOutputTensorData === 'function') {
        rawValues = (runner as any).getOutputTensorData(0);
      } else if (typeof (runner as any).getOutputData === 'function') {
        rawValues = (runner as any).getOutputData(0);
      } else {
        rawValues = outputTensor && typeof outputTensor.data === 'function'
          ? (outputTensor.data as any)()
          : (outputTensor ? outputTensor.data : null);
      }
      console.log('[TFLite Debug] Real Array Check:', Array.isArray(rawValues), rawValues?.constructor?.name);
      const debugSlice = rawValues
        ? Array.from(rawValues.subarray ? rawValues.subarray(0, 12) : (rawValues.slice ? rawValues.slice(0, 12) : rawValues))
        : 'null';
      console.log('[TFLite Debug] Real Slice:', debugSlice);

      // Model has internal NMS — output is [1, 300, 6] post-NMS candidates
      // Bypass manual CONF_THRESHOLD: model already filtered; skip only zero-confidence empty slots
      const numCandidates = outputShape[1] || 300;
      const stride = outputShape[2] || 6;

      for (let i = 0; i < numCandidates; i++) {
        const offset = i * stride;
        const confidence = rawValues[offset + 4];

        // Relax the gateway filter to capture both live detections explicitly
        if (isNaN(confidence) || confidence < 0.20) {
          continue;
        }

        // Extract the true class index predicted by the WebAssembly model graph
        const classIdx = Math.round(rawValues[offset + 5]);
        const className = MPOB_CLASSES[classIdx] || 'Unknown';

        // Map TFLite NMS format [ymin, xmin, ymax, xmax] to standard UI layout [x1, y1, x2, y2]
        const ny1 = parseFloat(rawValues[offset + 0].toFixed(6)); // ymin
        const nx1 = parseFloat(rawValues[offset + 1].toFixed(6)); // xmin
        const ny2 = parseFloat(rawValues[offset + 2].toFixed(6)); // ymax
        const nx2 = parseFloat(rawValues[offset + 3].toFixed(6)); // xmax

        industrialSummary[className] = (industrialSummary[className] ?? 0) + 1;

        console.log(`[TFLite Loop Success] Pushing candidate ${i}, Conf: ${confidence.toFixed(4)}, Class: ${classIdx}`);
        detections.push({
          bunch_id: detections.length + 1,
          class: className,
          confidence: parseFloat(confidence.toFixed(4)),
          is_health_alert: ['Abnormal', 'Empty_Bunch'].includes(className),
          norm_box: [nx1, ny1, nx2, ny2],
          box: [nx1 * 640, ny1 * 640, nx2 * 640, ny2 * 640],
        });
      }
    } else {
      // Default to standard local ONNX execution stream
      const ortSession = await initOnnxSession();
      const floatData = new Float32Array(tensor);
      const inputTensor = new ort.Tensor('float32', floatData, [1, 3, 640, 640]);

      const outputs = await ortSession.run({ [ortSession.inputNames[0]]: inputTensor });
      const outputKey = Object.keys(outputs)[0];
      const outputData = outputs[outputKey].data as Float32Array;
      const outputDims = outputs[outputKey].dims; // Shape [1, numCandidates, 6]

      const numCandidates = outputDims[1] as number;

      for (let i = 0; i < numCandidates; i++) {
        const offset = i * 6;
        const confidence = outputData[offset + 4];
        if (confidence < CONF_THRESHOLD) continue;

        const classIdx = Math.round(outputData[offset + 5]);
        const className = MPOB_CLASSES[classIdx] ?? 'Unknown';

        const nx1 = parseFloat(outputData[offset].toFixed(6));
        const ny1 = parseFloat(outputData[offset + 1].toFixed(6));
        const nx2 = parseFloat(outputData[offset + 2].toFixed(6));
        const ny2 = parseFloat(outputData[offset + 3].toFixed(6));

        industrialSummary[className] = (industrialSummary[className] ?? 0) + 1;

        detections.push({
          bunch_id: detections.length + 1,
          class: className,
          confidence: parseFloat(confidence.toFixed(4)),
          is_health_alert: ['Abnormal', 'Empty_Bunch'].includes(className),
          norm_box: [nx1, ny1, nx2, ny2],
          box: [nx1 * 640, ny1 * 640, nx2 * 640, ny2 * 640],
        });
      }
    }

    const inferenceMs = performance.now() - inferenceStart;

    postMessage({
      frameId,
      batchId,
      imageDataUrl,
      detections,
      inference_ms: parseFloat(inferenceMs.toFixed(2)),
      processing_ms: parseFloat(Math.max(0, performance.now() - processingStart).toFixed(2)),
      total_count: detections.length,
      industrial_summary: industrialSummary,
      source: 'wasm-local',
    });

  } catch (error: any) {
    console.error('[InferenceWorker] Local calculation stack crash:', error);
    postMessage({
      frameId,
      batchId,
      imageDataUrl,
      detections: [],
      inference_ms: 0,
      processing_ms: parseFloat(Math.max(0, performance.now() - processingStart).toFixed(2)),
      total_count: 0,
      industrial_summary: {},
      source: 'wasm-local',
      error: error.message,
    });
  }
});
