#!/usr/bin/env node
const fs = require('node:fs')
const path = require('node:path')
const sharp = require('sharp')
const ort = require('onnxruntime-node')

function listCandidates() {
  const cwd = process.cwd()
  const projectRoot = path.resolve(__dirname, '..')
  const priority = (process.env.U2NET_MODEL_PRIORITY ?? 'u2net_finetune.onnx,u2net.onnx,u2netp.onnx')
    .split(',')
    .map((entry) => entry.trim())
    .filter(Boolean)
  const explicit = process.env.U2NET_MODEL_PATH ? [process.env.U2NET_MODEL_PATH] : []
  const seen = new Set()
  const results = []

  const pushIfExists = (candidate) => {
    if (!candidate) return
    const absolute = path.isAbsolute(candidate) ? candidate : path.join(projectRoot, candidate)
    if (seen.has(absolute)) return
    try {
      if (fs.existsSync(absolute)) {
        seen.add(absolute)
        results.push(absolute)
      }
    } catch {
      // ignore
    }
  }

  explicit.forEach((entry) => pushIfExists(entry))

  for (const name of priority) {
    pushIfExists(path.join('models', name))
    pushIfExists(path.join(cwd, 'models', name))
    pushIfExists(path.join(__dirname, '..', 'models', name))
  }

  if (results.length === 0) {
    const fallback = path.join(projectRoot, 'models', 'u2netp.onnx')
    if (fs.existsSync(fallback)) results.push(fallback)
  }

  return results
}

function toCHW(rgbBuffer) {
  const pixels = rgbBuffer.length / 3
  const chw = new Float32Array(pixels * 3)
  for (let i = 0, p = 0; i < pixels; i++, p += 3) {
    const r = rgbBuffer[p] / 255
    const g = rgbBuffer[p + 1] / 255
    const b = rgbBuffer[p + 2] / 255
    chw[i] = r
    chw[i + pixels] = g
    chw[i + pixels * 2] = b
  }
  return chw
}

function sigmoidToMask(data) {
  const out = new Uint8Array(data.length)
  for (let i = 0; i < data.length; i++) {
    let v = data[i]
    if (!Number.isFinite(v)) v = 0
    v = 1 / (1 + Math.exp(-v))
    out[i] = Math.max(0, Math.min(255, Math.round(v * 255)))
  }
  return out
}

async function inferScaledMask(sharpImg, width, height, size, session, mode = 'letterbox') {
  let padLeft = 0, padTop = 0, padRight = 0, padBottom = 0
  let effectiveW = size
  let effectiveH = size
  let raw

  if (mode === 'letterbox') {
    const scale = size / Math.max(width, height)
    effectiveW = Math.max(1, Math.min(size, Math.round(width * scale)))
    effectiveH = Math.max(1, Math.min(size, Math.round(height * scale)))
    padLeft = Math.floor((size - effectiveW) / 2)
    padTop = Math.floor((size - effectiveH) / 2)
    padRight = Math.max(0, size - effectiveW - padLeft)
    padBottom = Math.max(0, size - effectiveH - padTop)

    raw = await sharpImg
      .clone()
      .resize(effectiveW, effectiveH, { fit: 'fill' })
      .extend({
        top: padTop,
        bottom: padBottom,
        left: padLeft,
        right: padRight,
        background: { r: 0, g: 0, b: 0 },
      })
      .raw()
      .toBuffer()
  } else {
    raw = await sharpImg.clone().resize(size, size, { fit: 'fill' }).raw().toBuffer()
  }

  const inputName = session.inputNames[0]
  const tensor = new ort.Tensor('float32', toCHW(raw), [1, 3, size, size])
  const outputs = await session.run({ [inputName]: tensor })
  const firstKey = Object.keys(outputs)[0]
  const maskSmall = sigmoidToMask(outputs[firstKey].data)

  let maskSharp = sharp(Buffer.from(maskSmall), { raw: { width: size, height: size, channels: 1 } })

  if (mode === 'letterbox' && (padLeft || padTop || padRight || padBottom)) {
    const extractWidth = Math.max(1, size - padLeft - padRight)
    const extractHeight = Math.max(1, size - padTop - padBottom)
    maskSharp = maskSharp.extract({ left: padLeft, top: padTop, width: extractWidth, height: extractHeight })
  }

  const up = await maskSharp
    .resize(width, height, { kernel: 'lanczos3' })
    .toColorspace('b-w')
    .raw()
    .toBuffer()

  return new Uint8Array(up)
}

