// @ts-check

import { PngNodeType } from "./constants.js";
import { areBytesEqual, chunkBytes } from "../common/bytes.js";
import crc32 from "./crc32.js";
/** @typedef {import("../../types/png.d.ts").PngNode} PngNode */

const PNG_SIGNATURE = new Uint8Array([137, 80, 78, 71, 13, 10, 26, 10]);

const textDecoder = new TextDecoder();

/**
 * @param {Uint8Array} bytes
 * @returns {null | PngNode} The root node of the PNG tree, or null if the PNG is invalid.
 */
export default (bytes) => {
  if (!isSignatureValid(bytes)) {
    console.warn("Invalid PNG signature");
    return null;
  }

  /** @type {PngNode[]} */
  const children = [{ type: PngNodeType.signature, bytes: PNG_SIGNATURE }];

  const view = new DataView(bytes.buffer, bytes.byteOffset, bytes.byteLength);

  for (let i = PNG_SIGNATURE.length; i < bytes.length;) {
    const lengthBytes = bytes.subarray(i, i + 4);
    if (lengthBytes.byteLength !== 4) {
      console.warn("Invalid data. Did the file end too early?");
      return null;
    }
    const length = view.getUint32(i);
    i += 4;

    const typeBytes = bytes.subarray(i, i + 4);
    const type = parseType(typeBytes);
    if (!type) {
      console.warn("Invalid chunk type)");
      return null;
    }
    i += 4;

    const data = bytes.subarray(i, i + length);
    if (data.length !== length) {
      console.warn("Invalid data. Did the file end too early?");
      return null;
    }
    i += length;

    const actualCrcBytes = bytes.subarray(i, i + 4);
    if (actualCrcBytes.length !== 4) {
      console.warn("Not enough bytes for the CRC. Did the file end too early?");
      return null;
    }
    const actualCrc = view.getUint32(i);
    const expectedCrc = crc32(typeBytes, data);
    if (actualCrc !== expectedCrc) {
      console.warn("Invalid chunk CRC");
      return null;
    }
    i += 4;

    const allBytes = bytes.subarray(i - length - 12, i);

    children.push(
      formatChunk(allBytes, lengthBytes, typeBytes, data, actualCrcBytes),
    );
  }

  return { type: PngNodeType.root, bytes, children };
};

/**
 * @param {Uint8Array} bytes
 * @returns {boolean}
 */
const isSignatureValid = (bytes) =>
  areBytesEqual(bytes.subarray(0, 8), PNG_SIGNATURE);

/**
 * @param {Uint8Array} bytes
 * @returns {null | string} The chunk type, or null if the chunk type is invalid.
 */
const parseType = (bytes) => {
  if (bytes.length !== 4) return null;
  const result = textDecoder.decode(bytes);
  if (!/^[A-Za-z]{4}$/.test(result)) return null;
  return result;
};

/**
 * @param {Uint8Array} allBytes All the bytes in the chunk.
 * @param {Uint8Array} lengthBytes
 * @param {Uint8Array} typeBytes
 * @param {Uint8Array} data
 * @param {Uint8Array} crcBytes
 * @returns {PngNode}
 */
