/**
 * JPEG conversion for images, backed by macOS `sips`.
 *
 * Two consumers share this converter:
 *   - transport optimization (`agent/image-optimize.ts`) downscales large
 *     images before provider calls;
 *   - storage normalization (attachment ingress + history hydration) converts
 *     HEIF/HEIC — which Chromium-based clients cannot decode — to JPEG.
 *
 * Conversion runs `sips` (a macOS builtin); on other platforms or on any
 * failure it returns null and callers keep the original bytes. Results are
 * cached on disk keyed by content hash + conversion options, so repeated
 * conversions of the same image (or daemon restarts) skip the sips call.
 */

import { createHash } from "node:crypto";
import {
  closeSync,
  existsSync,
  mkdirSync,
  openSync,
  readdirSync,
  readFileSync,
  readSync,
  renameSync,
  statSync,
  unlinkSync,
  writeFileSync,
} from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";

import { v4 as uuid } from "uuid";

const DEFAULT_JPEG_QUALITY = 80;

/** Full-resolution quality for stored attachment masters. */
const STORAGE_JPEG_QUALITY = 90;

const CACHE_MAX_ENTRIES = 500;

function getCacheDir(): string {
  return join(tmpdir(), "vellum-optimized-images");
}

/**
 * Structural completeness check for a JPEG payload: SOI marker (FF D8) at the
 * head and EOI marker (FF D9) at the tail. Every cache entry and every sips
 * output is a JPEG, so a payload failing this is truncated or empty (the
 * signature of a torn cache write) and must never reach a provider, which
 * rejects it with a 400 that wedges the conversation (the corrupt block is
 * resent on every subsequent turn).
 */
export function isCompleteJpeg(bytes: Uint8Array): boolean {
  return (
    bytes.length >= 4 &&
    bytes[0] === 0xff &&
    bytes[1] === 0xd8 &&
    bytes[bytes.length - 2] === 0xff &&
    bytes[bytes.length - 1] === 0xd9
  );
}

/**
 * Structural JPEG truncation check for payloads of arbitrary origin: walk the
 * marker segments from SOI and report whether a terminal EOI is reached.
 * Unlike {@link isCompleteJpeg} (exact tail framing, for sips output and
 * cache entries, which always end on EOI), this tolerates encoders that
 * append padding or metadata after the EOI.
 *
 * A raw byte search for FF D9 is not enough: length-delimited APP/COM
 * segments (an EXIF thumbnail is itself a complete embedded JPEG) may contain
 * FF D9, so a truncated file with an intact metadata prefix would pass.
 * Walking segment boundaries skips those payloads entirely; only an EOI at
 * the top level of the marker stream counts. Entropy-coded scan data is
 * traversed byte-wise, where FF is stuffed as FF 00 and restart markers
 * (D0-D7) continue the scan, so a marker byte there is unambiguous.
 *
 * The EOI only counts after at least one frame header (SOF) and one scan
 * (SOS) have been seen: a degenerate payload such as bare SOI+EOI carries no
 * image data and providers reject it.
 */
export function hasValidJpegStructure(bytes: Uint8Array): boolean {
  if (bytes.length < 4 || bytes[0] !== 0xff || bytes[1] !== 0xd8) {
    return false;
  }
  let sawFrame = false;
  let sawScan = false;
  let i = 2;
  while (i + 1 < bytes.length) {
    if (bytes[i] !== 0xff) {
      return false;
    }
    // FF fill bytes before a marker are legal padding.
    let j = i + 1;
    while (j < bytes.length && bytes[j] === 0xff) {
      j++;
    }
    if (j >= bytes.length) {
      return false;
    }
    const marker = bytes[j];
    i = j + 1;
    if (marker === 0xd9) {
      return sawFrame && sawScan;
    }
    // SOF0-SOF15 occupy C0-CF, excluding DHT (C4), JPG (C8), and DAC (CC).
    if (
      marker >= 0xc0 &&
      marker <= 0xcf &&
      marker !== 0xc4 &&
      marker !== 0xc8 &&
      marker !== 0xcc
    ) {
      sawFrame = true;
    }
    // Standalone markers carry no length field: repeated SOI, TEM, RST0-7.
    if (
      marker === 0xd8 ||
      marker === 0x01 ||
      (marker >= 0xd0 && marker <= 0xd7)
    ) {
      continue;
    }
    if (i + 1 >= bytes.length) {
      return false;
    }
    const segmentLength = (bytes[i] << 8) | bytes[i + 1];
    if (segmentLength < 2) {
      return false;
    }
    i += segmentLength;
    if (marker === 0xda) {
      sawScan = true;
      // SOS: entropy-coded data follows the header. Scan to the next real
      // marker; FF 00 (stuffed data byte) and FF D0-D7 (restart) stay inside
      // the scan.
      while (i + 1 < bytes.length) {
        if (
          bytes[i] === 0xff &&
          bytes[i + 1] !== 0x00 &&
          !(bytes[i + 1] >= 0xd0 && bytes[i + 1] <= 0xd7)
        ) {
          break;
        }
        i++;
      }
    }
  }
  return false;
}

