import { extname } from "node:path";
import { HarnessError } from "../errors/harness-error.js";

export type SupportedImageMediaType = "image/jpeg" | "image/png" | "image/webp";

const extensions: Record<SupportedImageMediaType, ReadonlySet<string>> = {
  "image/jpeg": new Set([".jpg", ".jpeg"]),
  "image/png": new Set([".png"]),
  "image/webp": new Set([".webp"]),
};

export function detectImageMediaType(bytes: Uint8Array): SupportedImageMediaType | undefined {
  if (
    bytes.length >= 8 &&
    bytes[0] === 0x89 && bytes[1] === 0x50 && bytes[2] === 0x4e && bytes[3] === 0x47 &&
    bytes[4] === 0x0d && bytes[5] === 0x0a && bytes[6] === 0x1a && bytes[7] === 0x0a
  ) return "image/png";
  if (bytes.length >= 3 && bytes[0] === 0xff && bytes[1] === 0xd8 && bytes[2] === 0xff) {
    return "image/jpeg";
  }
  if (
    bytes.length >= 12 &&
    Buffer.from(bytes.subarray(0, 4)).toString("ascii") === "RIFF" &&
    Buffer.from(bytes.subarray(8, 12)).toString("ascii") === "WEBP"
  ) return "image/webp";
  return undefined;
}

export function assertImageType(
  name: string,
  declaredType: string,
  bytes: Uint8Array,
): SupportedImageMediaType {
  if (!(declaredType in extensions)) {
    throw new HarnessError("IMAGE_TYPE_NOT_ALLOWED", "Only JPEG, PNG, and WebP images are allowed");
  }
  const mediaType = declaredType as SupportedImageMediaType;
  const extension = extname(name).toLowerCase();
  if (!extensions[mediaType].has(extension)) {
    throw new HarnessError("IMAGE_TYPE_NOT_ALLOWED", "Filename extension does not match an allowed image type");
  }
  const detected = detectImageMediaType(bytes);
  if (detected === undefined || detected !== mediaType) {
    throw new HarnessError("IMAGE_SIGNATURE_MISMATCH", "Image bytes do not match the declared type");
  }
  return mediaType;
}

export function extensionForMediaType(mediaType: SupportedImageMediaType): string {
  if (mediaType === "image/jpeg") return ".jpg";
  if (mediaType === "image/png") return ".png";
  return ".webp";
}