import { readFile, writeFile } from "fs/promises";
import { inflateSync } from "zlib";
import { PNG } from "pngjs";
import { Float16Array } from "@petamoriken/float16";
import { decompress as fpzipDecompress } from "../fpzip/decompress.js";

type Args = {
  in?: string;
  out?: string;
  raw?: boolean;
  profile?: DecoderProfile;
  concat?: string[];
  help?: boolean;
};

function parseArgs(argv: string[]): Args {
  const args: Args = {};

  for (let i = 0; i < argv.length; i++) {
    const a = argv[i];

    if (a === "--help" || a === "-h") {
      args.help = true;
      continue;
    }

    if (a === "--raw") {
      args.raw = true;
      continue;
    }

    if (a === "--profile") {
      args.profile = parseDecoderProfile(argv[++i]);
      continue;
    }

    if (a === "--in") {
      args.in = argv[++i];
      continue;
    }

    if (a === "--out") {
      args.out = argv[++i];
      continue;
    }

    if (a === "--concat") {
      args.concat = [];
      while (i + 1 < argv.length && !argv[i + 1].startsWith("--")) {
        args.concat.push(argv[++i]);
      }
      continue;
    }

    throw new Error(`Unknown arg: ${a}`);
  }

  return args;
}

function usage() {
  // Intentionally minimal and compatible with the existing Bazel tool flags.
  console.log(
    [
      "Usage:",
      "  npx tsx helpers/GRPCBin2PNG.ts --in <input.bin> --out <output.png>",
      "  npx tsx helpers/GRPCBin2PNG.ts --concat <part1.bin> <part2.bin> ... --out <output.png>",
      "",
      "Options:",
      "  --raw    Treat input as uncompressed Float16 tensor (no zip/fpzip)",
      "  --profile <direct-rgb|qwen-argb|flux2-latent32>",
      "  -h, --help",
    ].join("\n")
  );
}

type DecoderProfile = "direct-rgb" | "qwen-argb" | "flux2-latent32";

function parseDecoderProfile(value: string | undefined): DecoderProfile {
  if (!value || value === "direct-rgb") return "direct-rgb";
  if (value === "qwen-argb") return "qwen-argb";
  if (value === "flux2-latent32") return "flux2-latent32";
  throw new Error("Unsupported decoder profile: " + value);
}

function clampByte(v: number) {
  if (v < 0) return 0;
  if (v > 255) return 255;
  return v | 0;
}

function mapFloatToU8(v: number) {
  // NNC/Draw Things uses: FloatType(v) * 2 / 255 - 1 for encoding
  // Inverse: (v + 1) * 255 / 2 = (v + 1) * 127.5
  // Swift Int(...) truncates toward zero = floor for positive values — match that here.
  return clampByte(Math.floor((v + 1) * 127.5));
}

function mapAlphaToU8(v: number) {
  // Alpha is stored as a straight 0..1 value (not the -1..1 color encoding).
  return clampByte(Math.round(v * 255));
}

const flux2RgbProjection = [
  [0.0058, 0.0495, -0.0099, 0.2144, 0.0166, 0.0157, -0.0398, -0.0052, -0.3527, -0.0301, -0.0107, 0.0746, 0.0156, -0.0034, 0.0032, -0.0939, 0.0018, 0.0284, -0.0024, 0.1207, 0.0128, 0.0137, 0.0095, 0, -0.0465, 0.0095, 0.029, 0.022, -0.0332, -0.0085, -0.0076, -0.0111, -0.0329],
  [0.0113, 0.0443, 0.0096, 0.3009, -0.0039, 0.0103, 0.0902, 0.0095, -0.2712, -0.0356, 0.0078, 0.009, 0.0169, -0.004, 0.0181, -0.0008, 0.0043, 0.0056, -0.0022, -0.0026, 0.0101, -0.0072, 0.0092, -0.0077, -0.0204, 0.0012, -0.0034, 0.0169, -0.0457, 0.0389, 0.0003, -0.046, -0.0718],
  [0.0073, 0.0836, 0.0644, 0.3652, -0.0054, -0.016, -0.0235, 0.0109, -0.1666, -0.018, 0.0013, -0.0941, 0.007, -0.0114, 0.008, 0.0186, 0.0104, -0.0127, -0.003, 0.0065, 0.0142, -0.0007, -0.0059, -0.0049, -0.0312, -0.0066, 0.0025, -0.0048, -0.0468, 0.0609, -0.0043, -0.0614, -0.0851],
] as const;

