import fs from "node:fs";
import path from "node:path";
import sharp from "sharp";
import * as ort from "onnxruntime-node";

type SessionEntry = {
  key: string;
  file: string;
  session: ort.InferenceSession;
};

const sessionCache = new Map<string, SessionEntry>();

const parseSize = (value: string | undefined, fallback: number): number => {
  const parsed = Number(value);
  if (!Number.isFinite(parsed)) return fallback;
  return Math.max(64, Math.round(parsed));
};

const PRIMARY_SIZE = parseSize(process.env.U2NET_PRIMARY_SIZE, 320);
const SECONDARY_SIZE = parseSize(
  process.env.U2NET_SECONDARY_SIZE,
  Math.max(384, PRIMARY_SIZE),
);
const COVERAGE_THRESHOLD = Math.min(
  0.95,
  Math.max(0.02, Number(process.env.U2NET_FALLBACK_THRESHOLD ?? "0.14")),
);
const GROW_RADIUS = Math.max(0, Number(process.env.U2NET_GROW_RADIUS ?? "7"));
const GROW_SEED_THRESHOLD = Math.max(
  0,
  Math.min(255, Number(process.env.U2NET_GROW_SEED ?? "120")),
);
const GROW_MIN_COVERAGE = Math.min(
  1,
  Math.max(0, Number(process.env.U2NET_GROW_MIN_COVERAGE ?? "0.28")),
);
const USE_ALL_MODELS = (process.env.U2NET_USE_ALL_MODELS ?? "0") === "1";
const STRETCH_MIN_SIDE = Math.max(1, Number(process.env.U2NET_STRETCH_MIN_SIDE ?? "80"));

const toCHW = (raw: Buffer, width: number, height: number): Float32Array => {
  const pixels = width * height;
  const chw = new Float32Array(pixels * 3);
  for (let i = 0, p = 0; i < pixels; i++, p += 3) {
    const r = raw[p] / 255;
    const g = raw[p + 1] / 255;
    const b = raw[p + 2] / 255;
    chw[i] = r;
    chw[i + pixels] = g;
    chw[i + pixels * 2] = b;
  }
  return chw;
};

