#!/usr/bin/env node
// Gives a photo (or a video clip) a real film look, locally. No API, no cost per photo.
// Stills use sharp. Videos use ffmpeg with a 3D LUT built from the same preset, plus halation, grain and vignette.
//
// Usage:
//   node scripts/grade.mjs --input photo.jpg --preset teal-amber --output graded.webp
//   node scripts/grade.mjs --input photo.jpg --all --output-dir graded/          (one file per preset, to choose from)
//   node scripts/grade.mjs --input clip.mp4 --preset cinestill-800t --output graded.mp4
//   node scripts/grade.mjs --list
// Options:
//   --strength 0-1   blend with the original (default 1)
//   --grain 0-2      scale the preset's grain (default 1; 0 turns grain off)
//   --seed N         grain pattern seed (default 7), so a re-run gives the same file
//   --presets file   use another presets file (default assets/presets.json)
import { readFileSync, mkdirSync, writeFileSync, rmSync, mkdtempSync, existsSync } from 'node:fs'
import { dirname, join, extname, basename, resolve } from 'node:path'
import { fileURLToPath } from 'node:url'
import { spawnSync } from 'node:child_process'
import { tmpdir } from 'node:os'

const root = join(dirname(fileURLToPath(import.meta.url)), '..')
const usage = 'Usage: node scripts/grade.mjs --input photo.jpg --preset teal-amber --output graded.webp [--strength 1] [--grain 1] [--seed 7]\n       node scripts/grade.mjs --input photo.jpg --all --output-dir graded/\n       node scripts/grade.mjs --input clip.mp4 --preset teal-amber --output graded.mp4\n       node scripts/grade.mjs --list'

const VALUE_FLAGS = new Set(['input', 'output', 'output-dir', 'preset', 'strength', 'grain', 'seed', 'presets'])
const SWITCHES = new Set(['all', 'list'])
const fail = (message) => { console.error(`${message}\n${usage}`); process.exit(1) }
const options = {}
const args = process.argv.slice(2)
for (let i = 0; i < args.length; i += 1) {
  const arg = args[i]
  if (!arg.startsWith('--')) fail(`Unexpected argument "${arg}".`)
  const key = arg.slice(2)
  if (SWITCHES.has(key)) { options[key] = true; continue }
  if (!VALUE_FLAGS.has(key)) fail(`Unknown flag --${key}.`)
  const value = args[i + 1]
  if (value === undefined || value.startsWith('--')) fail(`--${key} needs a value.`)
  options[key] = value
  i += 1
}

const presetsPath = typeof options.presets === 'string' ? options.presets : join(root, 'assets/presets.json')
if (!existsSync(presetsPath)) fail(`Presets file not found: ${presetsPath}`)
let rawPresets
try { rawPresets = JSON.parse(readFileSync(presetsPath, 'utf8')) } catch { fail(`${presetsPath} is not valid JSON.`) }
const presets = Object.fromEntries(Object.entries(rawPresets).filter(([name]) => !name.startsWith('_')))

if (options.list) {
  for (const [name, p] of Object.entries(presets)) console.log(`${name.padEnd(16)} ${p.label ?? ''}. ${p.best_for ?? ''}`)
  process.exit(0)
}
if (!options.input) fail('--input is required.')
if (!options.all && (!options.preset || !options.output)) fail('Give --preset and --output, or --all with --output-dir.')
if (options.all && !options['output-dir']) fail('--all needs --output-dir.')
if (options.all && (options.output || options.preset)) fail('--all writes one file per preset into --output-dir; leave out --output and --preset.')
if (!options.all && options['output-dir']) fail('--output-dir only goes with --all. For one look, use --preset and --output.')
if (!existsSync(options.input)) fail(`Input not found: ${options.input}`)

