// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at https://mozilla.org/MPL/2.0/.

export interface BrowserImageSafetyLimits {
  /** Maximum compressed resource size eligible for automatic rendering. */
  maxBytes: number;
  /** Maximum supported width or height in pixels. */
  maxDimension: number;
  /** Maximum supported decoded pixel count. */
  maxPixels: number;
}

export const DEFAULT_BROWSER_IMAGE_SAFETY_LIMITS: Readonly<BrowserImageSafetyLimits> =
  {
    maxBytes: 32 * 1024 * 1024,
    maxDimension: 16_384,
    maxPixels: 64 * 1024 * 1024,
  };

export type BrowserImageRejectionReason =
  | 'unsupported-format'
  | 'byte-limit'
  | 'invalid-header'
  | 'dimension-limit'
  | 'pixel-limit';

export interface BrowserImageInspection {
  format?: 'png' | 'jpeg';
  extension?: 'png' | 'jpg';
  mediaType?: 'image/png' | 'image/jpeg';
  width?: number;
  height?: number;
  pixels?: number;
  browserRenderable: boolean;
  rejectionReason?: BrowserImageRejectionReason;
}

interface ImageDimensions {
  width: number;
  height: number;
}

const PNG_SIGNATURE = Uint8Array.of(
  0x89,
  0x50,
  0x4e,
  0x47,
  0x0d,
  0x0a,
  0x1a,
  0x0a
);
const PNG_IHDR_SIZE = 33;
const JPEG_HEADER_SCAN_BYTES = 1024 * 1024;
const JPEG_MARKER_LIMIT = 4_096;

/**
 * Inspect untrusted PNG/JPEG bytes without decoding pixels. The JPEG walk is
 * capped independently of the resource size so malformed marker streams
 * cannot turn preview eligibility checks into an unbounded scan.
 */
export function inspectBrowserImage(
  bytes: Uint8Array,
  limits: BrowserImageSafetyLimits = DEFAULT_BROWSER_IMAGE_SAFETY_LIMITS
): BrowserImageInspection {
  validateLimits(limits);

  let identity:
    | {
        format: 'png' | 'jpeg';
        extension: 'png' | 'jpg';
        mediaType: 'image/png' | 'image/jpeg';
      }
    | undefined;
  let dimensions: ImageDimensions | undefined;

  if (hasPngSignature(bytes)) {
    identity = {
      format: 'png',
      extension: 'png',
      mediaType: 'image/png',
    };
    dimensions = readPngDimensions(bytes);
  } else if (hasJpegSignature(bytes)) {
    identity = {
      format: 'jpeg',
      extension: 'jpg',
      mediaType: 'image/jpeg',
    };
    dimensions = readJpegDimensions(bytes);
  } else {
    return {
      browserRenderable: false,
      rejectionReason: 'unsupported-format',
    };
  }

  if (!dimensions) {
    return {
      ...identity,
      browserRenderable: false,
      rejectionReason: 'invalid-header',
    };
  }

  const { width, height } = dimensions;
  const pixels =
    width <= Math.floor(Number.MAX_SAFE_INTEGER / height)
      ? width * height
      : undefined;
  const details = { ...identity, width, height, pixels };

  if (bytes.byteLength > limits.maxBytes) {
    return {
      ...details,
      browserRenderable: false,
      rejectionReason: 'byte-limit',
    };
  }
  if (width > limits.maxDimension || height > limits.maxDimension) {
    return {
      ...details,
      browserRenderable: false,
      rejectionReason: 'dimension-limit',
    };
  }
  if (pixels === undefined || pixels > limits.maxPixels) {
    return {
      ...details,
      browserRenderable: false,
      rejectionReason: 'pixel-limit',
    };
  }
  return { ...details, browserRenderable: true };
}

