import { Injectable, OnModuleInit } from '@nestjs/common';
import * as ort from 'onnxruntime-web';
import { Jimp } from 'jimp';
import * as path from 'path';
import { IScannerProvider, InferenceTensor, ScanResult } from './scanner.interface';
import { postprocessShared } from './onnx-native.provider';

// Single-threaded WASM — safer on low-resource / ARM environments (Android/Termux)
ort.env.wasm.numThreads = 1;

@Injectable()
export class OnnxWasmProvider implements IScannerProvider, OnModuleInit {
  private session!: ort.InferenceSession;
  private readonly modelPath = path.join(process.cwd(), 'best.onnx');

  async onModuleInit() {
    try {
      this.session = await ort.InferenceSession.create(this.modelPath, {
        executionProviders: ['wasm'],
      });
      console.log('✅ [onnx-wasm] Inference session initialized:', this.modelPath);
    } catch (error) {
      console.error('❌ [onnx-wasm] Failed to initialize:', error);
      throw error;
    }
  }

  async preprocess(imageBuffer: Buffer): Promise<InferenceTensor> {
    const img = await Jimp.read(imageBuffer);
    img.resize({ w: 640, h: 640 });

    const pixels = img.bitmap.data;
    const imageSize = 640 * 640;
    const floatData = new Float32Array(3 * imageSize);

    for (let i = 0; i < imageSize; i++) {
      floatData[i] = pixels[i * 4] / 255.0;
      floatData[i + imageSize] = pixels[i * 4 + 1] / 255.0;
      floatData[i + 2 * imageSize] = pixels[i * 4 + 2] / 255.0;
    }

    return { data: floatData, dims: [1, 3, 640, 640] };
  }

  async inference(tensor: InferenceTensor): Promise<InferenceTensor> {
    const ortTensor = new ort.Tensor('float32', tensor.data, [1, 3, 640, 640]);
    const outputs = await this.session.run({ images: ortTensor });
    const out = outputs[Object.keys(outputs)[0]];
    return { data: out.data as Float32Array, dims: out.dims };
  }

  async postprocess(
    tensor: InferenceTensor,
    originalWidth: number,
    originalHeight: number,
    threshold = 0.25,
  ): Promise<ScanResult> {
    return postprocessShared(tensor, originalWidth, originalHeight, threshold);
  }
}