function readFromCache(key: string): Buffer | null {
  const cachePath = join(getCacheDir(), `${key}.jpg`);
  try {
    if (!existsSync(cachePath)) {
      return null;
    }
    const bytes = readFileSync(cachePath) as Buffer;
    if (!isCompleteJpeg(bytes)) {
      // Poisoned entry (torn write, disk full, etc.): drop it so the caller
      // re-converts and re-caches.
      unlinkSync(cachePath);
      return null;
    }
    return bytes;
  } catch {
    return null;
  }
}

function writeToCache(key: string, convertedBytes: Buffer): void {
  const dir = getCacheDir();
  // Write-then-rename so concurrent readers (the cache dir is shared across
  // daemon processes) never observe a partially written entry; rename within
  // the same directory is atomic. The pid+uuid suffix keeps concurrent
  // writers of the same key off each other's temp file.
  const tmpPath = join(dir, `${key}.${process.pid}-${uuid().slice(0, 8)}.tmp`);
  try {
    mkdirSync(dir, { recursive: true });
    writeFileSync(tmpPath, convertedBytes);
    renameSync(tmpPath, join(dir, `${key}.jpg`));
    evictIfNeeded(dir);
  } catch {
    // Cache write failure is non-fatal.
    try {
      unlinkSync(tmpPath);
    } catch {
      /* ignore */
    }
  }
}

function evictIfNeeded(dir: string): void {
  try {
    const entries = readdirSync(dir)
      .filter((f) => f.endsWith(".jpg"))
      .map((f) => {
        const full = join(dir, f);
        return { path: full, mtimeMs: statSync(full).mtimeMs };
      })
      .sort((a, b) => a.mtimeMs - b.mtimeMs);
    const excess = entries.length - CACHE_MAX_ENTRIES;
    for (let i = 0; i < excess; i++) {
      try {
        unlinkSync(entries[i]!.path);
      } catch {
        /* ignore */
      }
    }
  } catch {
    /* ignore */
  }
}

export interface ConvertToJpegOptions {
  /** Downscale so neither side exceeds this; omit to keep full resolution. */
  maxDimensionPx?: number;
  /**
   * Resample to exactly these pixel dimensions (up- or downscaling). The
   * caller is responsible for preserving aspect ratio. Mutually exclusive
   * with `maxDimensionPx`; when both are set this wins.
   */
  resizeToPx?: { width: number; height: number };
  /** JPEG quality 1-100 (default 80). */
  quality?: number;
}

async function runSips(
  inputBytes: Uint8Array,
  options: ConvertToJpegOptions,
): Promise<Buffer | null> {
  const stamp = `${Date.now()}-${uuid().slice(0, 8)}`;
  const srcPath = join(tmpdir(), `vellum-img-opt-${stamp}-src`);
  const outPath = join(tmpdir(), `vellum-img-opt-${stamp}-out.jpg`);
  try {
    writeFileSync(srcPath, inputBytes);
    // `-z` resamples to an exact height/width (the only sips mode that can
    // upscale); `--resampleHeightWidthMax` only caps the longest side.
    const args =
      options.resizeToPx != null
        ? [
            "-z",
            String(options.resizeToPx.height),
            String(options.resizeToPx.width),
          ]
        : options.maxDimensionPx != null
          ? ["--resampleHeightWidthMax", String(options.maxDimensionPx)]
          : [];
    args.push(
      "-s",
      "format",
      "jpeg",
      "-s",
      "formatOptions",
      String(options.quality ?? DEFAULT_JPEG_QUALITY),
      srcPath,
      "--out",
      outPath,
    );
    const proc = Bun.spawn(["sips", ...args], {
      stdout: "ignore",
      stderr: "ignore",
      timeout: 15_000,
    });
    await proc.exited;
    if (proc.exitCode !== 0) {
      return null;
    }
    const out = readFileSync(outPath) as Buffer;
    // A zero-exit sips can still leave a truncated file (disk full, timeout
    // kill racing the write). Returning it would poison the cache and every
    // downstream provider call. Treat it as a failed conversion instead.
    if (!isCompleteJpeg(out)) {
      return null;
    }
    return out;
  } catch {
    return null;
  } finally {
    try {
      unlinkSync(srcPath);
    } catch {
      /* ignore */
    }
    try {
      unlinkSync(outPath);
    } catch {
      /* ignore */
    }
  }
}

