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 { MPOB_CLASSES, HEALTH_ALERT_CLASSES } from '../constants/mpob-standards';
import { DetectionResult } from '../interfaces/palm-analysis.interface';

// Single-threaded WASM — ARM/Termux safe
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);
  }
}

export function postprocessShared(
  outputTensor: InferenceTensor,
  originalWidth: number,
  originalHeight: number,
  threshold: number,
): ScanResult {
  const data = outputTensor.data;

  const sampleRows = Math.min(5, outputTensor.dims[1]);
  const raw_tensor_sample: number[][] = [];
  for (let i = 0; i < sampleRows; i++) {
    const offset = i * 6;
    raw_tensor_sample.push([
      parseFloat(data[offset].toFixed(6)),
      parseFloat(data[offset + 1].toFixed(6)),
      parseFloat(data[offset + 2].toFixed(6)),
      parseFloat(data[offset + 3].toFixed(6)),
      parseFloat(data[offset + 4].toFixed(6)),
      parseFloat(data[offset + 5].toFixed(6)),
    ]);
  }

  const results: DetectionResult[] = [];
  const numCandidates = outputTensor.dims[1];

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

    const className = MPOB_CLASSES[Math.round(data[offset + 5])] || 'Unknown';

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

    results.push({
      bunch_id: results.length + 1,
      class: className,
      confidence: parseFloat(confidence.toFixed(4)),
      is_health_alert: HEALTH_ALERT_CLASSES.includes(className),
      norm_box: [nx1, ny1, nx2, ny2],
      box: [
        nx1 * originalWidth,
        ny1 * originalHeight,
        nx2 * originalWidth,
        ny2 * originalHeight,
      ],
    });
  }

  return { detections: results, raw_tensor_sample };
}