// Every number is checked before it reaches the pixel maths or an ffmpeg filter.
function number(value, name, lo, hi) {
  const n = Number(value)
  if (!Number.isFinite(n) || n < lo || n > hi) fail(`${name} must be a number from ${lo} to ${hi} (got "${value}").`)
  return n
}
const strength = number(options.strength ?? 1, '--strength', 0, 1)
const grainScale = number(options.grain ?? 1, '--grain', 0, 2)
const seed = Math.round(number(options.seed ?? 7, '--seed', 0, 2 ** 31))

const RANGES = {
  exposure: [-2, 2], temperature: [-1, 1], tint: [-1, 1], contrast: [0, 0.6], black_lift: [0, 0.2], white_point: [0.8, 1],
  shoulder: [0.5, 0.95], saturation: [0, 1.5], hue_axis: [-360, 360], hue_pull: [0, 0.8], diffusion: [0, 0.4], vignette: [0, 0.6],
}
function checkPreset(name, p) {
  if (!p || typeof p !== 'object') fail(`Unknown preset "${name}". Run with --list to see them.`)
  for (const [key, [lo, hi]] of Object.entries(RANGES)) if (p[key] !== undefined) number(p[key], `${name}.${key}`, lo, hi)
  for (const key of ['shadow_tint', 'highlight_tint']) {
    if (p[key] === undefined) continue
    if (!Array.isArray(p[key]) || p[key].length !== 3) fail(`${name}.${key} must be three numbers [r, g, b].`)
    p[key].forEach((v, i) => number(v, `${name}.${key}[${i}]`, -0.2, 0.2))
  }
  if (p.halation) {
    number(p.halation.threshold ?? 0.8, `${name}.halation.threshold`, 0.3, 0.99)
    number(p.halation.radius ?? 0.012, `${name}.halation.radius`, 0, 0.05)
    number(p.halation.strength ?? 0, `${name}.halation.strength`, 0, 1)
    const color = p.halation.color ?? [1, 0.4, 0.2]
    if (!Array.isArray(color) || color.length !== 3) fail(`${name}.halation.color must be three numbers [r, g, b] from 0 to 1.`)
    color.forEach((v, i) => number(v, `${name}.halation.color[${i}]`, 0, 1))
  }
  if (p.grain) {
    number(p.grain.amount ?? 0, `${name}.grain.amount`, 0, 0.15)
    number(p.grain.size ?? 1, `${name}.grain.size`, 0.3, 4)
    number(p.grain.chroma ?? 0, `${name}.grain.chroma`, 0, 1)
  }
  return p
}

function clamp(v, lo = 0, hi = 1) { return v < lo ? lo : v > hi ? hi : v }
const toLinear = (c) => (c <= 0.04045 ? c / 12.92 : ((c + 0.055) / 1.055) ** 2.4)
const toGamma = (c) => (c <= 0.0031308 ? c * 12.92 : 1.055 * Math.max(c, 0) ** (1 / 2.4) - 0.055)

// Tone curve: a gentle S for contrast, a smooth shoulder so highlights roll off instead of clipping,
// then the print range (lifted blacks, a white point just under pure white). Sampled into a lookup table.
function buildCurve(p) {
  const size = 4096
  const table = new Float32Array(size + 1)
  const c = p.contrast ?? 0
  const s = p.shoulder ?? 0.75
  const m = 0.3
  for (let i = 0; i <= size; i += 1) {
    const x = (i / size) * 1.5 // inputs up to 1.5 so exposure lifts still roll off
    let y = clamp(x, 0, 1)
    y = y - (c * Math.sin(2 * Math.PI * y)) / (2 * Math.PI)
    if (x > s) {
      const t = clamp((Math.min(x, 1.5) - s) / (1.5 - s), 0, 1)
      const h = (t ** 3 - 2 * t ** 2 + t) + (-2 * t ** 3 + 3 * t ** 2) + m * (t ** 3 - t ** 2)
      const ys = s - (c * Math.sin(2 * Math.PI * s)) / (2 * Math.PI)
      y = ys + (1 - ys) * h
    }
    table[i] = (p.black_lift ?? 0) + ((p.white_point ?? 1) - (p.black_lift ?? 0)) * clamp(y, 0, 1)
  }
  return (x) => {
    const f = clamp(x / 1.5, 0, 1) * size
    const i = Math.floor(f)
    const t = f - i
    return i >= size ? table[size] : table[i] * (1 - t) + table[i + 1] * t
  }
}