const runSigmoid = (data: Float32Array): Uint8Array => {
  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 runU2Net(
  session: ort.InferenceSession,
  raw: Buffer,
  width: number,
  height: number,
): Promise<Uint8Array> {
  const chw = toCHW(raw, width, height);
  const inputName = session.inputNames[0];
  const input = new ort.Tensor("float32", chw, [1, 3, height, width]);
  const results = await session.run({ [inputName]: input });
  const firstOut = results[Object.keys(results)[0]];
  const data = (firstOut as ort.Tensor).data as Float32Array;
  return runSigmoid(data);
}

const coverageRatio = (
  mask: Uint8Array,
  width: number,
  height: number,
  threshold = 110,
): number => {
  if (width <= 0 || height <= 0) return 0;
  const total = width * height;
  let solid = 0;
  for (let i = 0; i < total; i++) {
    if (mask[i] >= threshold) solid += 1;
  }
  return total > 0 ? solid / total : 0;
};

type MaskMode = "letterbox" | "stretch";

function projectRootFromStandaloneDirname(dirname: string) {
  return path.resolve(dirname, "../../../../");
}

function listModelCandidates(): string[] {
  const projectRoot = projectRootFromStandaloneDirname(__dirname);
  const cwd = process.cwd();
  const priority = process.env.U2NET_MODEL_PRIORITY
    ? process.env.U2NET_MODEL_PRIORITY.split(",")
        .map((entry) => entry.trim())
        .filter(Boolean)
    : ["u2net_finetune.onnx", "u2net.onnx", "u2netp.onnx"];

  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;
    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));
    pushIfExists(path.join(__dirname, "../../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 loadSession(modelPath: string): Promise<SessionEntry> {
  const cached = sessionCache.get(modelPath);
  if (cached) return cached;

  if (!fs.existsSync(modelPath)) {
    throw new Error(`U²-Net model missing: ${modelPath}`);
  }

  let session: ort.InferenceSession;
  try {
    session = await ort.InferenceSession.create(modelPath, {
      executionProviders: ["cpuExecutionProvider"],
      graphOptimizationLevel: "all",
    } as ort.InferenceSessionOptions);
  } catch (error: any) {
    const msg = String(error?.message || error);
    if (/backend not found|executionprovider/i.test(msg)) {
      session = await ort.InferenceSession.create(modelPath);
    } else {
      throw error;
    }
  }

  const entry: SessionEntry = { key: modelPath, file: modelPath, session };
  sessionCache.set(modelPath, entry);
  return entry;
}

async function getSessions(): Promise<SessionEntry[]> {
  const candidates = listModelCandidates();
  if (candidates.length === 0) {
    throw new Error(
      "No U²-Net models found. Place an ONNX file under models/ or set U2NET_MODEL_PATH.",
    );
  }

  const sessions: SessionEntry[] = [];
  for (const candidate of candidates) {
    try {
      sessions.push(await loadSession(candidate));
    } catch (error) {
      console.warn(`[u2net] Failed to load model ${candidate}:`, error);
    }
  }

  if (sessions.length === 0) {
    throw new Error(`Unable to load any U²-Net models. Checked: ${candidates.join(", ")}`);
  }

  return sessions;
}

const pickInputSize = (entry: SessionEntry, fallback: number): number => {
  const safeFallback = Math.max(64, Math.round(fallback));
  try {
    const session = entry.session as unknown as {
      inputNames?: string[];
      inputMetadata?: Record<string, { dimensions?: any[] }>;
    };
    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 */
  }
  return safeFallback;
};

const extractExpectedSize = (error: unknown): number | null => {
  const msg = String((error as 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 inferScaledMask(
  img: sharp.Sharp,
  width: number,
  height: number,
  size: number,
  entry: SessionEntry,
  mode: MaskMode,
): Promise<Uint8Array> {
  const attempt = async (target: number): Promise<Uint8Array> => {
    const safeTarget = Math.max(64, Math.round(target));
    let raw: Buffer;
    let padLeft = 0;
    let padTop = 0;
    let padRight = 0;
    let padBottom = 0;
    let effectiveW = safeTarget;
    let effectiveH = safeTarget;

    if (mode === "letterbox") {
      const scale = safeTarget / Math.max(width, height);
      effectiveW = Math.max(1, Math.min(safeTarget, Math.round(width * scale)));
      effectiveH = Math.max(1, Math.min(safeTarget, Math.round(height * scale)));
      padLeft = Math.floor((safeTarget - effectiveW) / 2);
      padTop = Math.floor((safeTarget - effectiveH) / 2);
      padRight = Math.max(0, safeTarget - effectiveW - padLeft);
      padBottom = Math.max(0, safeTarget - effectiveH - padTop);
      raw = await img
        .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 img.clone().resize(safeTarget, safeTarget, { fit: "fill" }).raw().toBuffer();
    }

    const maskSmall = await runU2Net(entry.session, raw, safeTarget, safeTarget);
    let maskSharp = sharp(Buffer.from(maskSmall), {
      raw: { width: safeTarget, height: safeTarget, channels: 1 },
    });

    if (mode === "letterbox" && (padLeft || padTop || padRight || padBottom)) {
      const extractWidth = Math.max(1, safeTarget - padLeft - padRight);
      const extractHeight = Math.max(1, safeTarget - 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);
  };

  try {
    return await attempt(size);
  } catch (error) {
    const expected = extractExpectedSize(error);
    if (expected && expected !== size) {
      return await attempt(expected);
    }
    if (size !== 320) {
      try {
        return await attempt(320);
      } catch {
        /* ignore: fall back to throwing below */
      }
    }
    throw error instanceof Error ? error : new Error(String(error));
  }
}

const mergeMaskMax = (base: Uint8Array, extra: Uint8Array) => {
  const n = Math.min(base.length, extra.length);
  for (let i = 0; i < n; i++) {
    if (extra[i] > base[i]) base[i] = extra[i];
  }
};

function softGrowMaskInPlace(
  mask: Uint8Array,
  width: number,
  height: number,
  radius: number,
  seedThreshold: number,
) {
  if (radius <= 0 || width <= 0 || height <= 0) return;
  const total = width * height;
  if (total === 0) return;

  const seeds: number[] = [];
  const seen = new Uint8Array(total);
  let maxVal = 0;
  for (let i = 0; i < total; i++) {
    const v = mask[i];
    if (v >= seedThreshold && !seen[i]) {
      seeds.push(i);
      seen[i] = 1;
    }
    if (v > maxVal) maxVal = v;
  }
  if (!seeds.length) return;

  const r = Math.max(1, Math.round(radius));
  const radiusNorm = r + 0.5;
  const kernel: Array<{ dx: number; dy: number; value: number }> = [];
  for (let dy = -r; dy <= r; dy++) {
    for (let dx = -r; dx <= r; dx++) {
      const dist = Math.hypot(dx, dy);
      if (dist > r + 0.01) continue;
      const weight = Math.max(0, 1 - dist / radiusNorm);
      if (weight <= 0) continue;
      kernel.push({ dx, dy, value: Math.min(255, Math.round(weight * 255)) });
    }
  }
  if (!kernel.length) return;

  for (const seed of seeds) {
    const x0 = seed % width;
    const y0 = Math.floor(seed / width);
    for (const k of kernel) {
      const nx = x0 + k.dx;
      const ny = y0 + k.dy;
      if (nx < 0 || ny < 0 || nx >= width || ny >= height) continue;
      const idx = ny * width + nx;
      if (k.value > mask[idx]) mask[idx] = k.value;
    }
  }

  if (maxVal < 220) {
    const scale = 220 / Math.max(1, maxVal);
    for (let i = 0; i < total; i++) {
      mask[i] = Math.min(255, Math.round(mask[i] * scale));
    }
  }
}

export async function inferU2Mask(originalJpegOrPng: Buffer): Promise<Uint8Array> {
  const img = sharp(originalJpegOrPng).removeAlpha();
  const meta = await img.metadata();
  const width = meta.width ?? 0;
  const height = meta.height ?? 0;
  if (!(width > 0 && height > 0)) throw new Error("Bild hat keine gültigen Dimensionen");

  const sessions = await getSessions();
  const primaryEntry = sessions[0];
  const primarySize = pickInputSize(primaryEntry, PRIMARY_SIZE);

  let mask = await inferScaledMask(img, width, height, primarySize, primaryEntry, "letterbox");
  let coverage = coverageRatio(mask, width, height);

  if (USE_ALL_MODELS && sessions.length > 1) {
    for (let i = 1; i < sessions.length; i++) {
      const entry = sessions[i];
      const size = pickInputSize(entry, PRIMARY_SIZE);
      const extra = await inferScaledMask(img, width, height, size, entry, "letterbox");
      mergeMaskMax(mask, extra);
    }
    coverage = coverageRatio(mask, width, height);
  }

  if (coverage < COVERAGE_THRESHOLD && SECONDARY_SIZE > PRIMARY_SIZE) {
    const targets = USE_ALL_MODELS ? sessions : [primaryEntry];
    for (const entry of targets) {
      const size = pickInputSize(entry, SECONDARY_SIZE);
      const hi = await inferScaledMask(img, width, height, size, entry, "letterbox");
      mergeMaskMax(mask, hi);
    }
    coverage = coverageRatio(mask, width, height);
  }

  if (coverage < COVERAGE_THRESHOLD && Math.min(width, height) >= STRETCH_MIN_SIDE) {
    const targets = USE_ALL_MODELS ? sessions : [primaryEntry];
    for (const entry of targets) {
      const size = pickInputSize(entry, PRIMARY_SIZE);
      const stretch = await inferScaledMask(img, width, height, size, entry, "stretch");
      mergeMaskMax(mask, stretch);
    }
    coverage = coverageRatio(mask, width, height);
  }

  if (coverage < GROW_MIN_COVERAGE && GROW_RADIUS > 0) {
    softGrowMaskInPlace(mask, width, height, GROW_RADIUS, GROW_SEED_THRESHOLD);
  }

  return mask;
}

export async function featheredDeocclude(
  original: Buffer,
  mask: Uint8Array,
  blurSigma = 8,
): Promise<Buffer> {
  const base = sharp(original).removeAlpha();
  const meta = await base.metadata();
  const W = meta.width ?? 0;
  const H = meta.height ?? 0;
  if (mask.length !== W * H) throw new Error("Maskengröße passt nicht zum Bild");

  const maskPng = await sharp(Buffer.from(mask), { raw: { width: W, height: H, channels: 1 } })
    .blur(2.0)
    .png()
    .toBuffer();

  const blurred = await base.clone().blur(blurSigma).png().toBuffer();
  const overlay: any = { input: blurred, blend: "over", mask: { input: maskPng } };

  const composited = await sharp(await base.png().toBuffer())
    .composite([overlay])
    .png()
    .toBuffer();

  return composited;
}