import { readdirSync, readFileSync, statSync } from 'node:fs';
import { join, relative } from 'node:path';
import { describe, expect, it } from 'vitest';

/**
 * GLSL pow(x, y) is undefined for x < 0 and on real GPUs it is NaN. A base
 * that is only "nearly" non-negative is not enough: an interpolated uv a hair
 * below 0 (a quad seen edge-on or clipped by the near plane) or 1 - dot() of
 * two unit vectors a hair above 1 happens a few times a minute. One NaN pixel
 * in the HDR target is spread by the bloom blurs and the lens streak across
 * the whole frame, which shows as a one-frame black flash (the traffic trail
 * and pool shaders did exactly this). So every pow() base in a shader must be
 * guarded on its face: abs(), max(), clamp() or smoothstep(), optionally
 * scaled by a positive literal.
 */
const ROOT = join(__dirname, '..', '..');
const SRC = join(ROOT, 'src');

function sources(dir: string): string[] {
  const out: string[] = [];
  for (const name of readdirSync(dir)) {
    const path = join(dir, name);
    if (statSync(path).isDirectory()) out.push(...sources(path));
    else if (/\.(ts|vue|glsl)$/.test(name)) out.push(path);
  }
  return out;
}

/** The first argument of the call whose '(' is at `open`, or null if unbalanced. */
function firstArgument(text: string, open: number): string | null {
  let depth = 0;
  for (let i = open; i < text.length; i += 1) {
    const ch = text[i];
    if (ch === '(') depth += 1;
    else if (ch === ')') {
      depth -= 1;
      if (depth === 0) return null;
    } else if (ch === ',' && depth === 1) return text.slice(open + 1, i);
  }
  return null;
}

const GUARD = /^(abs|max|clamp|smoothstep)\s*\(/;
const POSITIVE_SCALE = /^\s*[*/]\s*\d+(\.\d*)?\s*$/;

/** Is the base a single guarded call, optionally times or over a positive literal? */
function guardedPowBase(base: string): boolean {
  const trimmed = base.trim();
  const call = GUARD.exec(trimmed);
  if (!call) return false;
  let depth = 0;
  for (let i = call[0].length - 1; i < trimmed.length; i += 1) {
    if (trimmed[i] === '(') depth += 1;
    else if (trimmed[i] === ')') {
      depth -= 1;
      if (depth === 0) {
        const rest = trimmed.slice(i + 1);
        return rest.trim() === '' || POSITIVE_SCALE.test(rest);
      }
    }
  }
  return false;
}

/** Blank out comments, so prose about pow() is skipped but code on a line
 * that merely starts with a '*' continuation is still scanned. */
function withoutComments(text: string): string {
  return text.replace(/\/\*[\s\S]*?\*\/|(^|[^:])\/\/[^\n]*/g, (comment, lead: string | undefined) =>
    (lead ?? '') + comment.slice((lead ?? '').length).replace(/[^\n]/g, ' '),
  );
}

function unguardedPowCalls(source: string): string[] {
  const text = withoutComments(source);
  const found: string[] = [];
  const pattern = /(^|[^.\w])pow\s*\(/g;
  for (let match = pattern.exec(text); match; match = pattern.exec(text)) {
    const open = match.index + match[0].length - 1;
    if (text[open + 1] === ')') continue;
    const base = firstArgument(text, open);
    if (base === null) continue;
    if (!guardedPowBase(base)) found.push(base.trim());
  }
  return found;
}

describe('shader pow() safety', () => {
  it('accepts only bases that are non-negative on their face', () => {
    expect(guardedPowBase('abs(c.x) * 2.3')).toBe(true);
    expect(guardedPowBase('abs(vLocal.y - 0.92) / 0.18')).toBe(true);
    expect(guardedPowBase('clamp(vUv.y, 0.0, 1.0)')).toBe(true);
    expect(guardedPowBase(' max( ACESFilmicToneMapping( color ), 0.0 ) ')).toBe(true);
    expect(guardedPowBase('smoothstep(0.38, 0.498, radius)')).toBe(true);
    expect(guardedPowBase('vUv.y')).toBe(false);
    expect(guardedPowBase('1.0 - vUv.y')).toBe(false);
    expect(guardedPowBase('(vUv.x - 0.5) * 5.0')).toBe(false);
    expect(guardedPowBase('1.0 - abs(dot(n, toCamera))')).toBe(false);
    expect(guardedPowBase('abs(x) - 1.0')).toBe(false);
    expect(guardedPowBase('max(a, 0.0) * b')).toBe(false);
  });

  it('finds the traffic trail regression', () => {
    expect(unguardedPowCalls('float glow = across * pow(vUv.y, 1.6) * 0.8;')).toEqual(['vUv.y']);
    expect(unguardedPowCalls('// pow() of a negative base\nMath.pow(x, 2); value.pow(2);')).toEqual([]);
    expect(unguardedPowCalls('/* pow(x, 2) */ float a = 1.0\n  * exp(-pow((x - 6.6) * 1.1, 2.0));')).toEqual([
      '(x - 6.6) * 1.1',
    ]);
  });

  it('guards every pow() base in the scene shaders', () => {
    const offenders: string[] = [];
    for (const file of sources(SRC)) {
      for (const base of unguardedPowCalls(readFileSync(file, 'utf8'))) {
        offenders.push(`${relative(ROOT, file)}: pow(${base}, ...)`);
      }
    }
    expect(offenders).toEqual([]);
  });
});