/*
  author: Yagnik Poshiya
  organization: Webbrains Technologies Private Limited
*/

import { spawn } from 'node:child_process';
import { mkdtemp, readFile, rm, writeFile } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import path from 'node:path';

import ffmpegPath from 'ffmpeg-static';
import sharp from 'sharp';

/**
 * Server-side processing for animated avatar launchers (GIF/MP4).
 *
 * Many "GIF" exports (WhatsApp, screen recorders, etc.) are actually H.264
 * MP4s with the background baked into the pixels — H.264 cannot store an
 * alpha channel. This module converts every avatar upload into a transparent
 * animated GIF:
 *
 *   1. MP4 sources are first converted to animated GIF via ffmpeg.
 *   2. The dominant border color(s) of the first frame are detected
 *      (handles solid backdrops and two-tone "transparency" checkerboards).
 *   3. That background is removed with a flood fill that starts at the frame
 *      borders, so colors inside the character are never punched out.
 *   4. Frames are cropped to the union bounding box of visible pixels and
 *      scaled down so the result stays within the 1 MB avatar budget.
 *
 * GIFs that already contain transparency are passed through untouched.
 * When the background is too varied to key safely (e.g. photographic), the
 * asset is kept as-is rather than risking visible holes.
 */

/** Largest output dimension; matches the widget's avatar render cap. */
const MAX_AVATAR_DIMENSION_PX = 180;
/** Per-channel tolerance when matching a pixel against a key color. */
const KEY_COLOR_TOLERANCE = 20;
/** Dominant border colors must cover at least this fraction of the border. */
const BORDER_COVERAGE_THRESHOLD = 0.7;
/** Border already this transparent → asset needs no background removal. */
const TRANSPARENT_BORDER_FRACTION = 0.05;
/** Output budget (mirrors MAX_WIDGET_AVATAR_FILE_SIZE_BYTES). */
const OUTPUT_MAX_BYTES = 1024 * 1024;
/** Padding (px, pre-resize) kept around the cropped subject. */
const CROP_PADDING_PX = 6;

export type WidgetAvatarProcessResult =
  | { ok: true; gif: Buffer; transformed: boolean }
  | { ok: false; message: string };

type DecodedAnimation = {
  frames: Buffer[];
  width: number;
  height: number;
  delaysMs: number[];
  loop: number;
};

async function convertMp4ToGifBuffer(bytes: Buffer): Promise<Buffer> {
  if (!ffmpegPath) {
    throw new Error('ffmpeg binary is not available on this server.');
  }
  const dir = await mkdtemp(path.join(tmpdir(), 'ae-avatar-'));
  const inPath = path.join(dir, 'in.mp4');
  const outPath = path.join(dir, 'out.gif');
  try {
    await writeFile(inPath, bytes);
    await new Promise<void>((resolve, reject) => {
      // dither=none keeps flat backgrounds flat so border keying stays reliable.
      const filter =
        'split[s0][s1];[s0]palettegen=max_colors=256[p];[s1][p]paletteuse=dither=none';
      const proc = spawn(
        ffmpegPath as string,
        ['-y', '-i', inPath, '-vf', filter, '-f', 'gif', outPath],
        { stdio: ['ignore', 'ignore', 'pipe'] },
      );
      let stderr = '';
      proc.stderr.on('data', (d: Buffer) => {
        stderr += d.toString();
      });
      proc.on('error', reject);
      proc.on('close', (code) => {
        if (code === 0) resolve();
        else reject(new Error(`ffmpeg exited with code ${code}: ${stderr.slice(-300)}`));
      });
    });
    return await readFile(outPath);
  } finally {
    await rm(dir, { recursive: true, force: true });
  }
}