// The colour part of a look (no blur, no grain): exposure, white balance, tone curve, saturation,
// a pull of every hue towards the look's warm/cool axis, and split toning. Used per pixel and for the video LUT.
function makeColour(p) {
  const curve = buildCurve(p)
  const gain = 2 ** (p.exposure ?? 0)
  const temp = p.temperature ?? 0
  const tint = p.tint ?? 0
  const gr = gain * (1 + 0.12 * temp)
  const gg = gain * (1 - 0.06 * tint)
  const gb = gain * (1 - 0.12 * temp)
  const sat = p.saturation ?? 1
  const axis = ((p.hue_axis ?? 120) * Math.PI) / 180
  const pull = p.hue_pull ?? 0
  const st = p.shadow_tint ?? [0, 0, 0]
  const ht = p.highlight_tint ?? [0, 0, 0]
  return (r, g, b) => {
    r = curve(toGamma(toLinear(r) * gr))
    g = curve(toGamma(toLinear(g) * gg))
    b = curve(toGamma(toLinear(b) * gb))
    const y = 0.2126 * r + 0.7152 * g + 0.0722 * b
    let cb = b - y
    let cr = r - y
    const mag = Math.hypot(cb, cr)
    if (mag > 1e-5) {
      let theta = Math.atan2(cr, cb)
      let d = theta - axis
      d = Math.atan2(Math.sin(d), Math.cos(d))
      if (Math.abs(d) > Math.PI / 2) d = Math.atan2(Math.sin(d - Math.PI), Math.cos(d - Math.PI))
      const weight = clamp(mag / 0.08)
      theta -= pull * weight * d
      const m2 = mag * sat
      cb = m2 * Math.cos(theta)
      cr = m2 * Math.sin(theta)
    }
    r = y + cr
    b = y + cb
    g = (y - 0.2126 * r - 0.0722 * b) / 0.7152
    const ws = (1 - y) ** 3
    const wh = y ** 2
    return [clamp(r + ws * st[0] + wh * ht[0]), clamp(g + ws * st[1] + wh * ht[1]), clamp(b + ws * st[2] + wh * ht[2])]
  }
}

function rng(s) {
  let a = s >>> 0
  return () => {
    a = (a + 0x6d2b79f5) >>> 0
    let t = a
    t = Math.imul(t ^ (t >>> 15), t | 1)
    t ^= t + Math.imul(t ^ (t >>> 7), t | 61)
    return ((t ^ (t >>> 14)) >>> 0) / 4294967296
  }
}

async function noiseField(sharp, width, height, sigma, s) {
  const rand = rng(s)
  const raw = new Uint8Array(width * height)
  for (let i = 0; i < raw.length; i += 1) {
    const u = Math.max(rand(), 1e-9)
    const n = Math.sqrt(-2 * Math.log(u)) * Math.cos(2 * Math.PI * rand())
    raw[i] = clamp(128 + n * 40, 0, 255)
  }
  const blurred = await sharp(raw, { raw: { width, height, channels: 1 } }).blur(Math.max(0.3, sigma)).raw().toBuffer()
  let sum = 0
  let sq = 0
  for (let i = 0; i < blurred.length; i += 1) { sum += blurred[i]; sq += blurred[i] * blurred[i] }
  const mean = sum / blurred.length
  const sd = Math.sqrt(Math.max(sq / blurred.length - mean * mean, 1e-6))
  const out = new Float32Array(blurred.length)
  for (let i = 0; i < blurred.length; i += 1) out[i] = (blurred[i] - mean) / sd
  return out
}

