import test from 'node:test'
import assert from 'node:assert/strict'
import { EventEmitter } from 'node:events'
import { PassThrough, Readable } from 'node:stream'
import { mkdtempSync, mkdirSync, writeFileSync, readFileSync, existsSync } from 'node:fs'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import { upscaleOptions, upscaleDimensions, upscaleName } from '../shared/video-upscale.mjs'
import { createUpscaleHost } from '../scripts/upscale-host.mjs'
import { upscaleVideo } from '../scripts/upscale-worker.mjs'
import { execFile } from 'node:child_process'
import { promisify } from 'node:util'
const exec = promisify(execFile)

test('upscale defaults, delivery sizes, portrait and square aspect, and collision names', () => {
  assert.deepEqual(upscaleOptions(), { scale: 2, target: 'preserve', fps: 'keep', enhance: 'faithful' })
  assert.deepEqual(upscaleDimensions(1344, 768, upscaleOptions()), { width: 2688, height: 1536 })
  assert.deepEqual(upscaleDimensions(768, 1344, upscaleOptions({ target: '1080p' })), { width: 1098, height: 1920 })
  assert.deepEqual(upscaleDimensions(960, 960, upscaleOptions({ target: '4k' })), { width: 3840, height: 3840 })
  assert.equal(upscaleName('master', 4, 3), 'master_up4x-3.mp4')
  for (const value of [{ scale: 3 }, { fps: 30 }, { target: 'crop' }, { enhance: 'creative' }]) assert.throws(() => upscaleOptions(value))
})
function fixture() {
  const root = mkdtempSync(join(tmpdir(), 'aigen-upscale-'))
  const calls = [], child = new EventEmitter()
  child.stdout = new PassThrough(); child.stderr = new PassThrough()
  const host = createUpscaleHost({ root, leaseValid: token => token === 'lease', spawnProcess: (...args) => { calls.push(args); return child } })
  return { root, calls, child, host }
}
test('standalone upscale requires a lease, uploads locally, deduplicates, and reports OOM after exit', async () => {
  const f = fixture(), id = 'test-upscale-12345'
  await f.host.upload(id, Readable.from(Buffer.alloc(64)))
  assert.throws(() => f.host.start({ id }, 'invalid'), /reservation/)
  f.host.start({ id }, 'lease'); f.host.start({ id }, 'lease')
  assert.equal(f.calls.length, 1)
  assert.equal(f.calls[0][2].windowsHide, true)
  assert.match(f.calls[0][1][0], /upscale-worker/)
  assert.equal(f.host.busy(), true)
  f.child.stdout.write('AIGEN_EVENT {"error":"GPU out of ')
  f.child.stdout.write('memory"}\n')
  assert.equal(f.host.output(id), null)
  f.child.emit('close', 1)
  assert.equal(f.host.busy(), false)
  assert.equal(f.host.read(id).status, 'error')
  assert.match(f.host.read(id).error, /out of memory/)
  assert.equal(readFileSync(join(f.root, 'jobs', id, 'input.mp4')).length, 64)
})
test('cancel before submission is a durable tombstone and cannot later start a GPU worker', async () => {
  const f = fixture(), id = 'cancel-upscale-12345'
  await f.host.cancel(id)
  assert.equal(f.host.start({ id }, 'lease').status, 'cancelled')
  assert.equal(f.calls.length, 0)
  await assert.rejects(f.host.upload('../escape', Readable.from('x')), /Invalid/)
})
test('two-minute fractional-FPS master stays one job; ESRGAN precedes RIFE and audio is stream-copied', async () => {
  const root = mkdtempSync(join(tmpdir(), 'aigen-upscale-plan-'))
  mkdirSync(join(root, 'rife', 'rife-v4.6'), { recursive: true })
  const request = join(root, 'request.json'), events = [], calls = []
  writeFileSync(request, JSON.stringify({ root, parentPid: process.pid, fps: 50 }))
  let outputFrames = 0
  await upscaleVideo(request, { report: e => events.push(e), runCommand: async (exe, args) => {
    calls.push({ exe, args })
    if (exe.endsWith('ffprobe.exe')) return JSON.stringify({ streams: [{ codec_type: 'video', width: 1344, height: 768, avg_frame_rate: '30000/1001', duration: '120.02' }] })
    if (args.includes('libx264')) outputFrames += Number(args[args.indexOf('-frames:v') + 1])
    if (args.at(-1) === join(root, 'output.mp4')) writeFileSync(args.at(-1), 'mock result')
    return ''
  } })
  assert.equal(outputFrames, Math.round(Math.round(120.02 * 30000 / 1001) / (30000 / 1001) * 50))
  const engines = calls.filter(c => /ncnn-vulkan/.test(c.exe))
  assert.ok(engines.length > 100)
  for (let i = 0; i < engines.length; i += 2) { assert.match(engines[i].exe, /realesrgan/); assert.match(engines[i + 1].exe, /rife/) }
  const mux = calls.at(-1).args
  assert.equal(mux[mux.indexOf('-c:a') + 1], 'copy')
  assert.ok(mux.includes('1:a?'))
  assert.ok(!mux.includes('-shortest'))
  assert.equal(events.at(-1).progress, 100)
  assert.ok(existsSync(join(root, 'output.mp4')))
})