function computeCoverage(mask) {
  let solid = 0
  let sum = 0
  for (let i = 0; i < mask.length; i++) {
    const v = mask[i]
    sum += v
    if (v >= 128) solid += 1
  }
  const total = mask.length || 1
  return {
    solid,
    coverage: solid / total,
    mean: sum / (total * 255),
  }
}

function resolveSessionInputSize(session, requested) {
  const safeRequested = Number.isFinite(requested) && requested > 0 ? Math.round(requested) : 320
  try {
    const inputName = session.inputNames?.[0]
    const dims = inputName ? session.inputMetadata?.[inputName]?.dimensions : undefined
    if (Array.isArray(dims) && dims.length >= 4) {
      const h = dims[dims.length - 2]
      const w = dims[dims.length - 1]
      const candidate = typeof h === 'number' && h > 0 ? h : typeof w === 'number' && w > 0 ? w : null
      if (candidate && Number.isFinite(candidate) && candidate > 0) {
        return Math.max(64, Math.round(candidate))
      }
    }
  } catch {
    // ignore shape introspection errors
  }
  return Math.max(64, safeRequested)
}

function extractExpectedSize(error) {
  const msg = String((error && error.message) || error || '')
  const match = msg.match(/Expected:\s*(\d+)/i)
  if (match) {
    const parsed = Number(match[1])
    if (Number.isFinite(parsed) && parsed > 0) {
      return Math.max(64, Math.round(parsed))
    }
  }
  return null
}

async function runForModel(modelPath, imagePath) {
  const img = sharp(imagePath).removeAlpha()
  const meta = await img.metadata()
  const width = meta.width ?? 0
  const height = meta.height ?? 0
  if (!(width > 0 && height > 0)) throw new Error('Invalid image dimensions')

  let session
  const startLoad = Date.now()
  try {
    session = await ort.InferenceSession.create(modelPath, {
      executionProviders: ['cpuExecutionProvider'],
      graphOptimizationLevel: 'all',
    })
  } catch {
    session = await ort.InferenceSession.create(modelPath)
  }
  const loadMs = Date.now() - startLoad

  const requested = Number(process.env.U2NET_PRIMARY_SIZE ?? 480) || 480
  const initialSize = resolveSessionInputSize(session, requested)

  const attempt = async (target) => {
    const start = Date.now()
    const mask = await inferScaledMask(img, width, height, target, session, 'letterbox')
    return { mask, inferMs: Date.now() - start, size: target }
  }

  let inferResult
  try {
    inferResult = await attempt(initialSize)
  } catch (error) {
    const expected = extractExpectedSize(error)
    if (expected && expected !== initialSize) {
      inferResult = await attempt(expected)
    } else if (initialSize !== 320) {
      inferResult = await attempt(320)
    } else {
      throw error
    }
  }
  const stats = computeCoverage(inferResult.mask)

  return {
    modelPath,
    loadMs,
    inferMs: inferResult.inferMs,
    inputSize: inferResult.size,
    coverage: Number(stats.coverage.toFixed(4)),
    mean: Number(stats.mean.toFixed(4)),
  }
}

async function main() {
  const imagePath = process.argv[2] ?? path.join(__dirname, '..', 'node_modules', '@jimp', 'plugin-color', 'test', 'images', 'tiles.jpg')
  if (!fs.existsSync(imagePath)) {
    console.error(`Image not found: ${imagePath}`)
    process.exit(1)
  }

  const candidates = listCandidates()
  if (!candidates.length) {
    console.error('No candidates found')
    process.exit(1)
  }

  const results = []
  for (const modelPath of candidates) {
    try {
      results.push(await runForModel(modelPath, imagePath))
    } catch (error) {
      results.push({ modelPath, error: error instanceof Error ? error.message : String(error) })
    }
  }

  console.log(JSON.stringify({ imagePath, candidates, results }, null, 2))
}

main().catch((error) => {
  console.error(error)
  process.exit(1)
})