async function gradeStill(sharp, input, p, output) {
  const { data, info } = await sharp(input).rotate().removeAlpha().toColourspace('srgb').raw().toBuffer({ resolveWithObject: true })
  const { width, height } = info
  const n = width * height
  const long = Math.max(width, height)
  const colour = makeColour(p)
  const graded = new Float32Array(n * 3)
  const lumaMask = new Uint8Array(n)
  const h = p.halation ?? { strength: 0 }
  const cache = new Map()
  for (let i = 0; i < n; i += 1) {
    const key = (data[i * 3] << 16) | (data[i * 3 + 1] << 8) | data[i * 3 + 2]
    let c = cache.get(key)
    if (!c) {
      c = colour(data[i * 3] / 255, data[i * 3 + 1] / 255, data[i * 3 + 2] / 255)
      if (cache.size < 400000) cache.set(key, c)
    }
    graded[i * 3] = c[0]
    graded[i * 3 + 1] = c[1]
    graded[i * 3 + 2] = c[2]
    const y = 0.2126 * c[0] + 0.7152 * c[1] + 0.0722 * c[2]
    const t = clamp((y - (h.threshold ?? 0.8)) / (1 - (h.threshold ?? 0.8) + 1e-6))
    lumaMask[i] = Math.round(t * t * (3 - 2 * t) * 255)
  }

  // Diffusion: a soft glow, like a mist filter on the lens.
  if (p.diffusion > 0) {
    const eight = Buffer.alloc(n * 3)
    for (let i = 0; i < n * 3; i += 1) eight[i] = Math.round(graded[i] * 255)
    const blur = await sharp(eight, { raw: { width, height, channels: 3 } }).blur(Math.max(0.3, long * 0.012)).raw().toBuffer()
    for (let i = 0; i < n * 3; i += 1) graded[i] += p.diffusion * Math.max(0, blur[i] / 255 - graded[i])
  }
  // Halation: red-orange light bleeding around bright edges, from light bouncing off the film base.
  if (h.strength > 0) {
    const glow = await sharp(Buffer.from(lumaMask), { raw: { width, height, channels: 1 } }).blur(Math.max(0.3, long * (h.radius ?? 0.012))).raw().toBuffer()
    const col = h.color ?? [1, 0.4, 0.2]
    for (let i = 0; i < n; i += 1) {
      const g = (glow[i] / 255) * h.strength
      if (g <= 0) continue
      for (let k = 0; k < 3; k += 1) graded[i * 3 + k] += g * col[k] * (1 - graded[i * 3 + k])
    }
  }
  // Vignette and grain. Grain is strongest in the mid-tones, the way film grain shows.
  const grain = p.grain ?? { amount: 0 }
  const amount = (grain.amount ?? 0) * grainScale
  const sigma = (grain.size ?? 1) * (long / 2000) * 0.5
  const mono = amount > 0 ? await noiseField(sharp, width, height, sigma, seed) : null
  const chroma = amount > 0 && grain.chroma > 0 ? await noiseField(sharp, width, height, sigma * 1.6, seed + 101) : null
  const v = p.vignette ?? 0
  const cx = width / 2
  const cy = height / 2
  const rmax = Math.hypot(cx, cy)
  const out = Buffer.alloc(n * 3)
  for (let i = 0; i < n; i += 1) {
    const x = i % width
    const yy = (i - x) / width
    const rr = Math.hypot(x - cx, yy - cy) / rmax
    const vt = clamp((rr - 0.45) / 0.55)
    const vig = 1 - v * vt * vt * (3 - 2 * vt)
    let r = graded[i * 3] * vig
    let g = graded[i * 3 + 1] * vig
    let b = graded[i * 3 + 2] * vig
    if (mono) {
      const y = 0.2126 * r + 0.7152 * g + 0.0722 * b
      const w = Math.min(1, 0.25 + 3 * y * (1 - y))
      const gm = amount * w * mono[i]
      const gc = chroma ? amount * grain.chroma * w * chroma[i] : 0
      r += gm + gc
      g += gm
      b += gm - gc
    }
    const o = [r, g, b]
    for (let k = 0; k < 3; k += 1) {
      const orig = data[i * 3 + k] / 255
      out[i * 3 + k] = Math.round(clamp(orig + (o[k] - orig) * strength) * 255)
    }
  }
  let img = sharp(out, { raw: { width, height, channels: 3 } })
  const ext = extname(output).toLowerCase()
  if (ext === '.webp') img = img.webp({ quality: 86, smartSubsample: true })
  else if (ext === '.png') img = img.png()
  else img = img.jpeg({ quality: 92, mozjpeg: true })
  await img.toFile(output)
}

