import { open } from "node:fs/promises";
import type { FileHandle } from "@lmstudio/sdk";

export const IMAGE_LIMIT = 20 * 1024 * 1024;
export const TOTAL_IMAGE_LIMIT = 64 * 1024 * 1024;

// Inspect bytes rather than extensions: LM Studio's stored files may have no extension.
export function imageMime(bytes: Buffer): string {
  if (bytes.subarray(0, 8).equals(Buffer.from([137, 80, 78, 71, 13, 10, 26, 10]))) return "image/png";
  if (bytes[0] === 255 && bytes[1] === 216 && bytes[2] === 255) return "image/jpeg";
  if (/^GIF8[79]a$/.test(bytes.subarray(0, 6).toString("ascii"))) return "image/gif";
  if (bytes.subarray(0, 4).toString("ascii") === "RIFF" && bytes.subarray(8, 12).toString("ascii") === "WEBP") return "image/webp";
  if (bytes.subarray(0, 2).toString("ascii") === "BM") return "image/bmp";
  if (["49492a00", "4d4d002a"].includes(bytes.subarray(0, 4).toString("hex"))) return "image/tiff";
  if (bytes.subarray(4, 8).toString("ascii") === "ftyp" && /avif|avis/.test(bytes.subarray(8, 32).toString("ascii"))) return "image/avif";
  throw new Error("Unsupported image data. Use PNG, JPEG, WebP, GIF, BMP, TIFF or AVIF.");
}

export async function readImage(path: string, signal?: AbortSignal): Promise<Buffer> {
  signal?.throwIfAborted();
  const handle = await open(path, "r");
  try {
    const info = await handle.stat();
    if (!info.isFile() || info.size > IMAGE_LIMIT) throw new Error("Image must be a file of at most 20 MiB.");
    // A bounded read also handles a file growing after stat().
    const bytes = Buffer.alloc(info.size + 1);
    let size = 0;
    while (size < bytes.length) {
      signal?.throwIfAborted();
      const { bytesRead } = await handle.read(bytes, size, bytes.length - size, null);
      if (!bytesRead) break;
      size += bytesRead;
    }
    if (size > info.size) throw new Error("Image changed while reading; retry the request.");
    const result = bytes.subarray(0, size);
    imageMime(result);
    return result;
  } finally { await handle.close(); }
}

export async function readAttachment(file: FileHandle, signal: AbortSignal): Promise<Buffer> {
  signal.throwIfAborted();
  // getFilePath() is not implemented for SDK base64: handles.
  if (!file.identifier.startsWith("base64:")) return readImage(await file.getFilePath(), signal);
  const encoded = file.identifier.slice(7);
  if (encoded.length > Math.ceil(IMAGE_LIMIT / 3) * 4) throw new Error("Image exceeds 20 MiB.");
  const bytes = Buffer.from(encoded, "base64");
  if (bytes.length > IMAGE_LIMIT) throw new Error("Image exceeds 20 MiB.");
  imageMime(bytes);
  return bytes;
}

import { open } from "node:fs/promises";
import type { FileHandle } from "@lmstudio/sdk";

export const IMAGE_LIMIT = 20 * 1024 * 1024;
export const TOTAL_IMAGE_LIMIT = 64 * 1024 * 1024;

// Inspect bytes rather than extensions: LM Studio's stored files may have no extension.
export function imageMime(bytes: Buffer): string {
  if (bytes.subarray(0, 8).equals(Buffer.from([137, 80, 78, 71, 13, 10, 26, 10]))) return "image/png";
  if (bytes[0] === 255 && bytes[1] === 216 && bytes[2] === 255) return "image/jpeg";
  if (/^GIF8[79]a$/.test(bytes.subarray(0, 6).toString("ascii"))) return "image/gif";
  if (bytes.subarray(0, 4).toString("ascii") === "RIFF" && bytes.subarray(8, 12).toString("ascii") === "WEBP") return "image/webp";
  if (bytes.subarray(0, 2).toString("ascii") === "BM") return "image/bmp";
  if (["49492a00", "4d4d002a"].includes(bytes.subarray(0, 4).toString("hex"))) return "image/tiff";
  if (bytes.subarray(4, 8).toString("ascii") === "ftyp" && /avif|avis/.test(bytes.subarray(8, 32).toString("ascii"))) return "image/avif";
  throw new Error("Unsupported image data. Use PNG, JPEG, WebP, GIF, BMP, TIFF or AVIF.");
}

export async function readImage(path: string, signal?: AbortSignal): Promise<Buffer> {
  signal?.throwIfAborted();
  const handle = await open(path, "r");
  try {
    const info = await handle.stat();
    if (!info.isFile() || info.size > IMAGE_LIMIT) throw new Error("Image must be a file of at most 20 MiB.");
    // A bounded read also handles a file growing after stat().
    const bytes = Buffer.alloc(info.size + 1);
    let size = 0;
    while (size < bytes.length) {
      signal?.throwIfAborted();
      const { bytesRead } = await handle.read(bytes, size, bytes.length - size, null);
      if (!bytesRead) break;
      size += bytesRead;
    }
    if (size > info.size) throw new Error("Image changed while reading; retry the request.");
    const result = bytes.subarray(0, size);
    imageMime(result);
    return result;
  } finally { await handle.close(); }
}

export async function readAttachment(file: FileHandle, signal: AbortSignal): Promise<Buffer> {
  signal.throwIfAborted();
  // getFilePath() is not implemented for SDK base64: handles.
  if (!file.identifier.startsWith("base64:")) return readImage(await file.getFilePath(), signal);
  const encoded = file.identifier.slice(7);
  if (encoded.length > Math.ceil(IMAGE_LIMIT / 3) * 4) throw new Error("Image exceeds 20 MiB.");
  const bytes = Buffer.from(encoded, "base64");
  if (bytes.length > IMAGE_LIMIT) throw new Error("Image exceeds 20 MiB.");
  imageMime(bytes);
  return bytes;
}