import * as ort from 'onnxruntime-node';
import { existsSync } from 'node:fs';

const MODEL_PATH = process.env.INPAINT_MODEL_PATH || '/models/migan_pipeline_v2.onnx';

let sessionPromise: Promise<ort.InferenceSession | null> | null = null;

export async function getInpaintSession(): Promise<ort.InferenceSession | null> {
  if (!sessionPromise) {
    sessionPromise = (async () => {
      const candidates = [
        MODEL_PATH,
        '/home/locpham/RecipesCam/docker/frontend/public/models/migan_pipeline_v2.onnx',
        '../frontend/public/models/migan_pipeline_v2.onnx',
        './models/migan_pipeline_v2.onnx',
      ];
      let foundPath: string | null = null;
      for (const p of candidates) {
        if (existsSync(p)) {
          foundPath = p;
          break;
        }
      }
      if (!foundPath) {
        console.warn('[inpaintService] Model file not found in candidates:', candidates);
        return null;
      }

      console.log('[inpaintService] Loading ONNX model from:', foundPath);
      const t0 = Date.now();
      const session = await ort.InferenceSession.create(foundPath, {
        executionProviders: ['cpu'],
        graphOptimizationLevel: 'all',
      });
      console.log('[inpaintService] Model loaded in ' + (Date.now() - t0) + 'ms');
      return session;
    })().catch((err) => {
      console.error('[inpaintService] Failed to load ONNX model:', err);
      sessionPromise = null;
      return null;
    });
  }
  return sessionPromise;
}

export async function runServerInpaint(
  srcBytes: Uint8Array,
  maskBytes: Uint8Array,
  w: number,
  h: number
): Promise<Uint8Array | null> {
  const session = await getInpaintSession();
  if (!session) return null;

  const in0 = new ort.Tensor('uint8', srcBytes, [1, 3, h, w]);
  const in1 = new ort.Tensor('uint8', maskBytes, [1, 1, h, w]);

  const out = await session.run({
    [session.inputNames[0]]: in0,
    [session.inputNames[1]]: in1,
  });

  const tensor = out[session.outputNames[0]];
  const data = tensor.data as Uint8Array;
  const patch = new Uint8Array(4 * w * h);
  const plane = w * h;
  for (let i = 0; i < plane; i++) {
    patch[i * 4] = data[i];
    patch[i * 4 + 1] = data[plane + i];
    patch[i * 4 + 2] = data[2 * plane + i];
    patch[i * 4 + 3] = 255;
  }
  return patch;
}