async function decodeAnimatedGif(gif: Buffer): Promise<DecodedAnimation> {
  const image = sharp(gif, { animated: true });
  const meta = await image.metadata();
  const width = meta.width ?? 0;
  const pages = meta.pages ?? 1;
  const pageHeight = meta.pageHeight ?? meta.height ?? 0;
  if (!width || !pageHeight) {
    throw new Error('Could not read animation dimensions.');
  }
  const raw = await image.ensureAlpha().raw().toBuffer();
  const frameBytes = width * pageHeight * 4;
  const frames: Buffer[] = [];
  for (let i = 0; i < pages; i++) {
    frames.push(raw.subarray(i * frameBytes, (i + 1) * frameBytes));
  }
  const delaysRaw = Array.isArray(meta.delay) ? meta.delay : [];
  const delaysMs = frames.map((_, i) => {
    const d = Number(delaysRaw[i] ?? delaysRaw[0] ?? 100);
    return Number.isFinite(d) && d >= 20 ? d : 100;
  });
  return { frames, width, height: pageHeight, delaysMs, loop: Number(meta.loop ?? 0) };
}

function* borderOffsets(width: number, height: number): Generator<number> {
  for (let x = 0; x < width; x++) {
    yield x;
    yield (height - 1) * width + x;
  }
  for (let y = 1; y < height - 1; y++) {
    yield y * width;
    yield y * width + (width - 1);
  }
}

/**
 * Detects the background key colors from the first frame's border.
 * Returns null when the asset is already transparent or the border is too
 * varied to key safely.
 */
function detectBorderKeyColors(
  frame: Buffer,
  width: number,
  height: number,
): Array<[number, number, number]> | null {
  const buckets = new Map<number, { count: number; r: number; g: number; b: number }>();
  let total = 0;
  let transparent = 0;
  for (const off of borderOffsets(width, height)) {
    const i = off * 4;
    total += 1;
    if (frame[i + 3] < 128) {
      transparent += 1;
      continue;
    }
    const r = frame[i];
    const g = frame[i + 1];
    const b = frame[i + 2];
    const key = ((r >> 4) << 8) | ((g >> 4) << 4) | (b >> 4);
    const bucket = buckets.get(key);
    if (bucket) {
      bucket.count += 1;
      bucket.r += r;
      bucket.g += g;
      bucket.b += b;
    } else {
      buckets.set(key, { count: 1, r, g, b });
    }
  }
  if (total === 0) return null;
  if (transparent / total >= TRANSPARENT_BORDER_FRACTION) return null;

  const sorted = [...buckets.values()].sort((a, b) => b.count - a.count);
  const chosen: Array<[number, number, number]> = [];
  let covered = 0;
  let matched = false;
  for (const bucket of sorted.slice(0, 6)) {
    chosen.push([
      Math.round(bucket.r / bucket.count),
      Math.round(bucket.g / bucket.count),
      Math.round(bucket.b / bucket.count),
    ]);
    covered += bucket.count;
    if (covered / total >= BORDER_COVERAGE_THRESHOLD) {
      matched = true;
      break;
    }
  }
  if (!matched) return null;

  // Checkerboard tiles meet at anti-aliased seams whose colors sit between
  // the two tile tones; key the pairwise midpoints so the seam grid is
  // removed together with the tiles.
  const midpoints: Array<[number, number, number]> = [];
  for (let i = 0; i < chosen.length; i++) {
    for (let j = i + 1; j < chosen.length; j++) {
      midpoints.push([
        Math.round((chosen[i][0] + chosen[j][0]) / 2),
        Math.round((chosen[i][1] + chosen[j][1]) / 2),
        Math.round((chosen[i][2] + chosen[j][2]) / 2),
      ]);
    }
  }
  return [...chosen, ...midpoints];
}

function matchesKeyColor(
  r: number,
  g: number,
  b: number,
  keyColors: Array<[number, number, number]>,
): boolean {
  for (const [kr, kg, kb] of keyColors) {
    if (
      Math.abs(r - kr) <= KEY_COLOR_TOLERANCE &&
      Math.abs(g - kg) <= KEY_COLOR_TOLERANCE &&
      Math.abs(b - kb) <= KEY_COLOR_TOLERANCE
    ) {
      return true;
    }
  }
  return false;
}