/**
 * Convert an image to JPEG, optionally downscaling. Returns null when
 * conversion is unavailable (non-macOS) or fails; callers keep the original.
 */
export async function convertImageToJpeg(
  bytes: Uint8Array,
  options: ConvertToJpegOptions = {},
): Promise<Buffer | null> {
  const hash = createHash("sha256").update(bytes).digest("hex").slice(0, 16);
  // Options qualify the key so full-resolution storage conversions and
  // resized transport conversions of the same source never collide.
  const sizeKey =
    options.resizeToPx != null
      ? `${options.resizeToPx.width}x${options.resizeToPx.height}`
      : (options.maxDimensionPx ?? "full");
  const cacheKey = `${hash}-${sizeKey}-q${options.quality ?? DEFAULT_JPEG_QUALITY}`;

  const cached = readFromCache(cacheKey);
  if (cached) {
    return cached;
  }

  const converted = await runSips(bytes, options);
  if (!converted) {
    return null;
  }

  writeToCache(cacheKey, converted);
  return converted;
}

// HEIF container brands (ISO BMFF `ftyp` major brand). AVIF brands are
// deliberately absent: Chromium decodes AVIF natively, so it needs no
// normalization.
const HEIF_FTYP_BRANDS = new Set([
  "heic",
  "heix",
  "hevc",
  "hevx",
  "heif",
  "heim",
  "heis",
  "hevm",
  "hevs",
  "mif1",
  "msf1",
]);

// Brands whose canonical MIME is `image/heic`; the remaining HEIF brands
// (`heif`, `mif1`, `msf1`) are `image/heif`.
const HEIC_FTYP_BRANDS = new Set([
  "heic",
  "heix",
  "hevc",
  "hevx",
  "heim",
  "heis",
  "hevm",
  "hevs",
]);

function heifFtypBrand(bytes: Uint8Array): string | null {
  if (bytes.length < 12) {
    return null;
  }
  // ISO BMFF layout: bytes 4-8 are "ftyp", bytes 8-12 the major brand.
  if (
    bytes[4] !== 0x66 || // f
    bytes[5] !== 0x74 || // t
    bytes[6] !== 0x79 || // y
    bytes[7] !== 0x70 // p
  ) {
    return null;
  }
  const brand = String.fromCharCode(
    bytes[8]!,
    bytes[9]!,
    bytes[10]!,
    bytes[11]!,
  );
  return HEIF_FTYP_BRANDS.has(brand) ? brand : null;
}

/**
 * Content-based HEIF/HEIC detection. MIME metadata is unreliable here:
 * Chromium reports an empty `file.type` for `.heic`, which clients coerce to
 * `application/octet-stream`.
 */
export function isHeifImage(bytes: Uint8Array): boolean {
  return heifFtypBrand(bytes) !== null;
}

/**
 * Canonical HEIF MIME type for bytes whose container brand names one, and null
 * for everything else. For providers that read HEIF directly and key on the
 * declared media type (Gemini accepts `image/heic` and `image/heif`:
 * https://ai.google.dev/gemini-api/docs/image-understanding).
 */
export function heifImageMimeType(bytes: Uint8Array): string | null {
  const brand = heifFtypBrand(bytes);
  if (!brand) {
    return null;
  }
  return HEIC_FTYP_BRANDS.has(brand) ? "image/heic" : "image/heif";
}

/**
 * Sniff the actual image format from magic bytes, covering the formats
 * providers accept as image blocks. Returns null for anything unrecognized
 * (HEIF has its own detector, {@link isHeifImage}).
 *
 * Exists because clients derive the declared MIME from the file extension, so
 * a JPEG renamed to `.png` arrives as `image/png` — and providers reject an
 * image whose bytes disagree with the declared media type.
 */
export function sniffImageMimeType(bytes: Uint8Array): string | null {
  if (
    bytes.length >= 4 &&
    bytes[0] === 0x89 &&
    bytes[1] === 0x50 && // P
    bytes[2] === 0x4e && // N
    bytes[3] === 0x47 // G
  ) {
    return "image/png";
  }
  if (
    bytes.length >= 3 &&
    bytes[0] === 0xff &&
    bytes[1] === 0xd8 &&
    bytes[2] === 0xff
  ) {
    return "image/jpeg";
  }
  if (
    bytes.length >= 4 &&
    bytes[0] === 0x47 && // G
    bytes[1] === 0x49 && // I
    bytes[2] === 0x46 && // F
    bytes[3] === 0x38 // 8
  ) {
    return "image/gif";
  }
  if (
    bytes.length >= 12 &&
    bytes[0] === 0x52 && // R
    bytes[1] === 0x49 && // I
    bytes[2] === 0x46 && // F
    bytes[3] === 0x46 && // F
    bytes[8] === 0x57 && // W
    bytes[9] === 0x45 && // E
    bytes[10] === 0x42 && // B
    bytes[11] === 0x50 // P
  ) {
    return "image/webp";
  }
  return null;
}