function writeCube(p, file, size = 33) {
  const colour = makeColour(p)
  const lines = [`TITLE "film-look-grade"`, `LUT_3D_SIZE ${size}`]
  for (let b = 0; b < size; b += 1)
    for (let g = 0; g < size; g += 1)
      for (let r = 0; r < size; r += 1) {
        const c = colour(r / (size - 1), g / (size - 1), b / (size - 1))
        lines.push(c.map((v) => v.toFixed(5)).join(' '))
      }
  writeFileSync(file, lines.join('\n') + '\n')
}

function gradeVideo(input, p, output) {
  const probe = spawnSync('ffprobe', ['-v', 'error', '-print_format', 'json', '-show_streams', input], { encoding: 'utf8' })
  if (probe.error) { console.error('ffprobe was not found. Install ffmpeg 5 or newer (it includes ffprobe).'); process.exit(1) }
  if (probe.status !== 0) { console.error(`${input} is not a video ffmpeg can read. Use an MP4, MOV or WebM clip.`); process.exit(1) }
  const streams = JSON.parse(probe.stdout).streams ?? []
  const video = streams.find((s) => s.codec_type === 'video' && !['mjpeg', 'png', 'webp', 'bmp', 'gif'].includes(s.codec_name))
  const audio = streams.find((s) => s.codec_type === 'audio')
  const w = Number(video?.width)
  const hgt = Number(video?.height)
  if (!(w > 0 && hgt > 0)) { console.error(`${input} has no video picture to grade. Use an MP4, MOV or WebM clip.`); process.exit(1) }
  // HDR clips (PQ or HLG, e.g. iPhone HDR video) would come out washed out or clipped through an SDR look.
  if (['smpte2084', 'arib-std-b67'].includes(video.color_transfer)) {
    console.error(`${input} is HDR video (${video.color_transfer === 'smpte2084' ? 'PQ / HDR10' : 'HLG'}). This grade works on SDR (Rec.709) video only. Export an SDR copy first: on a phone, turn off HDR video or share as "Most compatible"; in an editor, export Rec.709 SDR. See "HDR and 10-bit video" in references/film-stocks.md.`)
    process.exit(1)
  }
  if (/1[02]le|1[02]be|p010/.test(video.pix_fmt ?? '')) console.error(`Note: ${input} is ${video.pix_fmt} (10-bit or 12-bit SDR). The graded file is 8-bit H.264 for the web.`)
  // MP4 holds AAC, MP3 and ALAC everywhere; anything else (Opus or Vorbis from WebM, PCM from MOV) is converted to AAC.
  const audioArgs = !audio ? [] : ['aac', 'mp3', 'alac'].includes(audio.codec_name) ? ['-c:a', 'copy'] : ['-c:a', 'aac', '-b:a', '192k']

  const work = mkdtempSync(join(tmpdir(), 'film-look-'))
  writeCube(p, join(work, 'look.cube'))
  const h = p.halation ?? { strength: 0 }
  const g = p.grain ?? { amount: 0 }
  const long = Math.max(w, hgt)
  const t = Math.round((h.threshold ?? 0.8) * 255)
  const col = h.color ?? [1, 0.4, 0.2]
  const grainStrength = Math.round((g.amount ?? 0) * grainScale * 255 * 0.7)
  const parts = [
    '[0:v]format=gbrp,lut3d=file=look.cube:interp=tetrahedral,split=3[base][hal][dif]',
    `[hal]colorlevels=rimin=${t / 255}:gimin=${t / 255}:bimin=${t / 255},gblur=sigma=${(long * (h.radius ?? 0.012)).toFixed(1)},colorchannelmixer=rr=${0.2126 * col[0]}:rg=${0.7152 * col[0]}:rb=${0.0722 * col[0]}:gr=${0.2126 * col[1]}:gg=${0.7152 * col[1]}:gb=${0.0722 * col[1]}:br=${0.2126 * col[2]}:bg=${0.7152 * col[2]}:bb=${0.0722 * col[2]}[halo]`,
    `[dif]gblur=sigma=${(long * 0.012).toFixed(1)}[soft]`,
    `[base][soft]blend=all_mode=lighten:all_opacity=${p.diffusion ?? 0}[a]`,
    `[a][halo]blend=all_mode=screen:all_opacity=${h.strength ?? 0}[b]`,
    `[b]vignette=angle=${((p.vignette ?? 0) * 2.2).toFixed(3)},format=yuv420p${grainStrength > 0 ? `,noise=c0s=${grainStrength}:c0f=t+u` : ''}[v]`,
  ]
  // ffmpeg runs inside the temp folder so the LUT is referenced by a plain name, with nothing to escape.
  const run = spawnSync('ffmpeg', ['-y', '-v', 'error', '-i', resolve(input), '-filter_complex', parts.join(';'), '-map', '[v]', '-map', '0:a?', '-c:v', 'libx264', '-crf', '18', '-preset', 'slow', '-pix_fmt', 'yuv420p', ...audioArgs, '-movflags', '+faststart', resolve(output)], { stdio: 'inherit', cwd: work })
  rmSync(work, { recursive: true, force: true })
  if (run.error) { console.error('ffmpeg was not found. Install ffmpeg 5 or newer.'); process.exit(1) }
  if (run.status !== 0) { console.error('ffmpeg could not grade this clip; see its error above.'); process.exit(1) }
}