/** Flood fill from the borders: clears alpha on background-connected pixels only. */
function removeBackgroundInPlace(
  frame: Buffer,
  width: number,
  height: number,
  keyColors: Array<[number, number, number]>,
): void {
  const seen = new Uint8Array(width * height);
  const queue = new Int32Array(width * height);
  let head = 0;
  let tail = 0;

  for (const off of borderOffsets(width, height)) {
    if (seen[off]) continue;
    const i = off * 4;
    if (frame[i + 3] >= 128 && !matchesKeyColor(frame[i], frame[i + 1], frame[i + 2], keyColors)) {
      continue;
    }
    seen[off] = 1;
    queue[tail++] = off;
  }

  while (head < tail) {
    const off = queue[head++];
    frame[off * 4 + 3] = 0;
    const x = off % width;
    const y = (off / width) | 0;
    const neighbors = [
      x + 1 < width ? off + 1 : -1,
      x - 1 >= 0 ? off - 1 : -1,
      y + 1 < height ? off + width : -1,
      y - 1 >= 0 ? off - width : -1,
    ];
    for (const n of neighbors) {
      if (n < 0 || seen[n]) continue;
      const i = n * 4;
      if (
        frame[i + 3] < 128 ||
        matchesKeyColor(frame[i], frame[i + 1], frame[i + 2], keyColors)
      ) {
        seen[n] = 1;
        queue[tail++] = n;
      }
    }
  }
}

function unionVisibleBoundingBox(
  frames: Buffer[],
  width: number,
  height: number,
): { left: number; top: number; right: number; bottom: number } | null {
  // Histogram of opaque pixels per column/row across all frames, so isolated
  // leftover speckles cannot inflate the crop box.
  const colCounts = new Uint32Array(width);
  const rowCounts = new Uint32Array(height);
  let totalOpaque = 0;
  for (const frame of frames) {
    for (let y = 0; y < height; y++) {
      for (let x = 0; x < width; x++) {
        if (frame[(y * width + x) * 4 + 3] >= 128) {
          colCounts[x] += 1;
          rowCounts[y] += 1;
          totalOpaque += 1;
        }
      }
    }
  }
  if (totalOpaque === 0) return null;

  // Trim each edge until ~0.3% of the opaque mass is consumed: leftover
  // speckles carry almost no mass, so they cannot inflate the crop box.
  const trimBudget = totalOpaque * 0.003;
  let acc = 0;
  let left = 0;
  while (left < width - 1 && acc + colCounts[left] <= trimBudget) acc += colCounts[left++];
  acc = 0;
  let right = width - 1;
  while (right > left && acc + colCounts[right] <= trimBudget) acc += colCounts[right--];
  acc = 0;
  let top = 0;
  while (top < height - 1 && acc + rowCounts[top] <= trimBudget) acc += rowCounts[top++];
  acc = 0;
  let bottom = height - 1;
  while (bottom > top && acc + rowCounts[bottom] <= trimBudget) acc += rowCounts[bottom--];
  if (right < left || bottom < top) return null;

  return {
    left: Math.max(0, left - CROP_PADDING_PX),
    top: Math.max(0, top - CROP_PADDING_PX),
    right: Math.min(width - 1, right + CROP_PADDING_PX),
    bottom: Math.min(height - 1, bottom + CROP_PADDING_PX),
  };
}