function decodeFlux2Pixel(values: ArrayLike<number>, base: number, rgba: Uint8Array, out: number) {
  for (let channel = 0; channel < 3; channel++) {
    const coefficients = flux2RgbProjection[channel];
    let value = coefficients[32];
    for (let index = 0; index < 32; index++) {
      value += coefficients[index] * (values[base + index] ?? 0);
    }
    rgba[out + channel] = clampByte(Math.trunc(value * 127.5 + 127.5));
  }
  rgba[out + 3] = 255;
}

function decodePixel(
  profile: DecoderProfile,
  channels: number,
  values: ArrayLike<number>,
  base: number,
  rgba: Uint8Array,
  out: number
) {
  if (profile === "flux2-latent32" && channels === 32) {
    decodeFlux2Pixel(values, base, rgba, out);
    return;
  }

  if (profile === "qwen-argb" && channels === 4) {
    rgba[out + 0] = mapFloatToU8(values[base + 1] ?? 0);
    rgba[out + 1] = mapFloatToU8(values[base + 2] ?? 0);
    rgba[out + 2] = mapFloatToU8(values[base + 3] ?? 0);
    rgba[out + 3] = mapAlphaToU8(values[base + 0] ?? 1);
    return;
  }

  rgba[out + 0] = mapFloatToU8(values[base + 0] ?? 0);
  rgba[out + 1] = mapFloatToU8(values[base + 1] ?? 0);
  rgba[out + 2] = mapFloatToU8(values[base + 2] ?? 0);
  rgba[out + 3] = 255;
}

function readHeaderU32LE(data: Uint8Array, word: number) {
  const o = word * 4;
  return (
    data[o] | (data[o + 1] << 8) | (data[o + 2] << 16) | (data[o + 3] << 24)
  );
}

async function main() {
  const args = parseArgs(process.argv.slice(2));
  if (args.help) {
    usage();
    return;
  }

  const profile = args.profile ?? "direct-rgb";

  if (!args.out) throw new Error("Missing --out");
  if (!args.in && (!args.concat || args.concat.length === 0)) {
    throw new Error("Missing --in or --concat");
  }

  const inputBuffers = args.concat?.length
    ? await Promise.all(args.concat.map((p) => readFile(p)))
    : [await readFile(args.in!)];

  const input = Uint8Array.from(Buffer.concat(inputBuffers));

  if (input.byteLength < 68) throw new Error("Input too small to be a tensor");

  const magic = readHeaderU32LE(input, 0);
  const height = readHeaderU32LE(input, 6) >>> 0;
  const width = readHeaderU32LE(input, 7) >>> 0;
  const channels = readHeaderU32LE(input, 8) >>> 0;

  const isCompressed = magic === 1012247 && !args.raw;
  const payload = input.subarray(68);

  let rgba: Uint8Array;

  if (isCompressed) {
    // Try zlib-inflate first (".zip"), then fpzip (".fpzip").
    let fpzipBytes: Uint8Array = payload;

    try {
      fpzipBytes = new Uint8Array(inflateSync(Buffer.from(payload)));
    } catch {
      // Not zlib-compressed (or failed inflate) - fall back to raw fpzip bytes.
    }

    let floats: Float32Array;
    try {
      floats = await fpzipDecompress(fpzipBytes);
    } catch {
      floats = await fpzipDecompress(payload);
    }

    const pixels = width * height;
    rgba = new Uint8Array(pixels * 4);

    for (let i = 0; i < pixels; i++) {
      const out = i * 4;

      if (channels === 1) {
        const v = mapFloatToU8(floats[i] ?? 0);
        rgba[out + 0] = v;
        rgba[out + 1] = v;
        rgba[out + 2] = v;
        rgba[out + 3] = 255;
        continue;
      }

      const base = i * channels;
      decodePixel(profile, channels, floats, base, rgba, out);
    }
  } else {
    // Raw Float16 payload.
    const values = (input.byteLength - 68) / 2;
    if (!Number.isInteger(values))
      throw new Error("Invalid raw Float16 payload length");

    const f16 = new Float16Array(input.buffer, input.byteOffset + 68, values);

    const pixels = width * height;
    rgba = new Uint8Array(pixels * 4);

    for (let i = 0; i < pixels; i++) {
      const out = i * 4;

      if (channels === 1) {
        const v = mapFloatToU8(f16[i] ?? 0);
        rgba[out + 0] = v;
        rgba[out + 1] = v;
        rgba[out + 2] = v;
        rgba[out + 3] = 255;
        continue;
      }

      const base = i * channels;
      decodePixel(profile, channels, f16, base, rgba, out);
    }
  }

  const png = new PNG({ width, height });
  png.data = Buffer.from(rgba);

  const outPng = (PNG as any).sync.write(png as any);
  await writeFile(args.out, outPng);
}

main().catch((err) => {
  console.error(err?.stack ?? String(err));
  process.exit(1);
});

