#!/usr/bin/env node

/**
 * Generate synthetic environment maps in long-lat (equirectangular) format
 * Outputs HDR format images for testing
 */

const fs = require('fs');
const path = require('path');

// HDR file writer (Radiance RGBE format)
class HDRWriter {
  constructor(width, height) {
    this.width = width;
    this.height = height;
    this.data = new Float32Array(width * height * 3);
  }

  setPixel(x, y, r, g, b) {
    const idx = (y * this.width + x) * 3;
    this.data[idx] = r;
    this.data[idx + 1] = g;
    this.data[idx + 2] = b;
  }

  // Convert float RGB to RGBE format
  rgbToRGBE(r, g, b) {
    const maxVal = Math.max(r, g, b);
    if (maxVal < 1e-32) {
      return [0, 0, 0, 0];
    }

    const exponent = Math.ceil(Math.log2(maxVal));
    const mantissa = Math.pow(2, -exponent);

    return [
      Math.round(r * mantissa * 255),
      Math.round(g * mantissa * 255),
      Math.round(b * mantissa * 255),
      exponent + 128
    ];
  }

  save(filename) {
    // Create HDR header
    const header = `#?RADIANCE\n# Generated by TinyUSDZ synthetic envmap generator\nFORMAT=32-bit_rle_rgbe\n\n-Y ${this.height} +X ${this.width}\n`;

    // Convert to RGBE format
    const rgbeData = new Uint8Array(this.width * this.height * 4);
    for (let i = 0; i < this.width * this.height; i++) {
      const r = this.data[i * 3];
      const g = this.data[i * 3 + 1];
      const b = this.data[i * 3 + 2];
      const rgbe = this.rgbToRGBE(r, g, b);
      rgbeData[i * 4] = rgbe[0];
      rgbeData[i * 4 + 1] = rgbe[1];
      rgbeData[i * 4 + 2] = rgbe[2];
      rgbeData[i * 4 + 3] = rgbe[3];
    }

    // Write file
    const headerBuffer = Buffer.from(header, 'ascii');
    const fullBuffer = Buffer.concat([headerBuffer, Buffer.from(rgbeData)]);
    fs.writeFileSync(filename, fullBuffer);
    console.log(`Generated: ${filename}`);
  }
}

// Convert long-lat coordinates to 3D direction vector
function lonLatToDir(lon, lat) {
  const theta = lon * Math.PI * 2; // 0 to 2π
  const phi = lat * Math.PI; // 0 to π

  return {
    x: Math.sin(phi) * Math.cos(theta),
    y: Math.cos(phi),
    z: Math.sin(phi) * Math.sin(theta)
  };
}

// Check if direction is on a specific cube face
function isOnCubeFace(dir, face) {
  const absX = Math.abs(dir.x);
  const absY = Math.abs(dir.y);
  const absZ = Math.abs(dir.z);
  const maxAbs = Math.max(absX, absY, absZ);

  const threshold = maxAbs * 0.98; // Slightly smaller to avoid edge artifacts

  switch(face) {
    case '+X': return dir.x > threshold;
    case '-X': return dir.x < -threshold;
    case '+Y': return dir.y > threshold;
    case '-Y': return dir.y < -threshold;
    case '+Z': return dir.z > threshold;
    case '-Z': return dir.z < -threshold;
    default: return false;
  }
}

// Generate RGB axis environment map
function generateRGBAxisEnvmap(width, height) {
  const hdr = new HDRWriter(width, height);

  for (let y = 0; y < height; y++) {
    for (let x = 0; x < width; x++) {
      const lon = x / width;
      const lat = y / height;
      const dir = lonLatToDir(lon, lat);

      let r = 0, g = 0, b = 0;

      // +X = Red, -X = Black
      if (dir.x > 0) r = dir.x;

      // +Y = Green, -Y = Black
      if (dir.y > 0) g = dir.y;

      // +Z = Blue, -Z = Black
      if (dir.z > 0) b = dir.z;

      hdr.setPixel(x, y, r, g, b);
    }
  }

  return hdr;
}

// Generate single-face white environment map
function generateSingleFaceEnvmap(width, height, face) {
  const hdr = new HDRWriter(width, height);

  for (let y = 0; y < height; y++) {
    for (let x = 0; x < width; x++) {
      const lon = x / width;
      const lat = y / height;
      const dir = lonLatToDir(lon, lat);

      const isWhite = isOnCubeFace(dir, face) ? 1.0 : 0.0;
      hdr.setPixel(x, y, isWhite, isWhite, isWhite);
    }
  }

  return hdr;
}

// Generate upper hemisphere white, lower hemisphere black
function generateHemisphereEnvmap(width, height) {
  const hdr = new HDRWriter(width, height);

  for (let y = 0; y < height; y++) {
    for (let x = 0; x < width; x++) {
      const lon = x / width;
      const lat = y / height;
      const dir = lonLatToDir(lon, lat);

      // Upper hemisphere: Y > 0 (white)
      // Lower hemisphere: Y <= 0 (black)
      const value = dir.y > 0 ? 1.0 : 0.0;
      hdr.setPixel(x, y, value, value, value);
    }
  }

  return hdr;
}

// Main generation
const width = 128;
const height = 64;

console.log(`Generating synthetic environment maps (${width}x${height})...\n`);

// Generate RGB axis envmap
console.log('Generating RGB axis environment map...');
const rgbAxisEnv = generateRGBAxisEnvmap(width, height);
rgbAxisEnv.save('env_synthetic_rgb_axes.hdr');

// Generate single-face envmaps
const faces = [
  { name: '+X', filename: 'env_synthetic_pos_x_red.hdr' },
  { name: '-X', filename: 'env_synthetic_neg_x.hdr' },
  { name: '+Y', filename: 'env_synthetic_pos_y_green.hdr' },
  { name: '-Y', filename: 'env_synthetic_neg_y.hdr' },
  { name: '+Z', filename: 'env_synthetic_pos_z_blue.hdr' },
  { name: '-Z', filename: 'env_synthetic_neg_z.hdr' }
];

console.log('\nGenerating single-face environment maps...');
for (const face of faces) {
  const env = generateSingleFaceEnvmap(width, height, face.name);
  env.save(face.filename);
}

// Generate hemisphere envmap
console.log('\nGenerating hemisphere environment map...');
const hemisphereEnv = generateHemisphereEnvmap(width, height);
hemisphereEnv.save('env_synthetic_hemisphere_upper_white.hdr');

console.log('\nAll synthetic environment maps generated successfully!');
console.log('\nFiles created:');
console.log('  - env_synthetic_rgb_axes.hdr (R=+X, G=+Y, B=+Z)');
faces.forEach(f => console.log(`  - ${f.filename}`));
console.log('  - env_synthetic_hemisphere_upper_white.hdr (Upper hemisphere white, lower black)');