/** Base64 variant of {@link sniffImageMimeType}; decodes only the head. */
export function sniffBase64ImageMimeType(dataBase64: string): string | null {
  // 16 base64 chars decode to the 12 bytes the longest signature needs.
  return sniffImageMimeType(Buffer.from(dataBase64.slice(0, 16), "base64"));
}

/**
 * On-disk variant of {@link sniffImageMimeType}, for file-backed attachments
 * that are never read into memory (recordings can exceed the upload limit).
 * Reads only the 12-byte head. Returns null when the file is unreadable or the
 * format is unrecognized, so callers keep the declared MIME.
 */
export function sniffImageFileMimeType(filePath: string): string | null {
  let fd: number;
  try {
    fd = openSync(filePath, "r");
  } catch {
    return null;
  }
  try {
    const head = Buffer.alloc(12);
    const bytesRead = readSync(fd, head, 0, 12, 0);
    return sniffImageMimeType(head.subarray(0, bytesRead));
  } catch {
    return null;
  } finally {
    closeSync(fd);
  }
}

const HEIF_FILENAME_RE = /\.(heic|heif)$/i;

/**
 * Filename-extension HEIF/HEIC detection, for call sites that only hold
 * attachment metadata (filename + MIME) and want to skip hydrating bytes for
 * non-candidate rows. Complements {@link isHeifImage}: Chromium reports an
 * empty MIME for `.heic`, so legacy rows can carry HEIC bytes under
 * `application/octet-stream` while the extension survives.
 */
export function isHeicFilename(filename: string): boolean {
  return HEIF_FILENAME_RE.test(filename.trim());
}

/** Rewrites a filename's extension to `.jpg` (e.g. `IMG_5487.HEIC` → `IMG_5487.jpg`). */
export function jpegFilenameFor(filename: string): string {
  const fallback = "attachment";
  const trimmed = filename.trim() || fallback;
  if (/\.jpe?g$/i.test(trimmed)) {
    return trimmed;
  }
  const withoutExtension = trimmed.replace(/\.[^./\\]+$/, "") || fallback;
  return `${withoutExtension}.jpg`;
}

export interface NormalizedImageBytes {
  mimeType: string;
  bytes: Uint8Array;
  converted: boolean;
}

/**
 * Normalize image bytes for storage: HEIF/HEIC becomes a full-resolution JPEG
 * master; a declared MIME that disagrees with the sniffed format is corrected
 * (bytes untouched); everything else (and any conversion failure) passes
 * through unchanged. Callers that persist a filename should rewrite it with
 * {@link jpegFilenameFor} when `converted` is true — a MIME-only correction
 * sets `converted: false` and keeps the filename.
 */
export async function normalizeImageBytes(
  mimeType: string,
  bytes: Uint8Array,
): Promise<NormalizedImageBytes> {
  if (isHeifImage(bytes)) {
    const converted = await convertImageToJpeg(bytes, {
      quality: STORAGE_JPEG_QUALITY,
    });
    if (!converted) {
      return { mimeType, bytes, converted: false };
    }
    return { mimeType: "image/jpeg", bytes: converted, converted: true };
  }
  const sniffed = sniffImageMimeType(bytes);
  if (sniffed && sniffed !== mimeType) {
    return { mimeType: sniffed, bytes, converted: false };
  }
  return { mimeType, bytes, converted: false };
}

export interface NormalizedImageBase64 {
  mimeType: string;
  dataBase64: string;
  converted: boolean;
}

/**
 * Base64 variant of {@link normalizeImageBytes}. Sniffs the decoded head
 * first so non-HEIF payloads skip the full decode.
 */
export async function normalizeImageBase64(
  mimeType: string,
  dataBase64: string,
): Promise<NormalizedImageBase64> {
  // 16 base64 chars decode to the 12 bytes the ftyp/signature sniffs need.
  const head = Buffer.from(dataBase64.slice(0, 16), "base64");
  if (isHeifImage(head)) {
    const normalized = await normalizeImageBytes(
      mimeType,
      Buffer.from(dataBase64, "base64"),
    );
    if (!normalized.converted) {
      return { mimeType, dataBase64, converted: false };
    }
    return {
      mimeType: normalized.mimeType,
      dataBase64: Buffer.from(normalized.bytes).toString("base64"),
      converted: true,
    };
  }
  const sniffed = sniffImageMimeType(head);
  if (sniffed && sniffed !== mimeType) {
    return { mimeType: sniffed, dataBase64, converted: false };
  }
  return { mimeType, dataBase64, converted: false };
}