const formatChunk = (allBytes, lengthBytes, typeBytes, data, crcBytes) => {
  /** @type {PngNodeType} */
  let type = PngNodeType.unknownChunk;

  /** @type {PngNode} */
  let chunkLengthNode = { type: PngNodeType.chunkLength, bytes: lengthBytes };
  /** @type {PngNode} */
  const chunkTypeNode = { type: PngNodeType.chunkType, bytes: typeBytes };
  /** @type {PngNode} */
  let chunkDataNode = { type: PngNodeType.chunkData, bytes: data };
  /** @type {PngNode} */
  const chunkCrcNode = { type: PngNodeType.chunkCrc, bytes: crcBytes };

  switch (textDecoder.decode(typeBytes)) {
    case "IHDR":
      type = PngNodeType.ihdr;
      chunkDataNode = {
        ...chunkDataNode,
        type: PngNodeType.ihdrChunkData,
        children: [
          { type: PngNodeType.ihdrWidth, bytes: data.subarray(0, 4) },
          { type: PngNodeType.ihdrHeight, bytes: data.subarray(4, 8) },
          { type: PngNodeType.ihdrBitDepth, bytes: data.subarray(8, 9) },
          { type: PngNodeType.ihdrColourType, bytes: data.subarray(9, 10) },
          {
            type: PngNodeType.ihdrCompressionMethod,
            bytes: data.subarray(10, 11),
          },
          { type: PngNodeType.ihdrFilterMethod, bytes: data.subarray(11, 12) },
          {
            type: PngNodeType.ihdrInterlaceMethod,
            bytes: data.subarray(12, 13),
          },
        ],
      };
      break;
    case "PLTE":
      type = PngNodeType.plte;
      chunkDataNode = {
        ...chunkDataNode,
        type: PngNodeType.plteChunkData,
        children: chunkBytes(data, 3).map((bytes) => ({
          type: PngNodeType.plteColor,
          bytes,
        })),
      };
      break;
    case "IDAT":
      type = PngNodeType.idat;
      chunkDataNode = { ...chunkDataNode, type: PngNodeType.idatChunkData };
      break;
    case "IEND":
      type = PngNodeType.iend;
      chunkLengthNode = {
        ...chunkLengthNode,
        type: PngNodeType.iendChunkLength,
      };
      break;
    case "tRNS":
      type = PngNodeType.trns;
      break;
    case "cHRM":
      type = PngNodeType.chrm;
      break;
    case "gAMA":
      type = PngNodeType.gama;
      break;
    case "iCCP":
      type = PngNodeType.iccp;
      break;
    case "sBIT":
      type = PngNodeType.sbit;
      break;
    case "sRGB":
      type = PngNodeType.srgb;
      break;
    case "cICP":
      type = PngNodeType.cicp;
      break;
    case "tEXt": {
      type = PngNodeType.text;
      const nullSeparatorIndex = data.indexOf(0);
      if (nullSeparatorIndex === -1) break;
      chunkDataNode = {
        ...chunkDataNode,
        type: PngNodeType.textData,
        children: [
          {
            type: PngNodeType.textKeyword,
            bytes: data.subarray(0, nullSeparatorIndex),
          },
          {
            type: PngNodeType.textNullSeparator,
            bytes: data.subarray(nullSeparatorIndex, nullSeparatorIndex + 1),
          },
          {
            type: PngNodeType.textString,
            bytes: data.subarray(nullSeparatorIndex + 1),
          },
        ],
      };
      break;
    }
    case "zTXt": {
      type = PngNodeType.ztxt;
      const nullSeparatorIndex = data.indexOf(0);
      if (nullSeparatorIndex === -1) break;
      chunkDataNode = {
        ...chunkDataNode,
        type: PngNodeType.ztxtData,
        children: [
          {
            type: PngNodeType.textKeyword,
            bytes: data.subarray(0, nullSeparatorIndex),
          },
          {
            type: PngNodeType.textNullSeparator,
            bytes: data.subarray(nullSeparatorIndex, nullSeparatorIndex + 1),
          },
          {
            type: PngNodeType.ztxtCompressionMethod,
            bytes: data.subarray(
              nullSeparatorIndex + 1,
              nullSeparatorIndex + 2,
            ),
          },
          {
            type: PngNodeType.ztxtString,
            bytes: data.subarray(nullSeparatorIndex + 2),
          },
        ],
      };
      break;
    }
    case "iTXt":
      type = PngNodeType.itxt;
      break;
    case "bKGD":
      type = PngNodeType.bkgd;
      break;
    case "hIST":
      type = PngNodeType.hist;
      break;
    case "pHYs":
      type = PngNodeType.phys;
      break;
    case "sPLT":
      type = PngNodeType.splt;
      break;
    case "eXIf":
      type = PngNodeType.exif;
      break;
    case "tIME":
      type = PngNodeType.time;
      break;
    case "acTL":
      type = PngNodeType.actl;
      break;
    case "fcTL":
      type = PngNodeType.fctl;
      break;
    case "fdAT":
      type = PngNodeType.fdat;
      break;
    case "oFFs":
      type = PngNodeType.offs;
      break;
    case "pCAL":
      type = PngNodeType.pcal;
      break;
    case "sCAL":
      type = PngNodeType.scal;
      break;
    case "gIFg":
      type = PngNodeType.gifg;
      break;
    case "gIFx":
      type = PngNodeType.gifx;
      break;
    case "gIFt":
      type = PngNodeType.gift;
      break;
    case "sTER":
      type = PngNodeType.ster;
      break;
    case "dSIG":
      type = PngNodeType.dsig;
      break;
    case "iDOT":
      type = PngNodeType.idot;
      break;
    default:
      console.warn(`Unknown chunk type ${type}`);
      break;
  }

  /** @type {PngNode[]} */
  const children = [];
  children.push(chunkLengthNode);
  children.push(chunkTypeNode);
  if (data.byteLength) children.push(chunkDataNode);
  children.push(chunkCrcNode);

  return { type, bytes: allBytes, children };
};