import { readFile, writeFile } from "fs/promises";
import { inflateSync } from "zlib";
import { PNG } from "pngjs";
import { Float16Array } from "@petamoriken/float16";
import { decompress as fpzipDecompress } from "../fpzip/decompress.js";

type Args = {
  in?: string;
  out?: string;
  raw?: boolean;
  profile?: DecoderProfile;
  concat?: string[];
  help?: boolean;
};

function parseArgs(argv: string[]): Args {
  const args: Args = {};

  for (let i = 0; i < argv.length; i++) {
    const a = argv[i];

    if (a === "--help" || a === "-h") {
      args.help = true;
      continue;
    }

    if (a === "--raw") {
      args.raw = true;
      continue;
    }

    if (a === "--profile") {
      args.profile = parseDecoderProfile(argv[++i]);
      continue;
    }

    if (a === "--in") {
      args.in = argv[++i];
      continue;
    }

    if (a === "--out") {
      args.out = argv[++i];
      continue;
    }

    if (a === "--concat") {
      args.concat = [];
      while (i + 1 < argv.length && !argv[i + 1].startsWith("--")) {
        args.concat.push(argv[++i]);
      }
      continue;
    }

    throw new Error(`Unknown arg: ${a}`);
  }

  return args;
}

function usage() {
  // Intentionally minimal and compatible with the existing Bazel tool flags.
  console.log(
    [
      "Usage:",
      "  npx tsx helpers/GRPCBin2PNG.ts --in <input.bin> --out <output.png>",
      "  npx tsx helpers/GRPCBin2PNG.ts --concat <part1.bin> <part2.bin> ... --out <output.png>",
      "",
      "Options:",
      "  --raw    Treat input as uncompressed Float16 tensor (no zip/fpzip)",
      "  --profile <direct-rgb|qwen-argb|flux2-latent32>",
      "  -h, --help",
    ].join("\n")
  );
}

type DecoderProfile = "direct-rgb" | "qwen-argb" | "flux2-latent32";

function parseDecoderProfile(value: string | undefined): DecoderProfile {
  if (!value || value === "direct-rgb") return "direct-rgb";
  if (value === "qwen-argb") return "qwen-argb";
  if (value === "flux2-latent32") return "flux2-latent32";
  throw new Error("Unsupported decoder profile: " + value);
}

function clampByte(v: number) {
  if (v < 0) return 0;
  if (v > 255) return 255;
  return v | 0;
}

function mapFloatToU8(v: number) {
  // NNC/Draw Things uses: FloatType(v) * 2 / 255 - 1 for encoding
  // Inverse: (v + 1) * 255 / 2 = (v + 1) * 127.5
  // Swift Int(...) truncates toward zero = floor for positive values — match that here.
  return clampByte(Math.floor((v + 1) * 127.5));
}

function mapAlphaToU8(v: number) {
  // Alpha is stored as a straight 0..1 value (not the -1..1 color encoding).
  return clampByte(Math.round(v * 255));
}

const flux2RgbProjection = [
  [0.0058, 0.0495, -0.0099, 0.2144, 0.0166, 0.0157, -0.0398, -0.0052, -0.3527, -0.0301, -0.0107, 0.0746, 0.0156, -0.0034, 0.0032, -0.0939, 0.0018, 0.0284, -0.0024, 0.1207, 0.0128, 0.0137, 0.0095, 0, -0.0465, 0.0095, 0.029, 0.022, -0.0332, -0.0085, -0.0076, -0.0111, -0.0329],
  [0.0113, 0.0443, 0.0096, 0.3009, -0.0039, 0.0103, 0.0902, 0.0095, -0.2712, -0.0356, 0.0078, 0.009, 0.0169, -0.004, 0.0181, -0.0008, 0.0043, 0.0056, -0.0022, -0.0026, 0.0101, -0.0072, 0.0092, -0.0077, -0.0204, 0.0012, -0.0034, 0.0169, -0.0457, 0.0389, 0.0003, -0.046, -0.0718],
  [0.0073, 0.0836, 0.0644, 0.3652, -0.0054, -0.016, -0.0235, 0.0109, -0.1666, -0.018, 0.0013, -0.0941, 0.007, -0.0114, 0.008, 0.0186, 0.0104, -0.0127, -0.003, 0.0065, 0.0142, -0.0007, -0.0059, -0.0049, -0.0312, -0.0066, 0.0025, -0.0048, -0.0468, 0.0609, -0.0043, -0.0614, -0.0851],
] as const;