const cpuFfmpeg = process.env.UPSCALE_TEST_FFMPEG || join(process.cwd(), '.data', 'video-test-tools', 'ffmpeg.exe')
for (const hasAudio of [true, false]) test(`CPU-only fixture: chunk boundary, frame count and ${hasAudio ? 'bit-identical audio' : 'video-only output'}`, { skip: !existsSync(cpuFfmpeg) }, async () => {
  const root = mkdtempSync(join(tmpdir(), 'aigen-upscale-cpu-'))
  mkdirSync(join(root, 'rife', 'rife-v4.6'), { recursive: true })
  const source = join(root, 'input.mp4'), result = join(root, 'output.mp4')
  const ffmpeg = async args => (await exec(cpuFfmpeg, args, { windowsHide: true, maxBuffer: 4_000_000 })).stdout
  await ffmpeg(['-v', 'error', '-f', 'lavfi', '-i', 'testsrc2=size=48x32:rate=24:duration=2.5', ...(hasAudio ? ['-f', 'lavfi', '-i', 'sine=frequency=440:duration=2.5'] : []), '-c:v', 'libx264', ...(hasAudio ? ['-c:a', 'aac'] : []), source])
  const original = readFileSync(source)
  writeFileSync(join(root, 'request.json'), JSON.stringify({ root, parentPid: process.pid, fps: hasAudio ? 50 : 'keep' }))
  await upscaleVideo(join(root, 'request.json'), { report: () => {}, runCommand: async (exe, args) => {
    if (exe === 'powershell.exe') return ''
    if (exe.endsWith('ffprobe.exe')) return JSON.stringify({ streams: [{ codec_type: 'video', width: 48, height: 32, avg_frame_rate: '24/1', duration: '2.5' }] })
    if (exe.endsWith('realesrgan-ncnn-vulkan.exe')) {
      // Substitute a CPU scaler only in this test. Never execute the neural engine.
      return ffmpeg(['-v', 'error', '-i', join(args[args.indexOf('-i') + 1], '%08d.png'), '-vf', 'scale=192:128', '-threads', '1', join(args[args.indexOf('-o') + 1], '%08d.png')])
    }
    if (exe.endsWith('rife-ncnn-vulkan.exe')) {
      return ffmpeg(['-v', 'error', '-framerate', '24', '-i', join(args[args.indexOf('-i') + 1], '%08d.png'), '-vf', 'fps=72', '-frames:v', args[args.indexOf('-n') + 1], '-threads', '1', join(args[args.indexOf('-o') + 1], '%08d.png')])
    }
    return ffmpeg(args)
  } })
  assert.deepEqual(readFileSync(source), original)
  const hashAudio = file => ffmpeg(['-v', 'error', '-i', file, '-map', '0:a:0', '-c:a', 'copy', '-f', 'hash', '-'])
  if (hasAudio) assert.equal(await hashAudio(result), await hashAudio(source))
  else await assert.rejects(hashAudio(result), /matches no streams/)
  const frames = await ffmpeg(['-v', 'error', '-i', result, '-map', '0:v:0', '-f', 'framemd5', '-'])
  assert.equal(frames.split('\n').filter(line => /^0,/.test(line)).length, hasAudio ? 125 : 60)
  assert.match(frames, /dimensions 0: 96x64/)
})