import fs from 'node:fs'
import path from 'node:path'

import { inferU2Mask } from '../lib/segmentation/u2net'

function computeCoverage(mask: Uint8Array): { solid: number; coverage: number; mean: number } {
  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 resolveModelCandidates(): string[] {
  const cwd = process.cwd()
  const projectRoot = path.resolve(__dirname, '..')
  const priority = (process.env.U2NET_MODEL_PRIORITY ?? 'u2net.onnx,u2netp_finetuned.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<string>()
  const results: string[] = []

  const pushIfExists = (candidate: string | null | undefined) => {
    if (!candidate) return
    const absolute = path.isAbsolute(candidate) ? candidate : path.join(projectRoot, candidate)
    if (seen.has(absolute)) return
    if (fs.existsSync(absolute)) {
      seen.add(absolute)
      results.push(absolute)
    }
  }

  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
}

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 buf = fs.readFileSync(imagePath)
  const started = Date.now()
  const mask = await inferU2Mask(buf)
  const elapsed = Date.now() - started
  const stats = computeCoverage(mask)

  console.log(JSON.stringify({
    imagePath,
    elapsedMs: elapsed,
    pixels: mask.length,
    coverage: Number(stats.coverage.toFixed(4)),
    mean: Number(stats.mean.toFixed(4)),
    modelCandidates: resolveModelCandidates()
  }, null, 2))
}

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