function readPngDimensions(bytes: Uint8Array): ImageDimensions | undefined {
  if (bytes.byteLength < PNG_IHDR_SIZE) return undefined;
  if (
    readU32Be(bytes, 8) !== 13 ||
    bytes[12] !== 0x49 ||
    bytes[13] !== 0x48 ||
    bytes[14] !== 0x44 ||
    bytes[15] !== 0x52
  ) {
    return undefined;
  }
  const width = readU32Be(bytes, 16);
  const height = readU32Be(bytes, 20);
  const bitDepth = bytes[24]!;
  const colorType = bytes[25]!;
  if (
    width === 0 ||
    height === 0 ||
    !validPngBitDepth(bitDepth, colorType) ||
    bytes[26] !== 0 ||
    bytes[27] !== 0 ||
    (bytes[28] !== 0 && bytes[28] !== 1)
  ) {
    return undefined;
  }
  return { width, height };
}

function validPngBitDepth(bitDepth: number, colorType: number): boolean {
  switch (colorType) {
    case 0:
      return [1, 2, 4, 8, 16].includes(bitDepth);
    case 2:
    case 4:
    case 6:
      return bitDepth === 8 || bitDepth === 16;
    case 3:
      return [1, 2, 4, 8].includes(bitDepth);
    default:
      return false;
  }
}

function readJpegDimensions(bytes: Uint8Array): ImageDimensions | undefined {
  const scanEnd = Math.min(bytes.byteLength, JPEG_HEADER_SCAN_BYTES);
  let offset = 2;
  let markers = 0;

  while (offset < scanEnd && markers < JPEG_MARKER_LIMIT) {
    if (bytes[offset] !== 0xff) return undefined;
    while (offset < scanEnd && bytes[offset] === 0xff) offset += 1;
    if (offset >= scanEnd) return undefined;
    const marker = bytes[offset++]!;
    markers += 1;

    if (marker === 0x00) return undefined;
    if (marker === 0xd9 || marker === 0xda) return undefined;
    if (isStandaloneJpegMarker(marker)) continue;
    if (offset + 2 > scanEnd) return undefined;

    const segmentLength = readU16Be(bytes, offset);
    if (segmentLength < 2 || offset + segmentLength > scanEnd) {
      return undefined;
    }
    if (isStartOfFrame(marker)) {
      if (segmentLength < 8) return undefined;
      const height = readU16Be(bytes, offset + 3);
      const width = readU16Be(bytes, offset + 5);
      const componentCount = bytes[offset + 7]!;
      if (
        width === 0 ||
        height === 0 ||
        componentCount === 0 ||
        segmentLength < 8 + componentCount * 3
      ) {
        return undefined;
      }
      return { width, height };
    }
    offset += segmentLength;
  }
  return undefined;
}

function isStandaloneJpegMarker(marker: number): boolean {
  return (
    marker === 0x01 || marker === 0xd8 || (marker >= 0xd0 && marker <= 0xd7)
  );
}

function isStartOfFrame(marker: number): boolean {
  return (
    marker >= 0xc0 &&
    marker <= 0xcf &&
    marker !== 0xc4 &&
    marker !== 0xc8 &&
    marker !== 0xcc
  );
}

function hasPngSignature(bytes: Uint8Array): boolean {
  return (
    bytes.byteLength >= PNG_SIGNATURE.byteLength &&
    PNG_SIGNATURE.every((value, index) => bytes[index] === value)
  );
}

function hasJpegSignature(bytes: Uint8Array): boolean {
  return (
    bytes.byteLength >= 3 &&
    bytes[0] === 0xff &&
    bytes[1] === 0xd8 &&
    bytes[2] === 0xff
  );
}

function readU16Be(bytes: Uint8Array, offset: number): number {
  return bytes[offset]! * 0x100 + bytes[offset + 1]!;
}

function readU32Be(bytes: Uint8Array, offset: number): number {
  return (
    bytes[offset]! * 0x100_0000 +
    bytes[offset + 1]! * 0x1_0000 +
    bytes[offset + 2]! * 0x100 +
    bytes[offset + 3]!
  );
}

function validateLimits(limits: BrowserImageSafetyLimits): void {
  for (const [name, value] of Object.entries(limits)) {
    if (!Number.isSafeInteger(value) || value <= 0) {
      throw new TypeError(`${name} must be a positive safe integer`);
    }
  }
}