const isVideo = /\.(mp4|mov|webm|m4v)$/i.test(options.input)
const isImage = /\.(jpe?g|png|webp|tiff?|avif)$/i.test(options.input)
if (!isVideo && !isImage) fail(`"${basename(options.input)}" is not a photo or video this script can read. Use a JPEG, PNG or WebP photo (convert HEIC to JPEG first), or an MP4, MOV or WebM clip.`)
if (options.output && !options.all) {
  const outExt = extname(options.output).toLowerCase()
  if (isVideo && outExt !== '.mp4') fail('A video clip is graded to an .mp4 file. Change --output to end in .mp4.')
  if (!isVideo && !['.jpg', '.jpeg', '.webp', '.png'].includes(outExt)) fail('A photo is graded to .jpg, .webp or .png. Change the --output extension.')
}
const jobs = options.all
  ? Object.keys(presets).map((name) => [name, join(options['output-dir'], `${basename(options.input, extname(options.input))}-${name}${isVideo ? '.mp4' : extname(options.input).toLowerCase() === '.png' ? '.png' : '.jpg'}`)])
  : [[options.preset, options.output]]
for (const [name] of jobs) checkPreset(name, presets[name])
if (options.all) mkdirSync(options['output-dir'], { recursive: true })

let sharp
if (!isVideo) {
  try {
    sharp = (await import('sharp')).default
  } catch {
    console.error('This script needs sharp: Node 20.9 or newer; run `npm install sharp` in the Skill folder.')
    process.exit(1)
  }
}
for (const [name, output] of jobs) {
  if (isVideo) gradeVideo(options.input, presets[name], output)
  else {
    try {
      await gradeStill(sharp, options.input, presets[name], output)
    } catch (error) {
      fail(`Could not read or write the photo (${error.message}). Use a JPEG, PNG or WebP photo and a writable output path.`)
    }
  }
  console.log(`${name}: ${output}`)
}