function decodeFlux2Pixel(values: ArrayLike<number>, base: number, rgba: Uint8Array, out: number) {
  for (let channel = 0; channel < 3; channel++) {
    const coefficients = flux2RgbProjection[channel];
    let value = coefficients[32];
    for (let index = 0; index < 32; index++) {
      value += coefficients[index] * (values[base + index] ?? 0);
    }
    rgba[out + channel] = clampByte(Math.trunc(value * 127.5 + 127.5));
  }
  rgba[out + 3] = 255;
}

function decodePixel(
  profile: DecoderProfile,
  channels: number,
  values: ArrayLike<number>,
  base: number,
  rgba: Uint8Array,
  out: number
) {
  if (profile === "flux2-latent32" && channels === 32) {
    decodeFlux2Pixel(values, base, rgba, out);
    return;
  }

  if (profile === "qwen-argb" && channels === 4) {
    rgba[out + 0] = mapFloatToU8(values[base + 1] ?? 0);
    rgba[out + 1] = mapFloatToU8(values[base + 2] ?? 0);
    rgba[out + 2] = mapFloatToU8(values[base + 3] ?? 0);
    rgba[out + 3] = mapAlphaToU8(values[base + 0] ?? 1);
    return;
  }

  rgba[out + 0] = mapFloatToU8(values[base + 0] ?? 0);
  rgba[out + 1] = mapFloatToU8(values[base + 1] ?? 0);
  rgba[out + 2] = mapFloatToU8(values[base + 2] ?? 0);
  rgba[out + 3] = 255;
}

function readHeaderU32LE(data: Uint8Array, word: number) {
  const o = word * 4;
  return (
    data[o] | (data[o + 1] << 8) | (data[o + 2] << 16) | (data[o + 3] << 24)
  );
}

async function main() {
  const args = parseArgs(process.argv.slice(2));
  if (args.help) {
    usage();
    return;
  }

  const profile = args.profile ?? "direct-rgb";

  if (!args.out) throw new Error("Missing --out");
  if (!args.in && (!args.concat || args.concat.length === 0)) {
    throw new Error("Missing --in or --concat");
  }

  const inputBuffers = args.concat?.length
    ? await Promise.all(args.concat.map((p) => readFile(p)))
    : [await readFile(args.in!)];

  const input = Uint8Array.from(Buffer.concat(inputBuffers));

  if (input.byteLength < 68) throw new Error("Input too small to be a tensor");

  const magic = readHeaderU32LE(input, 0);
  const height = readHeaderU32LE(input, 6) >>> 0;
  const width = readHeaderU32LE(input, 7) >>> 0;
  const channels = readHeaderU32LE(input, 8) >>> 0;

  const isCompressed = magic === 1012247 && !args.raw;
  const payload = input.subarray(68);

  let rgba: Uint8Array;

  if (isCompressed) {
    // Try zlib-inflate first (".zip"), then fpzip (".fpzip").
    let fpzipBytes: Uint8Array = payload;

    try {
      fpzipBytes = new Uint8Array(inflateSync(Buffer.from(payload)));
    } catch {
      // Not zlib-compressed (or failed inflate) - fall back to raw fpzip bytes.
    }

    let floats: Float32Array;
    try {
      floats = await fpzipDecompress(fpzipBytes);
    } catch {
      floats = await fpzipDecompress(payload);
    }

    const pixels = width * height;
    rgba = new Uint8Array(pixels * 4);

    for (let i = 0; i < pixels; i++) {
      const out = i * 4;

      if (channels === 1) {
        const v = mapFloatToU8(floats[i] ?? 0);
        rgba[out + 0] = v;
        rgba[out + 1] = v;
        rgba[out + 2] = v;
        rgba[out + 3] = 255;
        continue;
      }

      const base = i * channels;
      decodePixel(profile, channels, floats, base, rgba, out);
    }
  } else {
    // Raw Float16 payload.
    const values = (input.byteLength - 68) / 2;
    if (!Number.isInteger(values))
      throw new Error("Invalid raw Float16 payload length");

    const f16 = new Float16Array(input.buffer, input.byteOffset + 68, values);

    const pixels = width * height;
    rgba = new Uint8Array(pixels * 4);

    for (let i = 0; i < pixels; i++) {
      const out = i * 4;

      if (channels === 1) {
        const v = mapFloatToU8(f16[i] ?? 0);
        rgba[out + 0] = v;
        rgba[out + 1] = v;
        rgba[out + 2] = v;
        rgba[out + 3] = 255;
        continue;
      }

      const base = i * channels;
      decodePixel(profile, channels, f16, base, rgba, out);
    }
  }

  const png = new PNG({ width, height });
  png.data = Buffer.from(rgba);

  const outPng = (PNG as any).sync.write(png as any);
  await writeFile(args.out, outPng);
}

main().catch((err) => {
  console.error(err?.stack ?? String(err));
  process.exit(1);
});