export interface ImageDimensions {
  width: number
  height: number
}

function readUint32BE(bytes: Uint8Array, offset: number): number {
  return (
    ((bytes[offset] ?? 0) << 24) |
    ((bytes[offset + 1] ?? 0) << 16) |
    ((bytes[offset + 2] ?? 0) << 8) |
    (bytes[offset + 3] ?? 0)
  ) >>> 0
}

function readUint16BE(bytes: Uint8Array, offset: number): number {
  return (((bytes[offset] ?? 0) << 8) | (bytes[offset + 1] ?? 0)) >>> 0
}

function readUint24LE(bytes: Uint8Array, offset: number): number {
  return ((bytes[offset] ?? 0) | ((bytes[offset + 1] ?? 0) << 8) | ((bytes[offset + 2] ?? 0) << 16)) >>> 0
}

function parsePngDimensions(header: Uint8Array): ImageDimensions | null {
  if (header.length < 24) return null
  const pngSig = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]
  for (let i = 0; i < pngSig.length; i++) {
    if (header[i] !== pngSig[i]) return null
  }
  const chunkType = String.fromCharCode(header[12] ?? 0, header[13] ?? 0, header[14] ?? 0, header[15] ?? 0)
  if (chunkType !== 'IHDR') return null
  return {
    width: readUint32BE(header, 16),
    height: readUint32BE(header, 20),
  }
}

function parseJpegDimensions(header: Uint8Array): ImageDimensions | null {
  let offset = 2
  while (offset + 9 < header.length) {
    if (header[offset] !== 0xff) {
      offset++
      continue
    }
    const marker = header[offset + 1] ?? 0
    if (marker === 0xd8 || marker === 0xd9) {
      offset += 2
      continue
    }
    const isSof =
      (marker >= 0xc0 && marker <= 0xc3) ||
      (marker >= 0xc5 && marker <= 0xc7) ||
      (marker >= 0xc9 && marker <= 0xcb) ||
      (marker >= 0xcd && marker <= 0xcf)
    if (isSof) {
      return {
        height: readUint16BE(header, offset + 5),
        width: readUint16BE(header, offset + 7),
      }
    }
    const segmentLength = readUint16BE(header, offset + 2)
    if (segmentLength < 2) return null
    offset += 2 + segmentLength
  }
  return null
}

function parseWebpDimensions(header: Uint8Array): ImageDimensions | null {
  if (header.length < 30) return null
  const riff = String.fromCharCode(header[0] ?? 0, header[1] ?? 0, header[2] ?? 0, header[3] ?? 0)
  const webp = String.fromCharCode(header[8] ?? 0, header[9] ?? 0, header[10] ?? 0, header[11] ?? 0)
  if (riff !== 'RIFF' || webp !== 'WEBP') return null

  const chunk = String.fromCharCode(header[12] ?? 0, header[13] ?? 0, header[14] ?? 0, header[15] ?? 0)
  if (chunk === 'VP8X' && header.length >= 30) {
    const width = readUint24LE(header, 24) + 1
    const height = readUint24LE(header, 27) + 1
    return { width, height }
  }
  if (chunk === 'VP8 ' && header.length >= 30) {
    const width = readUint16BE(header, 26) & 0x3fff
    const height = readUint16BE(header, 28) & 0x3fff
    return { width, height }
  }
  if (chunk === 'VP8L' && header.length >= 25) {
    const bits = readUint32BE(header, 21)
    const width = (bits & 0x3fff) + 1
    const height = ((bits >> 14) & 0x3fff) + 1
    return { width, height }
  }
  return null
}

/** Parse PNG/WebP/JPEG dimensions from header bytes only. */
export function parseImageDimensions(header: Uint8Array): ImageDimensions | null {
  const png = parsePngDimensions(header)
  if (png) return png
  const webp = parseWebpDimensions(header)
  if (webp) return webp
  return parseJpegDimensions(header)
}