async function encodeAnimatedGif(
  frames: Buffer[],
  width: number,
  height: number,
  delaysMs: number[],
  loop: number,
  crop: { left: number; top: number; right: number; bottom: number },
  maxDimension: number,
  frameStep: number,
): Promise<Buffer> {
  const cropWidth = crop.right - crop.left + 1;
  const cropHeight = crop.bottom - crop.top + 1;
  const scale = Math.min(1, maxDimension / Math.max(cropWidth, cropHeight));
  const targetWidth = Math.max(1, Math.round(cropWidth * scale));
  const targetHeight = Math.max(1, Math.round(cropHeight * scale));

  const selected: Buffer[] = [];
  const selectedDelays: number[] = [];
  for (let i = 0; i < frames.length; i += frameStep) {
    selected.push(frames[i]);
    let delay = 0;
    for (let j = i; j < Math.min(i + frameStep, frames.length); j++) {
      delay += delaysMs[j] ?? 100;
    }
    selectedDelays.push(Math.max(20, Math.round(delay)));
  }

  const pngFrames = await Promise.all(
    selected.map((frame) =>
      sharp(frame, { raw: { width, height, channels: 4 } })
        .extract({ left: crop.left, top: crop.top, width: cropWidth, height: cropHeight })
        .resize(targetWidth, targetHeight)
        .png()
        .toBuffer(),
    ),
  );

  return sharp(pngFrames, { join: { animated: true } })
    .gif({ delay: selectedDelays, loop })
    .toBuffer();
}

/**
 * Converts a GIF/MP4 avatar upload into a transparent animated GIF.
 *
 * - `{ ok: true, transformed: false }` → asset was fine as-is (already
 *   transparent GIF); callers should store the original bytes.
 * - `{ ok: true, transformed: true }` → store the returned `gif` buffer.
 * - `{ ok: false }` → processing failed; callers may fall back to the
 *   original bytes.
 */
export async function processWidgetAvatarToTransparentGif(
  bytes: Buffer,
  sourceExt: 'gif' | 'mp4',
): Promise<WidgetAvatarProcessResult> {
  try {
    const gifSource = sourceExt === 'mp4' ? await convertMp4ToGifBuffer(bytes) : bytes;
    const { frames, width, height, delaysMs, loop } = await decodeAnimatedGif(gifSource);
    if (frames.length === 0) {
      return { ok: false, message: 'Animation has no frames.' };
    }

    const keyColors = detectBorderKeyColors(frames[0], width, height);

    if (!keyColors) {
      // Already transparent, or background too varied to key safely.
      if (sourceExt === 'gif') {
        return { ok: true, gif: bytes, transformed: false };
      }
      // MP4 must still become a GIF for the widget; keep pixels untouched.
      const crop = { left: 0, top: 0, right: width - 1, bottom: height - 1 };
      for (const attempt of [
        { maxDimension: MAX_AVATAR_DIMENSION_PX, frameStep: 1 },
        { maxDimension: 140, frameStep: 2 },
        { maxDimension: 110, frameStep: 2 },
      ]) {
        const gif = await encodeAnimatedGif(
          frames, width, height, delaysMs, loop, crop,
          attempt.maxDimension, attempt.frameStep,
        );
        if (gif.length <= OUTPUT_MAX_BYTES) {
          return { ok: true, gif, transformed: true };
        }
      }
      return { ok: false, message: 'Converted avatar exceeds the 1 MB limit.' };
    }

    for (const frame of frames) {
      removeBackgroundInPlace(frame, width, height, keyColors);
    }

    const crop = unionVisibleBoundingBox(frames, width, height);
    if (!crop) {
      return { ok: false, message: 'Avatar became fully transparent after background removal.' };
    }

    for (const attempt of [
      { maxDimension: MAX_AVATAR_DIMENSION_PX, frameStep: 1 },
      { maxDimension: 140, frameStep: 2 },
      { maxDimension: 110, frameStep: 2 },
    ]) {
      const gif = await encodeAnimatedGif(
        frames, width, height, delaysMs, loop, crop,
        attempt.maxDimension, attempt.frameStep,
      );
      if (gif.length <= OUTPUT_MAX_BYTES) {
        return { ok: true, gif, transformed: true };
      }
    }
    return { ok: false, message: 'Processed avatar exceeds the 1 MB limit.' };
  } catch (e) {
    const message = e instanceof Error ? e.message : 'Avatar processing failed.';
    return { ok: false, message };
  }
}